Studio: shareable per-checkpoint preview links (#6486)
* checkpoint preview endpoint * harden new preview endpoints * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review * Studio preview: pin adapter, guard streaming submit, robust copy-link Harden the public per-checkpoint preview surface: - Pin use_adapter=True in the preview payload sanitizer. Otherwise an unauthenticated /p caller can POST use_adapter=false, which calls disable_adapter_layers() on the shared in-memory model without restoring it; since load_model skips reloads for the same checkpoint, every later visitor (the page never sends the field) keeps getting base-model output instead of the fine-tuned checkpoint. Forcing it on also re-enables a previously disabled adapter and no-ops on merged checkpoints. - Ignore preview-page submits while a response is streaming. The send button was disabled but the Enter handler still called requestSubmit(), so a second request could start before the first reply landed in msgs and reorder the chat history. Both the keydown and submit handlers now honor the disabled button. - Keep the cloudflare-URL polling loop alive across transient startup fetch errors instead of letting one rejection halt it. - Build the copy-link from a backend preview_ref (output dir relative to outputs_root, gated on previewability and the two-segment /p route limit) so a nested output dir no longer copies a basename-only link that 404s. Expose preview_ref on training run summaries. Add route-level security tests (path traversal, payload sanitization, asset containment, CSP header, HTML title escaping, streaming lock held until drained) and preview_ref unit tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio preview: Safari-safe submit and adapter pin only for LoRA Follow-ups from cross-browser and route simulations: - Preview page: send the message from a shared send() helper called by both the form submit and the Enter key, instead of form.requestSubmit(). The latter throws on Safari < 16 and older iOS, which broke Enter-to-send there. Verified across Chromium, Firefox and WebKit with Playwright. - Only pin use_adapter=True when the resolved checkpoint is a LoRA adapter (adapter_config.json present); for a merged checkpoint strip it to None. A merged model has no adapter to toggle, so forcing it on only produced a per-request "not a PeftModel" warning. The cross-request base-model contamination fix still holds for LoRA previews. Add a merged-checkpoint test asserting use_adapter is stripped to None. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio preview: trim verbose comments Tighten comments across the preview routes, page, checkpoint helpers, and tests to short single-line notes; drop ones that just restate the code. No behavior change (verified comment/docstring-only with comment_tools.py check). * Harden preview routes for PR #6486 - Return a generic 400 detail on a rejected preview path so the public /p route never echoes the absolute install path (the real reason is logged server-side instead). - Strip confirm_tool_calls, session_id and rag_scope in the preview payload sanitizer so the public surface stays inert regardless of the tool gate. - Use Path.is_relative_to for the asset containment check, matching the rest of the codebase. - Add img-src 'self' and font-src 'self' to the preview page CSP. - Preview page: on a mid-stream error keep the streamed text, flag the break, and restore the prompt so the user can retry; drop the unused --font-sans var. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Nilay
GitHub
parent
c7c353d740
commit
e5cf956601
15 changed files with 1206 additions and 10 deletions
|
|
@ -57,6 +57,7 @@ studio = [
|
|||
"backend/requirements/**/*",
|
||||
"backend/plugins/**/*",
|
||||
"backend/assets/**/*.jinja",
|
||||
"backend/assets/**/*.html",
|
||||
"backend/core/data_recipe/oxc-validator/*.json",
|
||||
"backend/core/data_recipe/oxc-validator/*.mjs",
|
||||
]
|
||||
|
|
|
|||
398
studio/backend/assets/preview_page.html
Normal file
398
studio/backend/assets/preview_page.html
Normal file
|
|
@ -0,0 +1,398 @@
|
|||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8" />
|
||||
<meta name="viewport" content="width=device-width,initial-scale=1" />
|
||||
<title>__TITLE__ - Unsloth</title>
|
||||
<style>
|
||||
@font-face {
|
||||
font-family: "Hellix";
|
||||
src: url("/p/_assets/fonts/Hellix-Medium.woff") format("woff");
|
||||
font-weight: 500;
|
||||
font-display: swap;
|
||||
}
|
||||
@font-face {
|
||||
font-family: "Hellix";
|
||||
src: url("/p/_assets/fonts/Hellix-SemiBold.woff2") format("woff2");
|
||||
font-weight: 600;
|
||||
font-display: swap;
|
||||
}
|
||||
:root {
|
||||
color-scheme: light dark;
|
||||
--bg: #fefefd;
|
||||
--fg: #0d0d0d;
|
||||
--muted: #858279;
|
||||
--border: #ececec;
|
||||
--user-bubble: #f5f5f5;
|
||||
--primary: #17b88b;
|
||||
--composer-bg: #ffffff;
|
||||
--composer-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
|
||||
}
|
||||
@media (prefers-color-scheme: dark) {
|
||||
:root {
|
||||
--bg: #1a1b1e;
|
||||
--fg: #ececee;
|
||||
--muted: #96979b;
|
||||
--border: #3a3d42;
|
||||
--user-bubble: #2d2e32;
|
||||
--composer-bg: #2d2e32;
|
||||
--composer-shadow: none;
|
||||
}
|
||||
}
|
||||
* {
|
||||
box-sizing: border-box;
|
||||
}
|
||||
html,
|
||||
body {
|
||||
height: 100%;
|
||||
}
|
||||
body {
|
||||
margin: 0;
|
||||
background: var(--bg);
|
||||
color: var(--fg);
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
font:
|
||||
15.5px/1.6 "Inter",
|
||||
"Inter Variable",
|
||||
-apple-system,
|
||||
BlinkMacSystemFont,
|
||||
"Segoe UI",
|
||||
system-ui,
|
||||
sans-serif;
|
||||
-webkit-font-smoothing: antialiased;
|
||||
-moz-osx-font-smoothing: grayscale;
|
||||
}
|
||||
.heading {
|
||||
font-family: "Hellix", "Space Grotesk", system-ui, sans-serif;
|
||||
}
|
||||
header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 9px;
|
||||
padding: 14px 20px;
|
||||
}
|
||||
header img {
|
||||
width: 22px;
|
||||
height: 22px;
|
||||
border-radius: 50%;
|
||||
}
|
||||
.brand {
|
||||
font-family: "Hellix", "Space Grotesk", system-ui, sans-serif;
|
||||
font-weight: 600;
|
||||
font-size: 15px;
|
||||
}
|
||||
.model {
|
||||
margin-left: auto;
|
||||
max-width: 55%;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
font-size: 12.5px;
|
||||
color: var(--muted);
|
||||
}
|
||||
#log {
|
||||
flex: 1;
|
||||
overflow-y: auto;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
padding: 8px 16px 24px;
|
||||
}
|
||||
#thread {
|
||||
width: 100%;
|
||||
max-width: 46.5rem;
|
||||
margin: 0 auto;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
}
|
||||
.welcome {
|
||||
margin: auto;
|
||||
text-align: center;
|
||||
padding: 0 16px;
|
||||
animation: fade 0.25s ease-out;
|
||||
}
|
||||
.welcome h1 {
|
||||
margin: 0;
|
||||
font-weight: 500;
|
||||
font-size: 30px;
|
||||
letter-spacing: -0.02em;
|
||||
}
|
||||
.welcome p {
|
||||
margin: 0.55rem 0 0;
|
||||
color: var(--muted);
|
||||
font-size: 14px;
|
||||
}
|
||||
.msg {
|
||||
font-size: 15.5px;
|
||||
font-weight: 450;
|
||||
letter-spacing: 0.01em;
|
||||
word-wrap: break-word;
|
||||
white-space: pre-wrap;
|
||||
animation: fade 0.15s ease-out;
|
||||
}
|
||||
.user {
|
||||
align-self: flex-end;
|
||||
max-width: 80%;
|
||||
margin-top: 24px;
|
||||
padding: 10px 16px;
|
||||
border-radius: 24px;
|
||||
background: var(--user-bubble);
|
||||
}
|
||||
.assistant {
|
||||
align-self: stretch;
|
||||
margin-top: 16px;
|
||||
line-height: 1.75;
|
||||
}
|
||||
.dots {
|
||||
display: inline-flex;
|
||||
gap: 5px;
|
||||
align-items: center;
|
||||
height: 1.6em;
|
||||
}
|
||||
.dots i {
|
||||
width: 6px;
|
||||
height: 6px;
|
||||
border-radius: 50%;
|
||||
background: var(--muted);
|
||||
animation: blink 1.2s infinite;
|
||||
}
|
||||
.dots i:nth-child(2) {
|
||||
animation-delay: 0.18s;
|
||||
}
|
||||
.dots i:nth-child(3) {
|
||||
animation-delay: 0.36s;
|
||||
}
|
||||
.composer-wrap {
|
||||
padding: 6px 16px 16px;
|
||||
}
|
||||
form {
|
||||
width: 100%;
|
||||
max-width: 46.5rem;
|
||||
margin: 0 auto;
|
||||
}
|
||||
.composer {
|
||||
display: flex;
|
||||
align-items: flex-end;
|
||||
gap: 8px;
|
||||
padding: 8px 8px 8px 18px;
|
||||
border-radius: 28px;
|
||||
background: var(--composer-bg);
|
||||
box-shadow: var(--composer-shadow);
|
||||
}
|
||||
textarea {
|
||||
flex: 1;
|
||||
border: 0;
|
||||
outline: 0;
|
||||
resize: none;
|
||||
background: transparent;
|
||||
color: var(--fg);
|
||||
font: inherit;
|
||||
line-height: 1.5;
|
||||
max-height: 200px;
|
||||
padding: 8px 0;
|
||||
}
|
||||
textarea::placeholder {
|
||||
color: var(--muted);
|
||||
}
|
||||
.send {
|
||||
flex-shrink: 0;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
width: 36px;
|
||||
height: 36px;
|
||||
border: 0;
|
||||
border-radius: 50%;
|
||||
background: var(--primary);
|
||||
color: #fff;
|
||||
cursor: pointer;
|
||||
}
|
||||
.send:disabled {
|
||||
opacity: 0.4;
|
||||
cursor: default;
|
||||
}
|
||||
.foot {
|
||||
margin: 9px auto 0;
|
||||
max-width: 46.5rem;
|
||||
text-align: center;
|
||||
font-size: 11px;
|
||||
color: var(--muted);
|
||||
}
|
||||
@keyframes blink {
|
||||
0%,
|
||||
80%,
|
||||
100% {
|
||||
opacity: 0.25;
|
||||
}
|
||||
40% {
|
||||
opacity: 1;
|
||||
}
|
||||
}
|
||||
@keyframes fade {
|
||||
from {
|
||||
opacity: 0;
|
||||
transform: translateY(2px);
|
||||
}
|
||||
to {
|
||||
opacity: 1;
|
||||
transform: none;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header>
|
||||
<img src="/p/_assets/circle-logo-small.png" alt="" /><span class="brand"
|
||||
>Unsloth</span
|
||||
><span class="model">__TITLE__</span>
|
||||
</header>
|
||||
<main id="log">
|
||||
<div id="welcome" class="welcome">
|
||||
<h1 class="heading">Chat with your model</h1>
|
||||
<p>Fine-tuned with Unsloth</p>
|
||||
</div>
|
||||
<div id="thread"></div>
|
||||
</main>
|
||||
<div class="composer-wrap">
|
||||
<form id="f">
|
||||
<div class="composer">
|
||||
<textarea
|
||||
id="i"
|
||||
rows="1"
|
||||
autocomplete="off"
|
||||
placeholder="Message this model..."
|
||||
></textarea>
|
||||
<button id="b" class="send" aria-label="Send">
|
||||
<svg
|
||||
width="18"
|
||||
height="18"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2.2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M12 19V5" />
|
||||
<path d="M5 12l7-7 7 7" />
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
<div class="foot">Served by Unsloth Studio</div>
|
||||
</form>
|
||||
</div>
|
||||
<script>
|
||||
const base = location.pathname.replace(/\/+$/, "");
|
||||
const log = document.getElementById("log"),
|
||||
thread = document.getElementById("thread"),
|
||||
welcome = document.getElementById("welcome");
|
||||
const form = document.getElementById("f"),
|
||||
input = document.getElementById("i"),
|
||||
btn = document.getElementById("b");
|
||||
const msgs = [];
|
||||
const down = () => {
|
||||
log.scrollTop = log.scrollHeight;
|
||||
};
|
||||
function autosize() {
|
||||
input.style.height = "auto";
|
||||
input.style.height = Math.min(input.scrollHeight, 200) + "px";
|
||||
}
|
||||
input.addEventListener("input", autosize);
|
||||
input.addEventListener("keydown", (e) => {
|
||||
if (e.isComposing || e.keyCode === 229) return;
|
||||
if (e.key === "Enter" && !e.shiftKey) {
|
||||
e.preventDefault();
|
||||
// send() (not form.requestSubmit, unsupported on Safari < 16) guards the btn.
|
||||
send();
|
||||
}
|
||||
});
|
||||
function add(role) {
|
||||
const d = document.createElement("div");
|
||||
d.className = "msg " + role;
|
||||
thread.appendChild(d);
|
||||
down();
|
||||
return d;
|
||||
}
|
||||
async function send() {
|
||||
// One path for button + Enter; ignore while a request is in flight.
|
||||
if (btn.disabled) return;
|
||||
const content = input.value.trim();
|
||||
if (!content) return;
|
||||
if (welcome) welcome.style.display = "none";
|
||||
input.value = "";
|
||||
autosize();
|
||||
btn.disabled = true;
|
||||
msgs.push({ role: "user", content });
|
||||
add("user").textContent = content;
|
||||
const out = add("assistant");
|
||||
out.innerHTML = '<span class="dots"><i></i><i></i><i></i></span>';
|
||||
let acc = "";
|
||||
try {
|
||||
const r = await fetch(base + "/v1/chat/completions", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
model: "preview",
|
||||
messages: msgs,
|
||||
stream: true,
|
||||
}),
|
||||
});
|
||||
if (!r.ok) {
|
||||
out.textContent =
|
||||
"Error " + r.status + ": " + (await r.text()).slice(0, 300);
|
||||
msgs.pop();
|
||||
input.value = content; // restore the prompt so the user can retry
|
||||
autosize();
|
||||
btn.disabled = false;
|
||||
return;
|
||||
}
|
||||
const reader = r.body.getReader(),
|
||||
dec = new TextDecoder();
|
||||
let buf = "";
|
||||
for (;;) {
|
||||
const { value, done } = await reader.read();
|
||||
if (done) break;
|
||||
buf += dec.decode(value, { stream: true });
|
||||
let i;
|
||||
while ((i = buf.indexOf("\n")) >= 0) {
|
||||
const line = buf.slice(0, i).trim();
|
||||
buf = buf.slice(i + 1);
|
||||
if (!line.startsWith("data:")) continue;
|
||||
const data = line.slice(5).trim();
|
||||
if (data === "[DONE]") continue;
|
||||
try {
|
||||
const j = JSON.parse(data);
|
||||
const d =
|
||||
j.choices &&
|
||||
j.choices[0] &&
|
||||
j.choices[0].delta &&
|
||||
j.choices[0].delta.content;
|
||||
if (d) {
|
||||
acc += d;
|
||||
out.textContent = acc;
|
||||
down();
|
||||
}
|
||||
} catch (_) {}
|
||||
}
|
||||
}
|
||||
if (!acc) out.textContent = "";
|
||||
msgs.push({ role: "assistant", content: acc });
|
||||
} catch (err) {
|
||||
// Keep any streamed text, flag the break, restore the prompt for retry.
|
||||
out.textContent = acc ? acc + "\n\n[connection lost]" : "Network error, please retry.";
|
||||
msgs.pop();
|
||||
input.value = content;
|
||||
autosize();
|
||||
}
|
||||
btn.disabled = false;
|
||||
input.focus();
|
||||
}
|
||||
form.addEventListener("submit", (e) => {
|
||||
e.preventDefault();
|
||||
send();
|
||||
});
|
||||
autosize();
|
||||
input.focus();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
|
|
@ -282,6 +282,7 @@ from routes import (
|
|||
training_router,
|
||||
)
|
||||
from routes.llama import router as llama_router
|
||||
from routes.preview import router as preview_router
|
||||
from hub.routes import (
|
||||
inventory_router as hub_inventory_router,
|
||||
datasets_router as hub_datasets_router,
|
||||
|
|
@ -672,6 +673,7 @@ from utils.upload_limits import ( # noqa: E402
|
|||
_BODY_PROTECTED_PREFIXES = (
|
||||
"/v1/chat/completions",
|
||||
"/v1/completions",
|
||||
"/p/",
|
||||
"/api/inference",
|
||||
"/api/data-recipe",
|
||||
"/api/datasets",
|
||||
|
|
@ -885,6 +887,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["
|
|||
|
||||
# OpenAI-compatible: mount the inference router at /v1 for external tools.
|
||||
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
|
||||
app.include_router(preview_router, prefix = "/p", tags = ["preview"])
|
||||
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
|
||||
app.include_router(settings_router, prefix = "/api/settings", tags = ["settings"])
|
||||
app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp"])
|
||||
|
|
|
|||
|
|
@ -601,6 +601,8 @@ class TrainingRunSummary(BaseModel):
|
|||
loss_sparkline: Optional[List[float]] = None
|
||||
can_resume: bool = False
|
||||
resumed_later: bool = False
|
||||
has_preview_model: bool = False
|
||||
preview_ref: Optional[str] = None
|
||||
|
||||
|
||||
class TrainingRunUpdateRequest(BaseModel):
|
||||
|
|
|
|||
196
studio/backend/routes/preview.py
Normal file
196
studio/backend/routes/preview.py
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Per-checkpoint preview endpoints: /p/{run}[/{checkpoint}]/v1/..."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import html
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse
|
||||
from loggers import get_logger
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from auth.storage import DEFAULT_ADMIN_USERNAME
|
||||
from models.inference import ChatCompletionRequest, LoadRequest
|
||||
from routes.inference import load_model, openai_chat_completions
|
||||
from state.tool_policy import tools_force_disabled
|
||||
from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# Public (no key); resolve_preview_checkpoint pins `run` under outputs_root.
|
||||
# One model loads at a time, so serialize load+generate across previews.
|
||||
_preview_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _resolve_or_4xx(run: str, checkpoint: str | None):
|
||||
try:
|
||||
return resolve_preview_checkpoint(run, checkpoint)
|
||||
except ValueError as exc:
|
||||
# Detail can carry the absolute install path on a symlink escape; log it,
|
||||
# return a generic message on this public route.
|
||||
logger.warning("preview path rejected: %s", exc)
|
||||
raise HTTPException(status_code = 400, detail = "Invalid run or checkpoint")
|
||||
except FileNotFoundError as exc:
|
||||
raise HTTPException(status_code = 404, detail = str(exc))
|
||||
|
||||
|
||||
def _sanitize_preview_payload(
|
||||
payload: ChatCompletionRequest, is_lora: bool
|
||||
) -> ChatCompletionRequest:
|
||||
# Public surface: strip tools/MCP + provider routing (no host code / open proxy).
|
||||
# Normalize use_adapter (never trust the caller): pin True for LoRA, None for
|
||||
# merged. _apply_adapter_state mutates the shared model without restoring, so an
|
||||
# unpinned `false` would persist to later visitors who omit the field.
|
||||
return payload.model_copy(
|
||||
update = {
|
||||
"tools": None,
|
||||
"enable_tools": False,
|
||||
"enabled_tools": None,
|
||||
"mcp_enabled": False,
|
||||
"bypass_permissions": False,
|
||||
"confirm_tool_calls": False,
|
||||
"session_id": None,
|
||||
"rag_scope": None,
|
||||
"openai_code_exec_container_id": None,
|
||||
"anthropic_code_exec_container_id": None,
|
||||
"provider_id": None,
|
||||
"provider_type": None,
|
||||
"external_model": None,
|
||||
"encrypted_api_key": None,
|
||||
"provider_base_url": None,
|
||||
"use_adapter": True if is_lora else None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _unlock_after(body_iterator):
|
||||
# Hold the lock until the stream drains so another checkpoint can't swap mid-stream.
|
||||
try:
|
||||
async for chunk in body_iterator:
|
||||
yield chunk
|
||||
finally:
|
||||
_preview_lock.release()
|
||||
|
||||
|
||||
async def _serve_chat(
|
||||
run: str, checkpoint: str | None, payload: ChatCompletionRequest, request: Request
|
||||
):
|
||||
path = _resolve_or_4xx(run, checkpoint)
|
||||
is_lora = (path / "adapter_config.json").exists()
|
||||
payload = _sanitize_preview_payload(payload, is_lora)
|
||||
await _preview_lock.acquire()
|
||||
keep_locked = False
|
||||
try:
|
||||
await load_model(LoadRequest(model_path = str(path)), request, DEFAULT_ADMIN_USERNAME)
|
||||
# Beats a process-wide `--enable-tools` (enable_tools=False alone wouldn't).
|
||||
with tools_force_disabled():
|
||||
response = await openai_chat_completions(payload, request, DEFAULT_ADMIN_USERNAME)
|
||||
if isinstance(response, StreamingResponse):
|
||||
response.body_iterator = _unlock_after(response.body_iterator)
|
||||
keep_locked = True
|
||||
return response
|
||||
finally:
|
||||
if not keep_locked:
|
||||
_preview_lock.release()
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def list_previews(request: Request, current_subject: str = Depends(get_current_subject)):
|
||||
base = str(request.base_url)
|
||||
previews = []
|
||||
for target in list_preview_targets():
|
||||
ref = quote(target["ref"], safe = "/")
|
||||
previews.append({**target, "url": f"{base}p/{ref}/v1"})
|
||||
return {"object": "list", "data": previews}
|
||||
|
||||
|
||||
@router.post("/{run}/v1/chat/completions")
|
||||
async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request: Request):
|
||||
return await _serve_chat(run, None, payload, request)
|
||||
|
||||
|
||||
@router.post("/{run}/{checkpoint}/v1/chat/completions")
|
||||
async def preview_chat_checkpoint(
|
||||
run: str, checkpoint: str, payload: ChatCompletionRequest, request: Request
|
||||
):
|
||||
return await _serve_chat(run, checkpoint, payload, request)
|
||||
|
||||
|
||||
def _models_response(run: str, checkpoint: str | None):
|
||||
path = _resolve_or_4xx(run, checkpoint)
|
||||
model_id = run if not checkpoint else f"{run}/{checkpoint}"
|
||||
return {
|
||||
"object": "list",
|
||||
"data": [
|
||||
{
|
||||
"id": model_id,
|
||||
"object": "model",
|
||||
"created": int(path.stat().st_mtime),
|
||||
"owned_by": "unsloth-studio",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@router.get("/{run}/v1/models")
|
||||
async def preview_models_latest(run: str):
|
||||
return _models_response(run, None)
|
||||
|
||||
|
||||
@router.get("/{run}/{checkpoint}/v1/models")
|
||||
async def preview_models_checkpoint(run: str, checkpoint: str):
|
||||
return _models_response(run, checkpoint)
|
||||
|
||||
|
||||
# Serve logo/fonts here too: the SPA static mount is absent in --api-only (Tauri).
|
||||
_FRONTEND_DIST = (Path(__file__).resolve().parents[2] / "frontend" / "dist").resolve()
|
||||
_PREVIEW_ASSET_MEDIA_TYPES = {
|
||||
".png": "image/png",
|
||||
".woff": "font/woff",
|
||||
".woff2": "font/woff2",
|
||||
}
|
||||
|
||||
|
||||
@router.get("/_assets/{asset_path:path}")
|
||||
async def preview_asset(asset_path: str):
|
||||
target = (_FRONTEND_DIST / asset_path).resolve()
|
||||
media_type = _PREVIEW_ASSET_MEDIA_TYPES.get(target.suffix.lower())
|
||||
if media_type is None or not target.is_relative_to(_FRONTEND_DIST) or not target.is_file():
|
||||
raise HTTPException(status_code = 404, detail = "Not found")
|
||||
return FileResponse(target, media_type = media_type)
|
||||
|
||||
|
||||
# Self-contained public page; only the title is interpolated.
|
||||
_PREVIEW_PAGE_HTML = (
|
||||
Path(__file__).resolve().parent.parent / "assets" / "preview_page.html"
|
||||
).read_text(encoding = "utf-8")
|
||||
|
||||
_PREVIEW_PAGE_CSP = (
|
||||
"default-src 'self'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; "
|
||||
"img-src 'self'; font-src 'self'; connect-src 'self'; base-uri 'none'"
|
||||
)
|
||||
|
||||
|
||||
def _preview_page(run: str, checkpoint: str | None) -> HTMLResponse:
|
||||
_resolve_or_4xx(run, checkpoint)
|
||||
title = run if not checkpoint else f"{run}/{checkpoint}"
|
||||
page = _PREVIEW_PAGE_HTML.replace("__TITLE__", html.escape(title))
|
||||
return HTMLResponse(page, headers = {"Content-Security-Policy": _PREVIEW_PAGE_CSP})
|
||||
|
||||
|
||||
@router.get("/{run}", response_class = HTMLResponse)
|
||||
async def preview_page_latest(run: str):
|
||||
return _preview_page(run, None)
|
||||
|
||||
|
||||
@router.get("/{run}/{checkpoint}", response_class = HTMLResponse)
|
||||
async def preview_page_checkpoint(run: str, checkpoint: str):
|
||||
return _preview_page(run, checkpoint)
|
||||
|
|
@ -27,6 +27,7 @@ from storage.studio_db import (
|
|||
list_runs,
|
||||
update_run_display_name,
|
||||
)
|
||||
from utils.models.checkpoints import has_preview_model, preview_ref
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -42,7 +43,17 @@ async def list_training_runs(
|
|||
"""List training runs, newest first."""
|
||||
result = list_runs(limit = limit, offset = offset)
|
||||
return TrainingRunListResponse(
|
||||
runs = [TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) for r in result["runs"]],
|
||||
runs = [
|
||||
TrainingRunSummary(
|
||||
**{
|
||||
**r,
|
||||
"can_resume": can_resume_run(r),
|
||||
"has_preview_model": has_preview_model(r.get("output_dir")),
|
||||
"preview_ref": preview_ref(r.get("output_dir")),
|
||||
}
|
||||
)
|
||||
for r in result["runs"]
|
||||
],
|
||||
total = result["total"],
|
||||
)
|
||||
|
||||
|
|
@ -67,6 +78,8 @@ async def get_training_run_detail(run_id: str, current_subject: str = Depends(ge
|
|||
**{
|
||||
**{k: v for k, v in run.items() if k != "config_json"},
|
||||
"can_resume": can_resume_run(run),
|
||||
"has_preview_model": has_preview_model(run.get("output_dir")),
|
||||
"preview_ref": preview_ref(run.get("output_dir")),
|
||||
}
|
||||
),
|
||||
config = config,
|
||||
|
|
@ -98,6 +111,8 @@ async def update_training_run(
|
|||
**{
|
||||
**{k: v for k, v in refreshed.items() if k != "config_json"},
|
||||
"can_resume": can_resume_run(refreshed),
|
||||
"has_preview_model": has_preview_model(refreshed.get("output_dir")),
|
||||
"preview_ref": preview_ref(refreshed.get("output_dir")),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1095,7 +1095,10 @@ def run_server(
|
|||
app.state.server_port = port if port and port > 0 else None
|
||||
# Direct (non-tunnel) base for the API panel; resolve 0.0.0.0 to the LAN IP.
|
||||
if port and port > 0:
|
||||
_direct_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
||||
_direct_host = _resolve_external_ip() if host in ("0.0.0.0", "::") else host
|
||||
# Bracket IPv6 literals so the URL is valid (http://[2405:...]:port).
|
||||
if ":" in _direct_host and not _direct_host.startswith("["):
|
||||
_direct_host = f"[{_direct_host}]"
|
||||
app.state.server_url = f"http://{_direct_host}:{port}"
|
||||
else:
|
||||
app.state.server_url = None
|
||||
|
|
|
|||
|
|
@ -10,15 +10,34 @@ Set by `unsloth run` at startup; consulted by the inference route gates.
|
|||
False -> CLI forced tools off for every request.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
import contextvars
|
||||
from contextlib import contextmanager
|
||||
from typing import Iterator, Optional
|
||||
|
||||
_tool_policy: Optional[bool] = None
|
||||
|
||||
# Per-request hard-off so public surfaces refuse tools even under a CLI `--enable-tools`.
|
||||
_force_disabled: contextvars.ContextVar[bool] = contextvars.ContextVar(
|
||||
"tool_policy_force_disabled", default = False
|
||||
)
|
||||
|
||||
|
||||
def get_tool_policy() -> Optional[bool]:
|
||||
if _force_disabled.get():
|
||||
return False
|
||||
return _tool_policy
|
||||
|
||||
|
||||
@contextmanager
|
||||
def tools_force_disabled() -> Iterator[None]:
|
||||
"""Hard-disable server-side tools for the current async context."""
|
||||
token = _force_disabled.set(True)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_force_disabled.reset(token)
|
||||
|
||||
|
||||
def set_tool_policy(value: Optional[bool]) -> None:
|
||||
if value is not None and not isinstance(value, bool):
|
||||
raise TypeError(f"tool_policy must be Optional[bool], got {type(value).__name__}")
|
||||
|
|
|
|||
134
studio/backend/tests/test_preview.py
Normal file
134
studio/backend/tests/test_preview.py
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types as _types
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
from utils.models.checkpoints import (
|
||||
list_preview_targets,
|
||||
preview_ref,
|
||||
resolve_preview_checkpoint,
|
||||
)
|
||||
|
||||
|
||||
def _make_run(outputs: Path) -> tuple[Path, Path]:
|
||||
run = outputs / "unsloth_SmolLM-135M_1775412608"
|
||||
run.mkdir(parents = True)
|
||||
(run / "adapter_config.json").write_text(
|
||||
json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"})
|
||||
)
|
||||
ckpt = run / "checkpoint-60"
|
||||
ckpt.mkdir()
|
||||
(ckpt / "adapter_config.json").write_text(
|
||||
json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"})
|
||||
)
|
||||
return run, ckpt
|
||||
|
||||
|
||||
def _point_outputs_root_at(monkeypatch, outputs: Path) -> None:
|
||||
from utils.paths import storage_roots as _sr
|
||||
from utils.models import checkpoints as _ckpt
|
||||
|
||||
monkeypatch.setattr(_sr, "outputs_root", lambda: outputs)
|
||||
# checkpoints imported outputs_root by name; patch that alias too (preview_ref uses it).
|
||||
monkeypatch.setattr(_ckpt, "outputs_root", lambda: outputs)
|
||||
|
||||
|
||||
def test_resolve_main_adapter_and_checkpoint(tmp_path: Path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
run, ckpt = _make_run(outputs)
|
||||
_point_outputs_root_at(monkeypatch, outputs)
|
||||
|
||||
assert resolve_preview_checkpoint(run.name) == run
|
||||
assert resolve_preview_checkpoint(run.name, "checkpoint-60") == ckpt
|
||||
|
||||
|
||||
def test_resolve_missing_raises_not_found(tmp_path: Path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
_make_run(outputs)
|
||||
_point_outputs_root_at(monkeypatch, outputs)
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
resolve_preview_checkpoint("does-not-exist")
|
||||
(outputs / "empty").mkdir()
|
||||
with pytest.raises(FileNotFoundError):
|
||||
resolve_preview_checkpoint("empty")
|
||||
|
||||
|
||||
def test_resolve_rejects_traversal(tmp_path: Path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
_make_run(outputs)
|
||||
_point_outputs_root_at(monkeypatch, outputs)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
resolve_preview_checkpoint("..", "etc")
|
||||
|
||||
|
||||
def test_list_preview_targets_flattens_with_latest_flag(tmp_path: Path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
run, _ = _make_run(outputs)
|
||||
_point_outputs_root_at(monkeypatch, outputs)
|
||||
|
||||
targets = list_preview_targets(str(outputs))
|
||||
by_ref = {t["ref"]: t for t in targets}
|
||||
|
||||
assert by_ref[run.name]["is_latest"] is True
|
||||
assert by_ref[run.name]["checkpoint"] is None
|
||||
assert by_ref[f"{run.name}/checkpoint-60"]["is_latest"] is False
|
||||
assert by_ref[f"{run.name}/checkpoint-60"]["checkpoint"] == "checkpoint-60"
|
||||
assert all(t["base_model"] == "HuggingFaceTB/SmolLM-135M" for t in targets)
|
||||
|
||||
|
||||
def test_preview_ref_flat_run_is_basename(tmp_path: Path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
run, _ = _make_run(outputs)
|
||||
_point_outputs_root_at(monkeypatch, outputs)
|
||||
|
||||
assert preview_ref(str(run)) == run.name
|
||||
|
||||
|
||||
def test_preview_ref_preserves_one_level_nesting(tmp_path: Path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
_point_outputs_root_at(monkeypatch, outputs)
|
||||
nested = outputs / "experiments" / "run1"
|
||||
nested.mkdir(parents = True)
|
||||
(nested / "adapter_config.json").write_text("{}")
|
||||
|
||||
# /p route supports run/checkpoint, so a single level of nesting survives.
|
||||
assert preview_ref(str(nested)) == "experiments/run1"
|
||||
|
||||
|
||||
def test_preview_ref_none_for_unpreviewable_or_too_deep(tmp_path: Path, monkeypatch):
|
||||
outputs = tmp_path / "outputs"
|
||||
_point_outputs_root_at(monkeypatch, outputs)
|
||||
|
||||
# Missing / no model artifact -> not previewable.
|
||||
assert preview_ref(None) is None
|
||||
empty = outputs / "empty"
|
||||
empty.mkdir(parents = True)
|
||||
assert preview_ref(str(empty)) is None
|
||||
|
||||
# Too deep for the two-segment /p route -> no dead link.
|
||||
deep = outputs / "a" / "b" / "run"
|
||||
deep.mkdir(parents = True)
|
||||
(deep / "adapter_config.json").write_text("{}")
|
||||
assert preview_ref(str(deep)) is None
|
||||
|
||||
# Outside outputs_root -> None.
|
||||
outside = tmp_path / "elsewhere"
|
||||
outside.mkdir()
|
||||
(outside / "adapter_config.json").write_text("{}")
|
||||
assert preview_ref(str(outside)) is None
|
||||
293
studio/backend/tests/test_preview_routes.py
Normal file
293
studio/backend/tests/test_preview_routes.py
Normal file
|
|
@ -0,0 +1,293 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Security smoke for the public /p preview routes.
|
||||
|
||||
Exercises the route layer with a real ``preview_router`` while stubbing the
|
||||
expensive model calls (``load_model`` / ``openai_chat_completions``). Covers the
|
||||
public-surface guarantees: path-traversal rejection, request sanitization
|
||||
(tools / provider routing / use_adapter), asset-path containment, the page CSP
|
||||
header + HTML escaping, and that the preview lock is held until a streaming
|
||||
response is fully drained.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types as _types
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
# Mirror test_preview.py: the real `loggers` package pulls in heavy handlers.
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import StreamingResponse
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.preview as preview
|
||||
from models.inference import ChatCompletionRequest
|
||||
|
||||
|
||||
def _make_run(outputs: Path, name: str = "demorun") -> Path:
|
||||
run = outputs / name
|
||||
run.mkdir(parents = True)
|
||||
(run / "adapter_config.json").write_text(
|
||||
json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"})
|
||||
)
|
||||
ckpt = run / "checkpoint-1"
|
||||
ckpt.mkdir()
|
||||
(ckpt / "adapter_config.json").write_text("{}")
|
||||
return run
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def captured():
|
||||
return {}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(tmp_path, monkeypatch, captured):
|
||||
outputs = tmp_path / "outputs"
|
||||
_make_run(outputs)
|
||||
|
||||
# resolve_preview_checkpoint -> resolve_output_dir -> outputs_root().
|
||||
from utils.paths import storage_roots as _sr
|
||||
|
||||
monkeypatch.setattr(_sr, "outputs_root", lambda: outputs)
|
||||
|
||||
async def _fake_load_model(load_req, request, subject):
|
||||
captured["load_path"] = load_req.model_path
|
||||
return None
|
||||
|
||||
async def _fake_chat(payload, request, subject):
|
||||
captured["payload"] = payload
|
||||
return {"ok": True}
|
||||
|
||||
monkeypatch.setattr(preview, "load_model", _fake_load_model)
|
||||
monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(preview.router, prefix = "/p")
|
||||
app.dependency_overrides[preview.get_current_subject] = lambda: "admin"
|
||||
# raise_server_exceptions=False so a 5xx surfaces as a response, not a throw.
|
||||
return TestClient(app, raise_server_exceptions = False)
|
||||
|
||||
|
||||
# ── Page rendering ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_page_renders_with_csp(client):
|
||||
r = client.get("/p/demorun")
|
||||
assert r.status_code == 200
|
||||
assert "text/html" in r.headers["content-type"]
|
||||
csp = r.headers.get("content-security-policy", "")
|
||||
assert "default-src 'self'" in csp
|
||||
assert "base-uri 'none'" in csp
|
||||
|
||||
|
||||
def test_page_escapes_title(tmp_path, monkeypatch, captured):
|
||||
outputs = tmp_path / "outputs"
|
||||
# Run dir name carries an HTML-special char; the page must escape it.
|
||||
_make_run(outputs, name = "a<b")
|
||||
from utils.paths import storage_roots as _sr
|
||||
|
||||
monkeypatch.setattr(_sr, "outputs_root", lambda: outputs)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(preview.router, prefix = "/p")
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
|
||||
r = c.get("/p/a%3Cb")
|
||||
assert r.status_code == 200
|
||||
assert "a<b" not in r.text
|
||||
assert "a<b" in r.text
|
||||
|
||||
|
||||
def test_models_endpoint_shape(client):
|
||||
r = client.get("/p/demorun/v1/models")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["object"] == "list"
|
||||
assert body["data"][0]["id"] == "demorun"
|
||||
assert body["data"][0]["owned_by"] == "unsloth-studio"
|
||||
|
||||
|
||||
def test_list_previews_builds_urls(client, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
preview,
|
||||
"list_preview_targets",
|
||||
lambda: [{"ref": "demorun", "is_latest": True}],
|
||||
)
|
||||
r = client.get("/p")
|
||||
assert r.status_code == 200
|
||||
data = r.json()["data"]
|
||||
assert data[0]["url"].endswith("/p/demorun/v1")
|
||||
|
||||
|
||||
# ── Path traversal / containment ────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"path",
|
||||
[
|
||||
"/p/..", # parent segment as run
|
||||
"/p/%2e%2e/etc", # encoded traversal
|
||||
"/p/..%2f..%2fetc/v1/models", # encoded slash traversal
|
||||
"/p/does-not-exist", # unknown run
|
||||
],
|
||||
)
|
||||
def test_traversal_and_missing_rejected(client, path):
|
||||
r = client.get(path)
|
||||
assert r.status_code in (400, 404), (path, r.status_code)
|
||||
|
||||
|
||||
def test_chat_traversal_rejected(client):
|
||||
r = client.post(
|
||||
"/p/..%2f..%2fetc/v1/chat/completions",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code in (400, 404)
|
||||
|
||||
|
||||
# ── Asset containment ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"asset",
|
||||
[
|
||||
"../../../../etc/passwd", # escapes dist
|
||||
"secrets.txt", # non-allowlisted suffix
|
||||
"nope.png", # allowlisted suffix but missing
|
||||
],
|
||||
)
|
||||
def test_asset_path_contained(client, asset):
|
||||
r = client.get(f"/p/_assets/{asset}")
|
||||
assert r.status_code == 404
|
||||
|
||||
|
||||
# ── Request sanitization ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_chat_payload_sanitized(client, captured):
|
||||
r = client.post(
|
||||
"/p/demorun/v1/chat/completions",
|
||||
json = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "rm", "parameters": {}}}],
|
||||
"enable_tools": True,
|
||||
"enabled_tools": ["python"],
|
||||
"mcp_enabled": True,
|
||||
"bypass_permissions": True,
|
||||
"provider_id": "p1",
|
||||
"provider_type": "custom",
|
||||
"provider_base_url": "http://evil.example/v1",
|
||||
"external_model": "gpt-4o",
|
||||
"use_adapter": False,
|
||||
"confirm_tool_calls": True,
|
||||
"session_id": "abc",
|
||||
"rag_scope": {"project_id": "x"},
|
||||
},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
p = captured["payload"]
|
||||
assert isinstance(p, ChatCompletionRequest)
|
||||
# Tools / code-exec off.
|
||||
assert p.tools is None
|
||||
assert p.enable_tools is False
|
||||
assert p.enabled_tools is None
|
||||
assert p.mcp_enabled is False
|
||||
assert p.bypass_permissions is False
|
||||
# Tool-loop levers neutralized regardless of the tool gate.
|
||||
assert p.confirm_tool_calls is False
|
||||
assert p.session_id is None
|
||||
assert p.rag_scope is None
|
||||
# Provider routing stripped so /p can't proxy an arbitrary endpoint.
|
||||
assert p.provider_id is None
|
||||
assert p.provider_type is None
|
||||
assert p.provider_base_url is None
|
||||
assert p.external_model is None
|
||||
# Adapter pinned on for LoRA: a caller can't flip the shared backend to base.
|
||||
assert p.use_adapter is True
|
||||
# Loads the resolved checkpoint dir, not an attacker-supplied path.
|
||||
assert captured["load_path"].endswith("demorun")
|
||||
|
||||
|
||||
def test_merged_checkpoint_strips_use_adapter(tmp_path, monkeypatch, captured):
|
||||
# Merged (non-LoRA) checkpoint: no adapter to toggle, so use_adapter -> None.
|
||||
outputs = tmp_path / "outputs"
|
||||
merged = outputs / "mergedrun"
|
||||
merged.mkdir(parents = True)
|
||||
(merged / "config.json").write_text(json.dumps({"_name_or_path": "some/base"}))
|
||||
|
||||
from utils.paths import storage_roots as _sr
|
||||
|
||||
monkeypatch.setattr(_sr, "outputs_root", lambda: outputs)
|
||||
|
||||
async def _fake_load(load_req, request, subject):
|
||||
return None
|
||||
|
||||
async def _fake_chat(payload, request, subject):
|
||||
captured["payload"] = payload
|
||||
return {"ok": True}
|
||||
|
||||
monkeypatch.setattr(preview, "load_model", _fake_load)
|
||||
monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(preview.router, prefix = "/p")
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.post(
|
||||
"/p/mergedrun/v1/chat/completions",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}], "use_adapter": False},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert captured["payload"].use_adapter is None
|
||||
|
||||
|
||||
# ── Streaming lock lifetime ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_streaming_holds_lock_until_drained(tmp_path, monkeypatch, captured):
|
||||
outputs = tmp_path / "outputs"
|
||||
_make_run(outputs)
|
||||
from utils.paths import storage_roots as _sr
|
||||
|
||||
monkeypatch.setattr(_sr, "outputs_root", lambda: outputs)
|
||||
|
||||
async def _fake_load_model(load_req, request, subject):
|
||||
return None
|
||||
|
||||
async def _gen():
|
||||
yield b"data: {}\n\n"
|
||||
yield b"data: [DONE]\n\n"
|
||||
|
||||
async def _fake_chat(payload, request, subject):
|
||||
return StreamingResponse(_gen())
|
||||
|
||||
monkeypatch.setattr(preview, "load_model", _fake_load_model)
|
||||
monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat)
|
||||
|
||||
async def _run():
|
||||
assert not preview._preview_lock.locked()
|
||||
payload = ChatCompletionRequest(messages = [{"role": "user", "content": "hi"}])
|
||||
resp = await preview._serve_chat("demorun", None, payload, request = None)
|
||||
# Lock must still be held: a second checkpoint must not swap the backend
|
||||
# mid-stream.
|
||||
assert preview._preview_lock.locked()
|
||||
chunks = [c async for c in resp.body_iterator]
|
||||
# Released only after the stream fully drains.
|
||||
assert not preview._preview_lock.locked()
|
||||
return chunks
|
||||
|
||||
chunks = asyncio.run(_run())
|
||||
assert any(b"[DONE]" in c for c in chunks)
|
||||
assert not preview._preview_lock.locked()
|
||||
|
|
@ -125,6 +125,12 @@ def is_anthropic_path(path: str) -> bool:
|
|||
return path.startswith("/v1/messages")
|
||||
|
||||
|
||||
def wants_api_error_envelope(path: str) -> bool:
|
||||
"""True for the OpenAI/Anthropic-compatible surfaces: the ``/v1/*`` mount and
|
||||
the preview ``/p/<run>[/<ckpt>]/v1/*`` mount."""
|
||||
return path.startswith("/v1/") or (path.startswith("/p/") and "/v1/" in path)
|
||||
|
||||
|
||||
def error_body_for_path(
|
||||
path,
|
||||
message,
|
||||
|
|
@ -183,15 +189,16 @@ def _summarize_validation_errors(errors) -> tuple:
|
|||
def install_api_error_handlers(app) -> None:
|
||||
"""Register validation + HTTPException handlers that emit ``/v1/*`` envelopes.
|
||||
|
||||
Both handlers are global but only transform responses for paths starting with
|
||||
``/v1/``. Non-``/v1/`` paths reproduce FastAPI's default ``{"detail": ...}``
|
||||
behavior exactly so the Studio frontend keeps working.
|
||||
Both handlers are global but only transform responses for OpenAI/Anthropic-
|
||||
compatible surfaces (see :func:`wants_api_error_envelope`: the ``/v1/*`` mount
|
||||
and the preview ``/p/.../v1/*`` mount). Every other path reproduces FastAPI's
|
||||
default ``{"detail": ...}`` behavior exactly so the Studio frontend keeps working.
|
||||
"""
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
async def _handle_validation_error(request, exc):
|
||||
path = request.url.path
|
||||
if path.startswith("/v1/"):
|
||||
if wants_api_error_envelope(path):
|
||||
summary, param = _summarize_validation_errors(exc.errors())
|
||||
return JSONResponse(
|
||||
status_code = 400,
|
||||
|
|
@ -211,7 +218,7 @@ def install_api_error_handlers(app) -> None:
|
|||
# default http_exception_handler, which returns a bodiless Response.
|
||||
if not is_body_allowed_for_status_code(exc.status_code):
|
||||
return Response(status_code = exc.status_code, headers = headers)
|
||||
if path.startswith("/v1/"):
|
||||
if wants_api_error_envelope(path):
|
||||
detail = exc.detail
|
||||
# Already a fully-formed envelope: pass through untouched.
|
||||
if isinstance(detail, dict) and ("error" in detail or detail.get("type") == "error"):
|
||||
|
|
|
|||
|
|
@ -159,3 +159,64 @@ def scan_checkpoints(
|
|||
except Exception as e:
|
||||
logger.error(f"Error scanning checkpoints: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def _is_model_dir(path: Path) -> bool:
|
||||
return (path / "config.json").exists() or (path / "adapter_config.json").exists()
|
||||
|
||||
|
||||
def has_preview_model(output_dir: Optional[str]) -> bool:
|
||||
"""True when ``output_dir`` holds a previewable root model (what ``/p/{run}``
|
||||
resolves). A cancelled run keeps ``output_dir`` but saves no root adapter."""
|
||||
if not output_dir:
|
||||
return False
|
||||
path = Path(output_dir)
|
||||
return path.is_dir() and _is_model_dir(path)
|
||||
|
||||
|
||||
def preview_ref(output_dir: Optional[str]) -> Optional[str]:
|
||||
"""``/p`` ref (``run`` or ``run/checkpoint``) relative to outputs_root, or None.
|
||||
|
||||
Posix-joined so a nested output dir keeps a working link instead of collapsing
|
||||
to its basename. None when not previewable, outside outputs_root, or deeper than
|
||||
the two path segments the ``/p`` route matches (so the UI omits a dead link).
|
||||
"""
|
||||
if not has_preview_model(output_dir):
|
||||
return None
|
||||
try:
|
||||
rel = Path(output_dir).resolve().relative_to(outputs_root().resolve())
|
||||
except (ValueError, OSError):
|
||||
return None
|
||||
parts = rel.parts
|
||||
if not parts or len(parts) > 2:
|
||||
return None
|
||||
return "/".join(parts)
|
||||
|
||||
|
||||
def resolve_preview_checkpoint(run: str, checkpoint: Optional[str] = None) -> Path:
|
||||
relative = run if not checkpoint else f"{run}/{checkpoint}"
|
||||
path = resolve_output_dir(relative)
|
||||
if not path.is_dir() or not _is_model_dir(path):
|
||||
raise FileNotFoundError(
|
||||
f"No trained checkpoint at '{relative}'. Check the run/checkpoint name (see GET /p)."
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def list_preview_targets(outputs_dir: str = str(outputs_root())) -> List[dict]:
|
||||
targets: List[dict] = []
|
||||
for run_name, checkpoints, metadata in scan_checkpoints(outputs_dir):
|
||||
for display_name, path, loss in checkpoints:
|
||||
is_latest = display_name == run_name
|
||||
checkpoint = None if is_latest else Path(path).name
|
||||
targets.append(
|
||||
{
|
||||
"run": run_name,
|
||||
"checkpoint": checkpoint,
|
||||
"ref": run_name if is_latest else f"{run_name}/{checkpoint}",
|
||||
"is_latest": is_latest,
|
||||
"loss": loss,
|
||||
"base_model": metadata.get("base_model"),
|
||||
}
|
||||
)
|
||||
return targets
|
||||
|
|
|
|||
|
|
@ -24,7 +24,10 @@ import {
|
|||
useTrainingRuntimeStore,
|
||||
} from "@/features/training";
|
||||
import { formatDuration } from "@/features/studio/sections/progress-section-lib";
|
||||
import { fetchDeviceType, usePlatformStore } from "@/config/env";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { Delete02Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { type ReactElement, useCallback, useEffect, useRef, useState } from "react";
|
||||
|
|
@ -194,6 +197,28 @@ export function HistoryCardGrid({
|
|||
const [manualFetchInFlight, setManualFetchInFlight] = useState(false);
|
||||
const { resumeTrainingRunFromHistory } = useTrainingActions();
|
||||
const isStarting = useTrainingRuntimeStore((state) => state.isStarting);
|
||||
// Copy-link base: Cloudflare tunnel > LAN host:port > origin. The tunnel
|
||||
// registers shortly after startup, so poll (bounded) until it shows.
|
||||
const cloudflareUrl = usePlatformStore((s) => s.cloudflareUrl);
|
||||
const serverUrl = usePlatformStore((s) => s.serverUrl);
|
||||
useEffect(() => {
|
||||
if (cloudflareUrl) return;
|
||||
let cancelled = false;
|
||||
void (async () => {
|
||||
for (let attempt = 0; attempt < 12 && !cancelled; attempt++) {
|
||||
try {
|
||||
await fetchDeviceType({ force: true });
|
||||
} catch {
|
||||
// Ignore startup blips; copy-link falls back to serverUrl/origin.
|
||||
}
|
||||
if (cancelled || usePlatformStore.getState().cloudflareUrl) return;
|
||||
await new Promise((r) => setTimeout(r, 2500));
|
||||
}
|
||||
})();
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
}, [cloudflareUrl]);
|
||||
|
||||
const userControllerRef = useRef<AbortController | null>(null);
|
||||
const pollControllerRef = useRef<AbortController | null>(null);
|
||||
|
|
@ -362,6 +387,8 @@ export function HistoryCardGrid({
|
|||
const isRunning = run.status === "running";
|
||||
const canResume = run.can_resume && !wasContinued;
|
||||
const isResuming = resumeTarget === run.id;
|
||||
// Backend /p ref, gated on previewability + route-expressible depth.
|
||||
const canCopyPreview = !!run.preview_ref;
|
||||
return (
|
||||
<div
|
||||
role="button"
|
||||
|
|
@ -372,7 +399,7 @@ export function HistoryCardGrid({
|
|||
isRunning
|
||||
? "border-blue-400/50 dark:border-blue-500/30"
|
||||
: "border-border/60",
|
||||
canResume && "gap-2",
|
||||
(canResume || canCopyPreview) && "gap-2",
|
||||
)}
|
||||
onClick={() => onSelectRun(run.id)}
|
||||
onKeyDown={(e) => {
|
||||
|
|
@ -411,6 +438,38 @@ export function HistoryCardGrid({
|
|||
{isResuming ? t("studio.history.resuming") : t("studio.history.resumeTraining")}
|
||||
</Button>
|
||||
)}
|
||||
{canCopyPreview && (
|
||||
<Button
|
||||
type="button"
|
||||
size="xs"
|
||||
variant="outline"
|
||||
className="absolute bottom-3 right-4 h-6 rounded-full px-2.5 text-[11px] leading-none shadow-sm"
|
||||
onClick={async (e) => {
|
||||
e.stopPropagation();
|
||||
// Encode each segment but keep "/" so the /p route matches.
|
||||
const ref = (run.preview_ref ?? "")
|
||||
.split("/")
|
||||
.map(encodeURIComponent)
|
||||
.join("/");
|
||||
const base = (
|
||||
cloudflareUrl ??
|
||||
serverUrl ??
|
||||
window.location.origin
|
||||
).replace(/\/+$/, "");
|
||||
const url = `${base}/p/${ref}`;
|
||||
const ok = await copyToClipboard(url);
|
||||
toast[ok ? "success" : "error"](
|
||||
t(
|
||||
ok
|
||||
? "studio.history.previewLinkCopied"
|
||||
: "studio.history.previewLinkCopyFailed",
|
||||
),
|
||||
);
|
||||
}}
|
||||
>
|
||||
{t("studio.history.copyPreviewLink")}
|
||||
</Button>
|
||||
)}
|
||||
<div className="min-w-0">
|
||||
<p
|
||||
className="truncate text-sm font-medium"
|
||||
|
|
@ -434,7 +493,7 @@ export function HistoryCardGrid({
|
|||
</p>
|
||||
</div>
|
||||
{run.loss_sparkline && run.loss_sparkline.length >= 2 && (
|
||||
<div className={cn(canResume && "h-7 overflow-hidden")}>
|
||||
<div className={cn((canResume || canCopyPreview) && "h-7 overflow-hidden")}>
|
||||
<Sparkline
|
||||
values={run.loss_sparkline}
|
||||
id={run.id}
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@ export interface TrainingRunSummary {
|
|||
output_dir: string | null;
|
||||
can_resume: boolean;
|
||||
resumed_later: boolean;
|
||||
has_preview_model: boolean;
|
||||
preview_ref: string | null;
|
||||
duration_seconds: number | null;
|
||||
error_message: string | null;
|
||||
loss_sparkline: number[] | null;
|
||||
|
|
|
|||
|
|
@ -763,6 +763,9 @@ export const en = {
|
|||
running: "Training in progress",
|
||||
errored: "Training errored",
|
||||
},
|
||||
copyPreviewLink: "Copy preview link",
|
||||
previewLinkCopied: "Preview link copied",
|
||||
previewLinkCopyFailed: "Couldn't copy the link",
|
||||
},
|
||||
charts: {
|
||||
settings: "Chart Settings",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue