* Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm The OpenAI-compatible endpoints serve whichever GGUF is loaded and ignore the request model field, so an OpenAI client that changes model never reloads. Add an opt-in setting that, when a /v1 request names a downloaded local GGUF different from the loaded one, loads it before serving by reusing the existing /load path (its dedup, tensor fallback, and threading apply). Unknown names still serve the loaded model, so drop-in compatibility is preserved and no remote download is triggered. Also add an optional idle auto-unload (TTL keep-warm): a pure-ASGI middleware tracks in-flight inference requests so a stream is never unloaded mid-response, and a lifespan loop unloads the model after the configured idle seconds. Both settings default off and live in the app_settings store, exposed via GET/PUT /api/settings/openai-auto-switch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: variant-aware auto-switch, /v1/responses coverage, keep-warm load stamp Follow-ups from review of the opt-in OpenAI auto-switch path: 1. Variant-aware dedup. _maybe_auto_switch_model compared only the repo id, so requesting another quant of the loaded repo (e.g. Q4_K_M loaded, Q8_0 asked) was served by the old quant. Compare hf_variant too, matching /load dedup. 2. Streaming /v1/responses now calls the auto-switch hook. It went straight into _responses_stream and only checked is_loaded, so stream=True could serve the old model or 400. Non-streaming already routed through chat completions; the hook is idempotent once loaded. 3. resolve_local_gguf tries an exact id match before splitting a trailing :VARIANT, so local ids that contain a colon (e.g. a Windows path) resolve instead of being cut at the drive letter. 4. Idle keep-warm stamps activity on a load/swap transition. _last_active was only refreshed by inference requests, so a model loaded after the server sat idle past the TTL could be unloaded before its first request. Tests cover each case. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make the /v1/responses auto-switch test order-independent The new streaming-responses test passed in isolation but failed under the CI's randomized collection order with "object has no attribute 'state'": it passed a bare object() as the request and stubbed only one dispatcher, so an ordering where the real dispatcher ran hit request.state. Give the request a state and stub both dispatchers; the test still asserts the hook fires before dispatch. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: assert /v1/responses auto-switch wiring on source, not at runtime The behavioral version executed openai_responses and relied on stubbing its callees, which a randomized collection order in CI could defeat (the real dispatcher ran and hit request attributes). Assert on the function source that the hook precedes both dispatchers instead; the hook's runtime behavior is already covered by the direct _maybe_auto_switch_model tests. * Studio: auto-switch on /v1/embeddings, GGUF-only targets, idle-unload race gate Second-pass review follow-ups on the opt-in auto-switch path: 1. /v1/embeddings now calls the auto-switch hook before the loaded-state check, matching the other model-bearing OpenAI endpoints (the keep-warm middleware already treats embeddings as inference). 2. The resolver index is now GGUF-only. The local-model scanners also surface Transformers/safetensors repos; without a filter, auto-switch could unload the GGUF and route a request into the non-GGUF loader. _has_local_gguf checks a direct .gguf, a models-dir folder, and the HF-cache snapshots layout. 3. Idle keep-warm now holds an asyncio gate across the idle check and the unload, and a request bumps inflight under the same gate, so the loop can no longer unload in the window between "looks idle" and the kill. Tests cover each. Broader local-model source parity (LM Studio, Ollama, legacy caches, custom scan folders) is a follow-up; missing one of those today just falls through to the loaded model. * Studio: variant-aware local resolver, count_tokens + audio auto-switch coverage Third-pass review follow-ups on the opt-in auto-switch path: 1. The resolver is now variant-aware via list_local_gguf_variants. It indexes only the quants actually on disk, recursing snapshots and quant subdirs such as the nested per-quant folders, so a requested repo:VARIANT resolves only when that quant is local and a bare repo resolves to a concrete local quant. This fixes two gaps: the previous shallow glob rejected nested-variant GGUF repos, and a request for an uncached quant could send /load down the remote download path, breaking the local-only contract. 2. /v1/messages/count_tokens now auto-switches like its sibling /v1/messages, so a count uses the requested model's tokenizer. 3. /api/inference/audio/generate (direct GGUF TTS) is now tracked as in-flight inference, so the idle loop cannot unload the model mid-generation. Tests cover each. Two reviewer items are left as follow-ups: indexing the remaining local sources (LM Studio, Ollama, legacy/default caches, custom scan folders), which fails safe today by falling through to the loaded model; and fully serializing concurrent different-model requests, an inherent limit of the single-slot llama backend that the opt-in feature is not designed around. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: make local GGUF resolver fail-safe so a bad model name cannot 500 The auto-switch hook calls resolve_local_gguf without its own guard, and /v1/completions and /v1/embeddings pass body.get("model") through unchanged. A non-string model (e.g. {"model": 123}) or any internal scan failure would then raise out of the resolver and turn a request that would otherwise be served by the loaded model into a 500, breaking the drop-in compatibility the feature is built on. Guard the resolver at its boundary: reject non-string input up front and wrap the lookup so any failure returns None (fall through to the loaded model). Add regression tests for both paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: per-model launch flags for auto-switched GGUF models * Studio: list switch-eligible GGUFs in /v1/models when auto-switch is on * Studio: settings UI for OpenAI model auto-switch and idle auto-unload * Studio: show save error over the disabled-idle hint in auto-switch settings * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: address gemini review (case-insensitive /v1/models retrieve, idle-input empty guard) * Studio: address codex review (deterministic override args, exclude probe/embedding models from discovery) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep-warm count_tokens, gate idle on auto-switch, drop hidden models Three hardening fixes to the opt-in auto-switch path surfaced while reviewing the work that builds on it: 1. count_tokens keep-warm. /v1/messages/count_tokens counts via the loaded tokenizer and already auto-switches, but the keep-warm middleware did not track it, so idle auto-unload could free the model mid-count. It is now a tracked in-flight path. 2. "Off means unchanged" for idle unload. get_auto_unload_idle_seconds now reports 0 while auto-switch is disabled. Idle unload only makes sense with auto-switch on (an unloaded model returns only via the next request's swap), so a stray TTL can no longer trigger a destructive unload while the feature is off, keeping the disabled state identical to pre-feature behavior. 3. Hidden models are not switch targets. The resolver index now skips what Studio hides from its own pickers (the llama.cpp validation probe, RAG embedding weights) via _is_hidden_model, so they can never be auto-switched to by name. Tests added for each. * Studio: bare-id reuse, responses validation order, in-flight tracking Review follow-ups after folding in the per-model overrides and discovery work: 1. A bare model id (no :VARIANT) is now satisfied by any loaded quant of that repo. Previously a bare name resolved to the largest local quant, so it could force a slow reload when a different quant of the same repo was already serving. An explicit repo:VARIANT request still honors the quant. 2. /v1/responses now runs the auto-switch hook after the empty-input validation so a request that 400s can no longer trigger a multi-minute model load before being rejected. The hook still precedes both dispatchers, so streaming requests switch. 3. The keep-warm middleware now tracks in-flight requests whenever auto-switch is enabled rather than only when the idle TTL is already positive, so a stream that starts with the TTL at 0 is still protected if idle-unload is enabled mid-stream. Off still passes straight through. Tests added for each. * Studio: tighten auto-switch code comments Comment/docstring-only pass over the OpenAI auto-switch feature: collapse multi-line blocks, drop a comment that restated the gate it sits next to, and trim verbose docstrings on internal helpers while keeping the load-bearing rationale (concurrency, API behavior, drop-in compat, gotchas). No logic change: verified comment-only with the AST/printer signature check. * Studio: bind auto-switch locks per running loop Review follow-up. The auto-switch swap lock and the keep-warm unload gate were module-level asyncio.Lock objects. That is safe under the single uvicorn loop and on Python 3.10+ (the Lock resolves the running loop lazily on acquire), but a module-level Lock binds to one loop on pre-3.10, which can raise a loop mismatch in multi-loop runners. Resolve each lock through a per-loop accessor backed by a WeakKeyDictionary so every running loop gets its own Lock and stale loops are collected. No behavior change under the server's single loop. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch re-review fixes (body codes, coverage, swap, alias, tracking) Follow-ups from a second review pass over the opt-in OpenAI auto-switch feature: 1. OFF-state status codes: /v1/completions and /v1/embeddings moved the body read ahead of the loaded-state check, so a malformed/empty body with no model loaded returned 500 instead of the prior 503. A shared helper reads the body defensively (an unparseable/non-dict body yields no model), and the handler re-reads after the 503 gate to surface the original parse error exactly as before. OFF behavior is unchanged. 2. Local-model coverage: the resolver index only scanned ./models and the active HF cache, while the model picker also lists the legacy/default HF caches, LM Studio dirs, and user scan folders. A request for one of those named models silently served the loaded model instead. _build_index now scans the same roots (Ollama's symlink-creating scanner is skipped on the request path), and resolution is offloaded with asyncio.to_thread so the wider scan never blocks the event loop. 3. Swap vs in-flight stream: a cross-model swap killed the llama-server while another client was still streaming from it. The hook now tracks how many requests are streaming on the loaded model (in-flight minus those still inside the hook) and returns 409 instead of swapping while one is active. Concurrent same-model requests never reach this path, so they are unaffected. 4. Idle-unload + alias: after idle-unload freed the model, an unknown/alias name resolved to nothing and 503'd, though it served the active model before the TTL. Idle-unload now remembers the freed id and an alias request reloads it (only an already-local model, so no remote download), cleared once a model is loaded again. 5. In-flight tracking: the keep-warm middleware tracked in-flight only while the feature was on, so a stream started while off could be unloaded if idle-unload was enabled mid-stream. It now tracks on every inference path; counting is cheap and invisible to clients. Tests added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove stray async_task_outputs files committed by mistake * Studio: auto-switch review round 3 (revert swap guard, hardening) Addressing a third review pass: - Revert the cross-model swap guard. It counted keep-warm in-flight (which includes external-provider calls that never touch the local model) and so could 409 a local swap spuriously, and it still left a same-model request able to start streaming on the model a concurrent swap was unloading. A correct fix needs a request-lifetime reader/writer barrier; a partial guard was worse than the honest single-slot behavior, so concurrent different-model use is back to being serialized (documented), like llama-swap's single slot. - Non-string request model (e.g. {"model": 123} on a raw-body endpoint) is now treated as absent, so it falls through instead of raising in the membership checks once an idle-unload stash exists. - Idle-unload now stashes and replays the freed quant: an alias reload restores the exact (id, variant) that was freed rather than the largest local quant. - Anthropic /v1/messages validates max_tokens before the auto-switch hook, so a request that 400s never triggers a model load. - Keep-warm tracks a pending count for requests waiting on the unload gate, so the idle loop cannot unload the model out from under a request that is blocked on the gate but not yet counted as in-flight. - The idle-unload task is awaited after cancel on shutdown to avoid pending-task warnings. - The resolver's HF cache scan is None-safe and logs at debug instead of letting a bad root abort the whole index build. - upsert_app_setting_map_entry rolls back explicitly on error. Tests updated/added for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: keep saved idle-unload seconds when auto-switch is toggled off * Studio: auto-switch hardening (thread-safe lock maps, body validation) Defensive fixes from review: - Guard the per-loop WeakKeyDictionary get-or-create for both the unload gate and the auto-switch lock with a threading lock, since WeakKeyDictionary mutation is not thread-safe when two event loops run on different threads. - Build the resolver index under the cache lock so concurrent callers with an expired cache don't all run the multi-dir scan at once. - /v1/completions and /v1/embeddings return a clean 400 for a valid JSON body that is not an object (e.g. a list), instead of a 500 from body.get(...). - The keep-warm middleware only tracks POST requests (inference is always POST), so CORS preflight (OPTIONS) is not counted, and tolerates a None path. Tests added for the list-body 400 and the non-POST skip. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 4 (local-path load, swap guard, idle fixes) From a 10-reviewer pass: - HF-cache entries now load by a concrete local path, not the bare repo id. The resolver records a load_path (the snapshot dir for a models--* cache repo, the file/dir otherwise) so /load takes the local branch and can never trigger a download to satisfy a partial cache. The advertised loader_id (repo id) is kept as the launch-override key. resolve_local_gguf now returns (load_path, variant, loader_id). - Re-add a single-slot swap guard: a cross-model swap returns 409 model_switch_busy while another inference request is active rather than killing its stream (the caller is excluded from the count), and holds the keep-warm gate across the load so no new inference starts mid-swap. Concurrent same-model requests never reach this path. A residual spurious 409 is possible while a concurrent or external- provider request is active; that is the documented single-slot tradeoff. - Idle keep-warm tracks (model_identifier, hf_variant): reloading the same repo at a different quant counts as a fresh model, so it is not unloaded before one TTL. - Track Studio's own /api/inference/generate/stream so the idle loop can't unload the model mid-stream on that route. - A successful manual /load clears the idle-unload reload stash synchronously, not only on the next idle poll. Also merged origin/main (the branch had fallen behind, which would have reverted unrelated files on merge). Tests added/updated for each. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 5 (concurrency, identity, load gate) From a 10-reviewer pass (9 request-changes, 1 approve): - Concurrent same-target requests load once instead of each returning 409. The count-based busy guard could not tell "another request wants the same model" (safe, load once) from "another request is using the loaded model" (refuse). Track in-flight auto-switch requests per (target, variant) and subtract same-target waiters from the busy count; a cross-model swap still 409s while a genuinely different request is active. - Fix the identity confusion introduced when round 4 began loading by concrete local path: the backend identifier became a filesystem path. Record the advertised repo id on the backend after an auto-switch load and use it so (a) a model loaded manually by repo id is recognized as already serving (no spurious reswap/409), (b) /v1/models reports the repo id, never a host path or a duplicate, and (c) the idle-unload stash keeps the override keyed by the repo id, so an alias reload after TTL keeps the user's saved launch flags. - Gate the manual /load route with the keep-warm lifecycle gate so idle auto-unload can't unload a model mid-load. load_model now wraps _load_model_impl in the gate; auto-switch calls _load_model_impl directly since it already holds the gate. - Restore default-off parity on Anthropic /v1/messages: an unloaded backend with auto-switch disabled 503s before the max_tokens 400 check, as it did pre-feature. When the feature is on, request-shape validation still runs before any load. Tests added for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 6 (concurrency ordering, leaks, unload gate) From a second 10-reviewer pass (8 request-changes, 2 approve): - Same-target concurrency: register a waiter by the raw requested model before the (slow) resolve, and exclude pending requests from the swap busy count. The middleware counts a concurrent same-model request as in-flight before it resolves and joins the resolved-target waiter map, so the prior fix could still 409 it. The guard now subtracts max(same resolved-target, same raw-request) waiters and ignores pending (a pending request is blocked in the middleware, not generating, so a swap can't interrupt it). - External-provider requests no longer block a local swap. The keep-warm middleware counts every inference-path POST, but external-provider chat returns before the auto-switch hook and never touches the local GGUF. The chat handler now untracks itself before proxying, so its in-flight stream can't trip model_switch_busy on a concurrent local auto-switch. The middleware skips its own end-decrement for an untracked request. - Manual /unload is gated like load and idle-unload: it holds the lifecycle gate and returns 409 rather than tearing down llama-server while an inference request is in flight. - Response model id no longer leaks the load path. /v1/models already advertised the repo id; chat, completions, embeddings, Anthropic messages, and audio response bodies now use the same _llama_public_model_id helper instead of the concrete on-disk model_identifier. - Chat completions validates the non-system-message requirement before the auto-switch hook (as /responses and /messages already do), so an invalid request can't swap the resident model before returning 400. Tests added for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch review round 7 (teardown policy, Unsloth-active swap, training) From a third 10-reviewer pass (9 request-changes, 1 approve), all on the same asymmetric-teardown theme. Resolved per the intended policy that only automatic paths defer to an active stream; deliberate user actions stay interrupting: - Revert the manual /unload in-flight guard added last round. A manual /load or /unload is a deliberate action and tears down immediately, as before; only the automatic idle-unload loop and auto-switch defer to an active request. This removes the asymmetry the reviewers flagged (manual /load, the /unload Unsloth branch, and the opposite-backend swaps inside _load_model_impl) by not extending the guard to deliberate paths, rather than spreading it. - Auto-switch now refuses a swap whenever another inference request is in flight, not only when a GGUF is already loaded. _load_model_impl also unloads an active Unsloth/transformers backend before loading a GGUF, so the busy guard must cover that case too; otherwise an Unsloth stream could be killed by an auto-switch. - Refuse API-initiated training while inference is active. When Studio is driven as an inference API (sk-unsloth key auth), POST /api/training/start returns 409 if a request is in flight, since training frees VRAM by unloading the chat model and would kill the stream. The Studio UI (session auth) still starts training and coexists/frees VRAM as before. A mixed UI+API session is not yet special-cased. Adds auth.authentication.authenticated_via_api_key. Tests added/updated for each; full backend suite diff vs baseline is unchanged. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: add UNSLOTH_MODEL_IDLE_TTL env override for idle-unload Borrowed from PR 6517: a startup env var that sets the idle-unload TTL without the settings UI. Unlike the stored setting (gated on auto-switch), the env value is a standalone default that enables idle-unload even with auto-switch off, for headless/container deploys. An explicit UI/API value still overrides it and stays gated. The settings GET reflects the env default when nothing is stored. * Studio: auto-switch fixes from review (paths, embeddings input, env idle reload) - /v1/models advertises a client-facing alias instead of a filesystem path: the ./models and LM Studio scanners report the on-disk path as the model id, so the index now prefers model_id/display_name as the advertised/override id and keeps the concrete path internal as load_path, still resolvable by path. - /v1/embeddings validates input before auto-switch: a request with a model but no input now 400s before the hook (like chat/responses/messages), so an invalid embeddings request cannot unload or swap the resident model. - Standalone UNSLOTH_MODEL_IDLE_TTL reloads the freed model: the hook now runs when auto-switch or idle-unload is active, and with auto-switch off it skips the resolver and only restores the idle-unloaded model, so the first idle timeout no longer leaves later /v1 requests with nothing loaded. - Do not resurrect a stale GGUF over an active Unsloth model: the reload-stash path bails when a non-GGUF backend is loaded, so an unknown /v1 name cannot tear down a live Transformers/Unsloth model. - Defensive HF cache scan: each cache root's resolve/dedup is wrapped so a missing or malformed root skips that root rather than aborting the index. - Single-model retrieve checks the id is a string before lowercasing. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix automatic-load asymmetry, audio reload, preview, idle timer The standalone UNSLOTH_MODEL_IDLE_TTL reload is a second automatic-load trigger, but several validate-before-switch guards and reload hooks only checked the auto-switch toggle. Add a shared _automatic_model_load_may_run() (auto-switch on, or idle TTL > 0) and route every guard through it. - /v1/completions validates prompt before any automatic load (it was the one model-bearing route with no pre-check). - /v1/chat/completions and /v1/embeddings pre-checks gate on the shared predicate so a standalone idle TTL cannot reload then reject. - /v1/messages no longer 503s before the reload hook can restore an idle-freed model when auto-switch is off. - Raw completions/embeddings with no model field pass a non-empty sentinel so the idle-stash reload runs, restoring the legacy "omit model, use loaded" path. - /api/inference/audio/generate gains the reload hook (after message validation) so an idle-freed audio GGUF is restored. - Public preview opts out of auto-switch via a request-scope flag, so a caller's model field cannot swap away from the pinned checkpoint; preview chat streams are now matched by _is_inference_path so idle-unload cannot kill them. - Keep-warm no longer stamps activity on request start, and external-provider untracking decrements without restamping, so periodic external traffic can no longer keep the local GGUF warm forever. Merges origin/main (the branch had fallen behind, which also brought in the preview route the review flagged). * Studio: surface model auto-switch in the API tab and demo it in examples The OpenAI auto-switch toggle previously lived only in Settings -> General. Add the same toggle to the API tab's usage-examples panel (it shares the settings cache), and make the examples reflect it: when on, the Python examples append a second call naming a different downloaded GGUF (so the model field visibly selects which model serves), and the curl examples gain a one-line note. Reuses the existing settings API client and i18n keys. * Studio: harden OpenAI auto-switch reload-only path and Anthropic tool validation - Omitted-model raw-body requests pass a reload-only sentinel so the idle-stash reload still restores an idle-freed model, but the resolver never matches a downloaded GGUF literally named "default". - Reject malformed Anthropic client tools before _maybe_auto_switch_model so an invalid request can no longer evict the loaded model. * Studio: extend auto-switch reload-only and tool validation to schema endpoints - Schema-backed endpoints (chat completions, responses, count_tokens, messages, audio) defaulted an omitted model to "default" and passed it to the switch hook, so a downloaded GGUF named "default" could be swapped to. Route the hook through a helper that switches only on an explicitly set model, else reload-only. - Propagate the explicit-set status when building the chat request from a Responses request, so the non-streaming chat re-check stays reload-only too. - Validate Responses function tools before the switch hook so a malformed tool returns 400 without evicting the loaded model. * Studio: serialize auto-switch swaps across event loops with a process-wide gate The auto-switch lock is a per-event-loop asyncio.Lock, so two /v1 swaps on different loops in one process could both pass it and race the single model slot (the backend and _load_model_impl are process-wide). Add a process-wide threading gate around the swap, acquired off the loop so a cross-loop wait never blocks it, layered with the existing per-loop lock. Add a cross-loop test that fails without the gate (two slow loads overlap) and passes with it. * Studio: make the auto-switch swap gate wait cancellation-safe _acquire_swap_gate awaited asyncio.to_thread(lock.acquire) when another loop held the process-wide gate. to_thread cancellation doesn't stop the worker thread, so a /v1 request cancelled mid-wait (client disconnect during a cross-loop swap) would have its thread acquire the gate after the fact, while the finally that releases it never runs -- permanently deadlocking later auto-switch swaps. Poll a non-blocking acquire off a short asyncio.sleep instead: it still keeps the wait off the loop and serializes across loops, but a cancel now lands during the sleep, when the gate is not held, so nothing leaks. Add a test that deadlocks the to_thread variant (it times out) and passes with the poll. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: validate modality and tool-confirmation before auto-switch Two more request shapes could load a named GGUF and only then 400, evicting the resident model: - An image request naming a different text-only GGUF. The switch hook now takes require_vision and rejects a swap to a non-vision target before loading it; a GGUF's vision capability is its companion mmproj, knowable without a load, and matches the post-load guard. Only the resolver branch is checked, never the reload-stash restore. - confirm_tool_calls=true with stream=false and local tools. /v1/chat/completions now rejects that shape before the hook, mirroring the local tool path's bypass_permissions exemption and intent signal. The vision probe threads the ambient HF token to keep the capability-probe invariant. Reload-only and idle-reload paths are unaffected. * Studio: extend validate-before-switch and make the lifecycle gate process-wide - /v1/messages/count_tokens now rejects malformed client tools before the switch hook, like /messages (shared _validate_anthropic_client_tools helper), so a count request can't evict the loaded model. - /v1/chat/completions rejects a malformed tool_choice forcing object (a {"type":"function","function":{}} with no name) before the switch hook. - The inference lifecycle gate that blocks new inference during a swap is now process-wide (a poll-acquired threading lock, cancellation-safe), not a per-loop asyncio lock, so a request on another event loop can't start inference while a swap tears the single backend down. - Usage examples no longer hard-code a switch-demo repo most users lack; the model is an explicit placeholder the user replaces. * Studio: extend the auto-switch modality guard to /v1/responses and /v1/messages The pre-load vision check that guards /v1/chat/completions now also runs on /v1/responses and /v1/messages, so an image request naming a text-only GGUF is rejected before the swap and never evicts the resident vision model. Run the vision capability probe off the event loop. Make the /v1/models retrieve loaded fast-path case-insensitive, and never advertise a host path from the resolver. Remove the dead list_switch_eligible_ids helper, superseded by the /v1/models catalog. * Studio: filter /v1/models to GGUF, per-loop catalog lock, reject system-only Responses Address review findings on the auto-switch path: - /v1/models advertises only GGUF models the API can actually switch to; a safetensors/LoRA entry would be selectable but never loadable via llama.cpp. - The /v1/models catalog cache uses a per-loop lock (like the auto-switch path) so a second event loop awaiting it can't hang in a multi-loop process. - /v1/responses rejects system/developer-only input before the switch, mirroring chat, so an invalid request can't evict the resident model. - _build_index guards each scan source on its own so one bad root drops only that source; the vision probe logs a real detection failure instead of swallowing it. * Studio: list cached GGUFs in /v1/models by inspecting files, not model_format The HF-cache scanner leaves model_format unset for GGUF snapshots, so the previous model_format == "gguf" filter dropped every downloaded HF-cache GGUF from /v1/models and the retrieve fallback. Decide GGUF-ness from the on-disk files via the resolver (info_has_local_gguf) instead, run off the event loop, so the catalog advertises exactly what /v1 can serve. * Studio: fix /v1/messages/count_tokens route binding plus auto-switch review fixes The @router.post decorator for /messages/count_tokens had been separated from anthropic_count_tokens by the _validate_anthropic_client_tools helper, so the route bound to the validator and dropped its auth dependency. Move the decorator back onto the handler. Add route-binding tests asserting each /v1 endpoint maps to its handler with the auth dependency, so a decorator/handler split is caught at the route level (the direct-call tests missed it). Also from review: - update_openai_auto_switch writes both settings keys in one transaction so a PUT can't leave one updated and the other stale (drop the now-unused single setters). - max_seq_length override rejects 0 at the boundary (ge=1) instead of accepting then silently dropping it. - Document that embeddings auto-switch is best-effort: GGUF pooling has no cheap pre-load probe like vision's mmproj, so a guard would false-reject GGUF embedders. - Add a positive idle-unload test (loop frees the model and stashes it for reload). * Studio: validate Responses tool_choice + Anthropic mixed tools before switch, filter Ollama from catalog More auto-switch review findings: - /v1/responses rejects a forcing-function tool_choice with no name before the switch, mirroring chat, so a malformed request can't evict the resident model. - /v1/messages rejects mixing Anthropic server tools with custom client tools before the switch (the check depends only on the payload, so it moves up cleanly). - /v1/models no longer advertises Ollama-link models: info_has_local_gguf excludes .studio_links / ollama_links entries, which the resolver skips and can't switch to, so an advertised id never silently falls through. * Studio: guard chat audio input before switch; surface env-backed idle unload in settings UI A chat request carrying audio_base64 rides the same companion mmproj projector as a vision request, so a text-only target cannot serve it either. Flag require_vision for audio input as well so the multimodal probe runs before the switch and a rejected request never evicts the working model. Generalize the reject message to cover image and audio. The settings response now reports idle_unload_active (effective TTL > 0) so the UI can distinguish idle-unload that is active via the UNSLOTH_MODEL_IDLE_TTL env var from the case where it needs the toggle enabled. * Studio: harden auto-switch eviction guards (count_tokens vision, TTS reload-only, mmproj/stash) Four eviction/correctness fixes on the opt-in /v1 auto-switch path: - /v1/messages/count_tokens now carries the same require_vision guard as /messages, so an image count naming a text-only GGUF can't evict a loaded vision model for a swap that can't serve the request. - /audio/generate is now reload-only. A local GGUF's audio-input capability is not a cheap pre-load probe (the companion mmproj signal can't tell an audio projector from a vision one, and codec TTS ships no projector), so resolving the client model could load a text/vision-only target and evict the working audio model before the audio check fails. Only the idle-stash restore runs here; switching TTS models is an explicit /load. - The resolver no longer treats a standalone mmproj .gguf as a servable model. _scan_models_dir's standalone-file pass does not filter mmproj the way its directory scan does, so /v1/models could advertise a projector and a switch could load it over the real weights. - A non-GGUF (Transformers/Unsloth) load and a deliberate /unload now clear the idle reload stash, so a manual load/unload is never superseded by a stale idle-freed GGUF that the next /v1 request resurrects. * Studio: report advertised repo id consistently after an auto-switch Two model-id reporting fixes so an auto-switched cached HF GGUF is named by its repo id everywhere, not its snapshot path: - Streamed /v1/responses envelopes now derive the model id from _llama_public_model_id (which prefers _openai_advertised_id) instead of the raw model_identifier. After an auto-switch the identifier is the snapshot path while the repo id lives in _openai_advertised_id, so the stream used to report a snapshot basename while /v1/models, chat completions, and non-streaming Responses all reported the repo id. - When an advertised alias already resolves to the loaded model (a model loaded by local path, requested by its repo or LM Studio id), the already-serving early return now records the alias as the advertised id, so /v1/models and responses report the alias and mark it loaded instead of the path-derived basename. Resolver branch only; safe lock-free because an in-flight request blocks any concurrent swap via the single-slot busy guard. * Studio: validate request shapes before auto-switch (prompt/input/audio/mcp confirm) Four more validate-before-switch guards so a deterministic client error never evicts the resident model on the opt-in /v1 auto-switch path: - /v1/completions rejects an object/number prompt (only a string or array is valid) before the switch, instead of loading the named GGUF and letting llama-server reject the shape afterward. - /v1/embeddings rejects an object/number input the same way. - Chat rejects an oversized audio_base64 upload (413) before the switch. The size cap is a cheap, target-independent length check; the decode itself stays post-switch to avoid decoding a valid upload twice. - The chat confirm-without-stream pre-switch guard now mirrors the tool loop's actual enablement: _effective_enable_tools (honoring a CLI --enable-tools policy) and mcp_enabled (which opens the tool loop on its own but defers to a CLI --disable-tools policy). Previously a confirm+no-stream request with only mcp_enabled slipped past and 400'd after the swap. * Studio: fix model-id retrieval, streaming n>1, resolver cache TTL, keep-warm auth Four fixes from review: - GET /v1/models/{id} legacy raw-path fallback now maps the raw identifier to the same public id its /v1/models entry uses. After an auto-switch load the identifier is the snapshot path while the entry is keyed by the advertised repo id, so a client that cached the old absolute path no longer 404s on a model that is in fact loaded. - stream=true with n>1 is now rejected before the switch. Only the non-streaming GGUF path returns multiple choices, so streaming n>1 is invalid on every local serving path; both fields are known pre-switch, so it must not load model B only to 400 and evict model A. Non-streaming n>1 stays post-switch where the serving path decides. - The resolver index cache is stamped after _build_index, not with the pre-scan timestamp. On installs with enough local models for the multi-root scan to exceed the 5s TTL, the cache was stored already expired and every request rebuilt it. - The keep-warm middleware no longer stamps model activity for 401/403 responses. It runs before FastAPI auth, so unauthenticated probes used to refresh the idle timer without touching llama.cpp; they now decrement the in-flight count without keeping the model warm. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
1338 lines
51 KiB
Python
1338 lines
51 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Main FastAPI application for Unsloth UI Backend
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import threading
|
|
from pathlib import Path as _Path
|
|
import asyncio
|
|
from dataclasses import asdict
|
|
|
|
# Suppress C-level dependency warnings globally
|
|
os.environ["PYTHONWARNINGS"] = "ignore"
|
|
|
|
# Pin GPU index ordering to PCI bus id before any torch import creates a CUDA
|
|
# context. Without this, torch/CUDA default to FASTEST_FIRST while nvidia-smi
|
|
# (and Studio's VRAM probes) use PCI-bus order, so a GPU index chosen from
|
|
# nvidia-smi data can resolve to a different physical card via
|
|
# CUDA_VISIBLE_DEVICES. setdefault so an explicit user override wins. See
|
|
# utils/hardware/hardware.py for the full rationale; set here too so the entry
|
|
# process is covered before its heavy ML imports.
|
|
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
|
|
|
|
# Windows terminals default to the active system code page. Reconfigure
|
|
# stdout/stderr before the startup banner so non-ASCII output cannot crash the
|
|
# backend process.
|
|
if sys.platform == "win32":
|
|
for _win_stream in (sys.stdout, sys.stderr):
|
|
if _win_stream is not None and hasattr(_win_stream, "reconfigure"):
|
|
try:
|
|
_win_stream.reconfigure(encoding = "utf-8", errors = "replace")
|
|
except Exception:
|
|
pass
|
|
del _win_stream
|
|
|
|
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
|
|
# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
|
|
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
|
|
if sys.platform == "win32":
|
|
# Retained at module scope; os.add_dll_directory returns a handle that
|
|
# removes the search-path entry when garbage collected.
|
|
_ROCM_DLL_HANDLES: list = []
|
|
|
|
def _add_rocm_dll_dirs() -> None:
|
|
candidates = []
|
|
# 1. HIP_PATH / ROCM_PATH set by the AMD HIP SDK installer
|
|
for _var in ("HIP_PATH", "ROCM_PATH"):
|
|
_val = os.environ.get(_var)
|
|
if _val:
|
|
candidates.append(os.path.join(_val, "bin"))
|
|
# 2. AMD installer: C:\Program Files\AMD\ROCm\<ver>\bin, newest first.
|
|
_default_root = os.path.join(
|
|
os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
|
|
)
|
|
|
|
def _ver_key(name: str) -> tuple:
|
|
# Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string
|
|
parts = []
|
|
for chunk in name.split("."):
|
|
try:
|
|
parts.append((0, int(chunk)))
|
|
except ValueError:
|
|
parts.append((1, chunk))
|
|
return tuple(parts)
|
|
|
|
try:
|
|
if os.path.isdir(_default_root):
|
|
for _ver in sorted(os.listdir(_default_root), key = _ver_key, reverse = True):
|
|
_bin = os.path.join(_default_root, _ver, "bin")
|
|
if os.path.isdir(_bin):
|
|
candidates.append(_bin)
|
|
except OSError:
|
|
pass
|
|
for _d in candidates:
|
|
if os.path.isdir(_d):
|
|
try:
|
|
_ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
|
|
except (OSError, AttributeError):
|
|
pass
|
|
|
|
_add_rocm_dll_dirs()
|
|
del _add_rocm_dll_dirs
|
|
|
|
# ── Windows AMD ROCm: make hipInfo.exe resolvable for subprocess probes ──
|
|
# bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via PATH at import
|
|
# time; the AMD torch wheel ships it in the venv Scripts dir, which is on
|
|
# PATH only when the venv is activated -- Studio launches python directly.
|
|
# Without this, every bitsandbytes import logs a scary (but harmless)
|
|
# "Could not detect ROCm GPU architecture: [WinError 2]" ERROR + WARNING.
|
|
# Gated on the file existing: only AMD ROCm wheels ship hipInfo.exe, so
|
|
# NVIDIA/CPU hosts are untouched. os.add_dll_directory above does not help
|
|
# here -- subprocess PATH resolution ignores DLL search directories.
|
|
_scripts_dir = os.path.dirname(sys.executable)
|
|
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")):
|
|
import shutil as _shutil
|
|
if not _shutil.which("hipinfo.exe"):
|
|
os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "")
|
|
del _shutil
|
|
del _scripts_dir
|
|
|
|
# ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
|
|
# bitsandbytes derives the rocm<ver>.dll name from torch.version.hip, but the
|
|
# wheel ships rocm72.dll, so the server crashes ("Configured ROCm binary not
|
|
# found") without this. Detect the shipped DLL (mirrors worker.py); gate on
|
|
# the rocm bnb DLL rather than torch.version.hip to avoid importing torch on
|
|
# every Windows host.
|
|
# Values seeded by the installer's sitecustomize.py are redetectable
|
|
# defaults; explicit caller values remain authoritative.
|
|
if (
|
|
"BNB_ROCM_VERSION" not in os.environ
|
|
or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"
|
|
):
|
|
import glob as _glob
|
|
import logging as _logging
|
|
|
|
_bnb_rocm_ver = None
|
|
_found_rocm_bnb = False
|
|
try:
|
|
import importlib.util as _ilu
|
|
_bnb_spec = _ilu.find_spec("bitsandbytes")
|
|
# submodule_search_locations (not spec.origin) handles editable installs
|
|
if _bnb_spec and _bnb_spec.submodule_search_locations:
|
|
import re as _re_bnb
|
|
|
|
_all_vers_main: list[str] = []
|
|
for _pkg_dir in _bnb_spec.submodule_search_locations:
|
|
for _dll in _glob.glob(os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")):
|
|
_found_rocm_bnb = True
|
|
_km = _re_bnb.search(
|
|
r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll)
|
|
)
|
|
if _km:
|
|
_all_vers_main.append(_km.group(1))
|
|
if _all_vers_main:
|
|
_bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v))
|
|
except Exception as _e:
|
|
_logging.getLogger(__name__).warning(
|
|
"Windows ROCm: BNB DLL detection failed (%s); leaving BNB_ROCM_VERSION as is",
|
|
_e,
|
|
)
|
|
# Only when a ROCm bnb DLL actually exists: HIP_PATH/ROCM_PATH alone
|
|
# (HIP SDK on a CUDA/CPU box) must not force a ROCm backend onto a
|
|
# non-ROCm bitsandbytes, which raises at import. DLL unparsable -> "72".
|
|
if _found_rocm_bnb:
|
|
_bnb_rocm_ver_final = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
|
|
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
|
|
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
|
|
_logging.getLogger(__name__).info(
|
|
"Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)",
|
|
_bnb_rocm_ver_final,
|
|
)
|
|
|
|
# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ──────
|
|
# In WSL the AMD GPU is reached via the ROCDXG bridge (librocdxg.so over
|
|
# /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_DETECTION=1 is set BEFORE
|
|
# torch touches the GPU. A worker launched outside a login shell (e.g.
|
|
# `wsl.exe -d Ubuntu-24.04 python ...`) misses the installer's persisted env
|
|
# and silently falls back to CPU. Set it here, gated to no-op unless BOTH
|
|
# /dev/dxg AND librocdxg.so exist -- native Linux ROCm, NVIDIA, macOS and
|
|
# Windows are unaffected.
|
|
elif sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ:
|
|
try:
|
|
if os.path.exists("/dev/dxg") and any(
|
|
os.path.exists(os.path.join(_p, "librocdxg.so"))
|
|
for _p in ("/opt/rocm/lib", "/opt/rocm/lib64")
|
|
):
|
|
os.environ["HSA_ENABLE_DXG_DETECTION"] = "1"
|
|
import logging as _logging
|
|
_logging.getLogger(__name__).info(
|
|
"WSL ROCm: set HSA_ENABLE_DXG_DETECTION=1 (librocdxg bridge present)"
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# Put backend dir on sys.path so _platform_compat is importable when main.py
|
|
# is launched directly (e.g. `uvicorn main:app`).
|
|
_backend_dir = str(_Path(__file__).parent)
|
|
if _backend_dir not in sys.path:
|
|
sys.path.insert(0, _backend_dir)
|
|
|
|
# `uvicorn main:app` bypasses run.py; seed thread caps here too.
|
|
from utils.cpu_threads import configure_cpu_threads
|
|
|
|
try:
|
|
configure_cpu_threads()
|
|
except ValueError as exc:
|
|
_raw = os.environ.get("UNSLOTH_CPU_THREADS")
|
|
raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}") from None
|
|
|
|
# Anaconda/conda-forge Python: seed platform._sys_version_cache before any
|
|
# library import triggers attrs -> rich -> structlog -> platform crash.
|
|
# See: https://github.com/python/cpython/issues/102396
|
|
import _platform_compat # noqa: F401
|
|
|
|
# Direct `uvicorn main:app` launches bypass run.py, so re-export here too
|
|
# (mirrors run.py). Required BEFORE the unsloth-zoo import below, whose
|
|
# LLAMA_CPP_DEFAULT_DIR binding is import-time.
|
|
from utils.paths.storage_roots import studio_root as _studio_root
|
|
|
|
try:
|
|
_LEGACY_STUDIO_ROOT = (_Path.home() / ".unsloth" / "studio").resolve()
|
|
except (OSError, ValueError):
|
|
_LEGACY_STUDIO_ROOT = _Path.home() / ".unsloth" / "studio"
|
|
try:
|
|
_STUDIO_ROOT_RESOLVED = _studio_root().resolve()
|
|
except (OSError, ValueError):
|
|
_STUDIO_ROOT_RESOLVED = _studio_root()
|
|
if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
|
|
if not os.environ.get("UNSLOTH_STUDIO_HOME"):
|
|
os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED)
|
|
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
|
|
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp")
|
|
|
|
# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth`
|
|
# does) so its lazy submodule imports (export, hardware, mlx) and the
|
|
# DiffusionGemma runner never trip the install guard on a clean install.
|
|
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
|
|
|
|
import hashlib
|
|
import mimetypes
|
|
import re as _re
|
|
import shutil
|
|
import warnings
|
|
from contextlib import asynccontextmanager
|
|
from importlib.metadata import PackageNotFoundError, version as package_version
|
|
from typing import Optional
|
|
from urllib.parse import urlparse
|
|
|
|
|
|
_STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$")
|
|
|
|
|
|
def _read_studio_install_id() -> str:
|
|
"""Per-install opaque id at $STUDIO_HOME/share/studio_install_id.
|
|
|
|
Returns "" when absent or not a 64-char lowercase-hex token; then
|
|
/api/health emits "" and the launcher accepts any healthy backend.
|
|
Carries no install-path info (matters when Studio runs -H 0.0.0.0)."""
|
|
try:
|
|
token = (_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
|
|
except (OSError, ValueError):
|
|
return ""
|
|
return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else ""
|
|
|
|
|
|
_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()
|
|
|
|
|
|
def _studio_root_id() -> str:
|
|
"""Same-install discriminator for /api/health (cached at import).
|
|
|
|
Empty when no installer token is present; the launcher treats "" as
|
|
"accept any healthy backend"."""
|
|
return _STUDIO_ROOT_ID_CACHE
|
|
|
|
|
|
# Fix broken Windows registry MIME types: some installs map .js to text/plain,
|
|
# which mimetypes (hence StaticFiles) inherits and browsers reject for ES
|
|
# modules. add_type() before StaticFiles forces correct types.
|
|
if sys.platform == "win32":
|
|
mimetypes.add_type("application/javascript", ".js")
|
|
mimetypes.add_type("text/css", ".css")
|
|
|
|
# Suppress dependency warnings in production
|
|
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
|
|
warnings.filterwarnings("ignore")
|
|
# Or be more specific:
|
|
# warnings.filterwarnings("ignore", category=DeprecationWarning)
|
|
# warnings.filterwarnings("ignore", module="triton.*")
|
|
|
|
from fastapi import Depends, FastAPI, HTTPException, Query, Request
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.responses import FileResponse, HTMLResponse, Response
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
|
|
from routes import (
|
|
auth_router,
|
|
chat_history_router,
|
|
data_recipe_router,
|
|
datasets_router,
|
|
export_router,
|
|
inference_router,
|
|
inference_studio_router,
|
|
mcp_servers_router,
|
|
models_router,
|
|
providers_router,
|
|
rag_router,
|
|
training_history_router,
|
|
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,
|
|
)
|
|
from hub.schemas.downloads import TransportCapabilities
|
|
from hub.utils.download_registry import (
|
|
get_download_transport_capabilities,
|
|
reap_orphan_workers as reap_hub_orphan_workers,
|
|
terminate_active_downloads as terminate_hub_downloads,
|
|
)
|
|
from routes.settings import router as settings_router
|
|
from routes.prompts import router as prompts_router
|
|
from auth import storage
|
|
from auth.authentication import get_current_subject
|
|
from utils.hardware import (
|
|
detect_hardware,
|
|
get_device,
|
|
DeviceType,
|
|
get_backend_visible_gpu_info,
|
|
)
|
|
import utils.hardware.hardware as _hw_module
|
|
|
|
from utils.cache_cleanup import clear_unsloth_compiled_cache
|
|
from utils.lifespan_shutdown import run_lifespan_shutdown
|
|
from utils.native_path_leases import native_path_leases_supported
|
|
from utils.update_status import (
|
|
get_studio_install_source_status,
|
|
get_studio_update_status,
|
|
)
|
|
from utils.studio_version import get_studio_version
|
|
from utils.api_errors import install_api_error_handlers
|
|
|
|
|
|
def get_unsloth_version() -> str:
|
|
try:
|
|
return package_version("unsloth")
|
|
except PackageNotFoundError:
|
|
pass
|
|
|
|
version_file = _Path(__file__).resolve().parents[2] / "unsloth" / "models" / "_utils.py"
|
|
try:
|
|
for line in version_file.read_text(encoding = "utf-8").splitlines():
|
|
if line.startswith("__version__ = "):
|
|
return line.split("=", 1)[1].strip().strip('"').strip("'")
|
|
except OSError:
|
|
pass
|
|
return "dev"
|
|
|
|
|
|
UNSLOTH_VERSION = get_unsloth_version()
|
|
STUDIO_VERSION = get_studio_version()
|
|
|
|
|
|
def _load_desktop_owner() -> dict[str, str] | None:
|
|
token = os.environ.pop("UNSLOTH_STUDIO_DESKTOP_OWNER_TOKEN", "")
|
|
kind = os.environ.pop("UNSLOTH_STUDIO_DESKTOP_OWNER_KIND", "")
|
|
if kind != "tauri" or not token:
|
|
return None
|
|
return {
|
|
"kind": "tauri",
|
|
"token_sha256": hashlib.sha256(token.encode("utf-8")).hexdigest(),
|
|
}
|
|
|
|
|
|
_DESKTOP_OWNER = _load_desktop_owner()
|
|
|
|
# The Tauri desktop app runs the backend on the owner's own machine, so local
|
|
# stdio MCP servers are safe there. setdefault lets an explicit "0" opt out.
|
|
if _DESKTOP_OWNER:
|
|
os.environ.setdefault("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
|
|
|
|
|
|
def _desktop_owner() -> dict[str, str] | None:
|
|
return _DESKTOP_OWNER
|
|
|
|
|
|
def _start_helper_precache_if_enabled() -> None:
|
|
"""Start optional Helper LLM GGUF pre-cache only after explicit opt-in."""
|
|
try:
|
|
from utils.helper_precache_settings import should_preload_helper_on_startup
|
|
if not should_preload_helper_on_startup():
|
|
return
|
|
except Exception:
|
|
return
|
|
|
|
import threading
|
|
|
|
def _precache():
|
|
try:
|
|
from utils.datasets.llm_assist import precache_helper_gguf
|
|
precache_helper_gguf()
|
|
except Exception:
|
|
pass # non-critical
|
|
|
|
threading.Thread(target = _precache, daemon = True, name = "helper-gguf-precache").start()
|
|
|
|
|
|
def _run_llama_cpp_startup_probes(app: FastAPI) -> None:
|
|
"""llama.cpp capability (MTP support) + freshness (release age) probes.
|
|
|
|
Runs OFF the startup critical path (see _start_llama_cpp_probes_if_enabled).
|
|
Both are cached and freshness has a 24h disk TTL, but on a cold/expired cache
|
|
the freshness check makes a blocking GitHub request, and on macOS the first
|
|
`llama-server --help` exec can stall on Gatekeeper verification -- neither must
|
|
ever gate `Application startup complete`. Writes app.state only; nothing reads
|
|
those values synchronously at startup (the status routes call
|
|
check_prebuilt_freshness directly at request time), so populating them late is
|
|
safe.
|
|
"""
|
|
try:
|
|
from core.inference.llama_cpp import LlamaCppBackend
|
|
from utils.llama_cpp_freshness import (
|
|
check_prebuilt_freshness,
|
|
format_stale_warning,
|
|
)
|
|
|
|
_bin = LlamaCppBackend._find_llama_server_binary()
|
|
_caps = LlamaCppBackend.probe_server_capabilities(_bin)
|
|
app.state.llama_cpp_capabilities = _caps
|
|
_freshness = check_prebuilt_freshness(_bin)
|
|
app.state.llama_cpp_freshness = _freshness
|
|
|
|
import structlog as _structlog
|
|
|
|
_log = _structlog.get_logger(__name__)
|
|
if _caps.get("found") and not _caps.get("supports_mtp"):
|
|
_msg = (
|
|
"llama.cpp prebuilt lacks MTP support "
|
|
"(--spec-type mtp/draft-mtp). Run `unsloth studio update`. "
|
|
"MTP GGUFs will load without speculative decoding."
|
|
)
|
|
_log.warning(_msg)
|
|
print(f"WARNING: {_msg}", flush = True)
|
|
if _freshness.get("stale"):
|
|
_msg = format_stale_warning(_freshness)
|
|
_log.warning(_msg)
|
|
print(f"WARNING: {_msg}", flush = True)
|
|
except Exception as _probe_exc:
|
|
import structlog as _structlog
|
|
_structlog.get_logger(__name__).debug("llama.cpp startup probes failed: %s", _probe_exc)
|
|
|
|
|
|
def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None:
|
|
"""Run the llama.cpp startup probes on a daemon thread, off the startup
|
|
critical path so they never delay `Application startup complete`. Skipped
|
|
entirely when update checks are disabled, so a fully offline boot makes no
|
|
background network calls."""
|
|
if os.environ.get("UNSLOTH_DISABLE_UPDATE_CHECK") == "1":
|
|
return
|
|
|
|
threading.Thread(
|
|
target = _run_llama_cpp_startup_probes,
|
|
args = (app,),
|
|
daemon = True,
|
|
name = "llama-cpp-startup-probe",
|
|
).start()
|
|
|
|
|
|
def _warm_rag_embedder() -> None:
|
|
"""Warm RAG embeddings without blocking backend readiness."""
|
|
try:
|
|
from storage import rag_db
|
|
|
|
if not rag_db.RAG_AVAILABLE:
|
|
return
|
|
from core.rag import embeddings
|
|
|
|
embeddings.warm()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
|
|
|
|
import time as _time
|
|
|
|
_lifespan_started = _time.perf_counter()
|
|
import structlog as _structlog
|
|
|
|
_lifespan_log = _structlog.get_logger(__name__)
|
|
clear_unsloth_compiled_cache()
|
|
|
|
# Remove stale .venv_overlay from old versions; switching now uses .venv_t5/.
|
|
overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay"
|
|
if overlay_dir.is_dir():
|
|
shutil.rmtree(overlay_dir, ignore_errors = True)
|
|
|
|
# Detect hardware first — sets the DEVICE global used everywhere.
|
|
detect_hardware()
|
|
|
|
_lifespan_log.info(
|
|
"lifespan hardware detection completed in %.1fms",
|
|
(_time.perf_counter() - _lifespan_started) * 1000,
|
|
)
|
|
|
|
# Apple Silicon with MLX missing => Train/Export are greyed out (chat-only).
|
|
# Reinstall mlx by name on a background thread (off the critical path) and
|
|
# re-detect, so a reinstall/update that dropped mlx self-heals. No-op
|
|
# elsewhere; opt out with UNSLOTH_DISABLE_MLX_AUTOREPAIR=1.
|
|
try:
|
|
from utils.mlx_repair import start_mlx_autorepair_if_needed
|
|
start_mlx_autorepair_if_needed()
|
|
except Exception as _mlx_exc:
|
|
import structlog as _structlog
|
|
_structlog.get_logger(__name__).debug("mlx autorepair skipped: %s", _mlx_exc)
|
|
|
|
# Reap workers/runs orphaned by a previous crash before new work starts.
|
|
try:
|
|
from storage.studio_db import cleanup_orphaned_runs
|
|
cleanup_orphaned_runs()
|
|
except Exception as exc:
|
|
_lifespan_log.warning("cleanup_orphaned_runs failed at startup: %s", exc)
|
|
|
|
reap_hub_orphan_workers()
|
|
|
|
# llama.cpp probes: capability (MTP support) + freshness (release age).
|
|
# These used to run inline here and could block `Application startup complete`
|
|
# for tens of seconds on macOS (cold GitHub freshness cache / slow network, and
|
|
# Gatekeeper verifying the unsigned binary on first `--help` exec). They only
|
|
# write app.state and nothing reads it synchronously at startup, so run them on
|
|
# a daemon thread off the startup critical path (mirrors the helper-precache and
|
|
# RAG-warm threads). Default to None until the thread populates them.
|
|
app.state.llama_cpp_capabilities = None
|
|
app.state.llama_cpp_freshness = None
|
|
_start_llama_cpp_probes_if_enabled(app)
|
|
|
|
try:
|
|
from storage.rag_db import reconcile_orphaned_ingestion_jobs
|
|
reconcile_orphaned_ingestion_jobs()
|
|
except Exception as exc:
|
|
_lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc)
|
|
|
|
_start_helper_precache_if_enabled()
|
|
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
|
|
|
|
# Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set).
|
|
from core.inference.llama_keepwarm import idle_unload_loop
|
|
|
|
app.state.idle_unload_task = asyncio.create_task(idle_unload_loop())
|
|
|
|
# Initialize RSA key pair for API key encryption (external providers).
|
|
from core.inference.key_exchange import init_key_pair
|
|
|
|
init_key_pair()
|
|
_lifespan_log.info(
|
|
"lifespan pre-auth setup completed in %.1fms",
|
|
(_time.perf_counter() - _lifespan_started) * 1000,
|
|
)
|
|
|
|
if storage.ensure_default_admin():
|
|
bootstrap_pw = storage.get_bootstrap_password()
|
|
app.state.bootstrap_password = bootstrap_pw
|
|
|
|
bootstrap_path = storage.DB_PATH.parent / ".bootstrap_password"
|
|
print("\n" + "=" * 60)
|
|
print("DEFAULT ADMIN ACCOUNT CREATED")
|
|
print(f" username: {storage.DEFAULT_ADMIN_USERNAME}")
|
|
print(f" password saved to: {bootstrap_path}")
|
|
print(" Open the Studio UI to sign in and change it.")
|
|
print("=" * 60 + "\n")
|
|
else:
|
|
app.state.bootstrap_password = storage.get_bootstrap_password()
|
|
|
|
_lifespan_log.info(
|
|
"lifespan startup completed in %.1fms",
|
|
(_time.perf_counter() - _lifespan_started) * 1000,
|
|
)
|
|
yield
|
|
|
|
_idle_task = getattr(app.state, "idle_unload_task", None)
|
|
if _idle_task is not None:
|
|
_idle_task.cancel()
|
|
try:
|
|
await _idle_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
from core.inference.llama_http import aclose as _close_llama_http
|
|
|
|
await _close_llama_http()
|
|
|
|
await run_lifespan_shutdown(
|
|
terminate_hub_downloads,
|
|
clear_unsloth_compiled_cache,
|
|
_hw_module,
|
|
)
|
|
|
|
|
|
app = FastAPI(
|
|
title = "Unsloth UI Backend",
|
|
version = UNSLOTH_VERSION,
|
|
description = "Backend API for Unsloth UI - Training and Model Management",
|
|
lifespan = lifespan,
|
|
)
|
|
|
|
from loggers.config import LogConfig
|
|
from loggers.handlers import LoggingMiddleware
|
|
|
|
logger = LogConfig.setup_logging(
|
|
service_name = "unsloth-studio-backend",
|
|
env = os.getenv("ENVIRONMENT_TYPE", "production"),
|
|
)
|
|
|
|
app.add_middleware(LoggingMiddleware)
|
|
|
|
|
|
# img/media-src allow any https origin so HF model-card assets render (mirrors
|
|
# tauri.conf.json); scripts/frames/connect-src stay same-origin + HF.
|
|
from starlette.datastructures import MutableHeaders # noqa: E402
|
|
|
|
|
|
_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
|
|
_ARTIFACT_PREVIEW_FRAME_PATH = "/api/inference/artifact-preview-frame"
|
|
|
|
|
|
# /content is Colab's working directory — more reliable than env vars, which
|
|
# aren't always set depending on Colab runtime version.
|
|
import importlib.util as _importlib_util
|
|
|
|
_IS_COLAB = os.path.isdir("/content") and (
|
|
bool(os.environ.get("COLAB_BACKEND_URL"))
|
|
or bool(os.environ.get("COLAB_JUPYTER_IP"))
|
|
or _importlib_util.find_spec("google.colab") is not None
|
|
)
|
|
|
|
|
|
def _build_csp(script_nonce: "str | None" = None) -> str:
|
|
script_src = "script-src 'self'"
|
|
if script_nonce:
|
|
script_src += f" 'nonce-{script_nonce}'"
|
|
# Colab parent frames span multi-level *.prod.colab.dev subdomains (CSP
|
|
# wildcards match one level only) and null-origin iframes; use '*' since
|
|
# Colab is already a sandboxed single-user environment.
|
|
frame_ancestors = "*" if _IS_COLAB else "'none'"
|
|
|
|
# In Colab, the kernel/output scaffolding injects scripts and fetch/WS from
|
|
# *.prod.colab.dev and *.googleusercontent.com, so widen script-src and
|
|
# connect-src for those. Scripts still use a nonce, not 'unsafe-inline'.
|
|
if _IS_COLAB:
|
|
script_src += " https://*.prod.colab.dev https://*.googleusercontent.com"
|
|
connect_src = (
|
|
"'self' blob: data: "
|
|
"https://huggingface.co https://datasets-server.huggingface.co "
|
|
"https://*.prod.colab.dev wss://*.prod.colab.dev "
|
|
"https://*.googleusercontent.com wss://*.googleusercontent.com"
|
|
)
|
|
else:
|
|
connect_src = "'self' https://huggingface.co https://datasets-server.huggingface.co"
|
|
|
|
return (
|
|
"default-src 'self'; "
|
|
"img-src 'self' data: blob: https:; "
|
|
"media-src 'self' data: blob: https:; "
|
|
f"connect-src {connect_src}; "
|
|
"style-src 'self' 'unsafe-inline'; "
|
|
f"{script_src}; "
|
|
"font-src 'self' data:; "
|
|
"frame-src 'self'; "
|
|
f"frame-ancestors {frame_ancestors}; "
|
|
"form-action 'self'; "
|
|
"base-uri 'self'"
|
|
)
|
|
|
|
|
|
class SecurityHeadersMiddleware:
|
|
"""Set baseline security headers; splice per-response inline-script nonces into CSP.
|
|
|
|
Pure ASGI (not BaseHTTPMiddleware) so streaming responses are not wrapped in
|
|
an anyio stream. Header logic mirrors the prior version exactly via
|
|
MutableHeaders on the response-start message.
|
|
"""
|
|
|
|
def __init__(self, app):
|
|
self.app = app
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
path = scope.get("path", "")
|
|
|
|
async def send_wrapper(message):
|
|
if message["type"] == "http.response.start":
|
|
# ASGI headers are an iterable; coerce to a list so MutableHeaders
|
|
# can mutate in place even if a server sends a tuple or omits it.
|
|
raw = message.setdefault("headers", [])
|
|
if not isinstance(raw, list):
|
|
raw = list(raw)
|
|
message["headers"] = raw
|
|
headers = MutableHeaders(raw = raw)
|
|
# Strip the internal nonce hand-off header so it never reaches the client
|
|
nonce = headers.get(_CSP_SCRIPT_NONCE_HEADER)
|
|
if nonce is not None:
|
|
del headers[_CSP_SCRIPT_NONCE_HEADER]
|
|
headers.setdefault("Content-Security-Policy", _build_csp(nonce))
|
|
# Omit X-Frame-Options in Colab: CSP frame-ancestors handles it, and
|
|
# DENY would block serve_kernel_port_as_iframe regardless of CSP.
|
|
if not _IS_COLAB and path != _ARTIFACT_PREVIEW_FRAME_PATH:
|
|
headers.setdefault("X-Frame-Options", "DENY")
|
|
headers.setdefault("X-Content-Type-Options", "nosniff")
|
|
headers.setdefault("Referrer-Policy", "no-referrer")
|
|
headers.setdefault(
|
|
"Permissions-Policy",
|
|
"camera=(), microphone=(self), geolocation=()",
|
|
)
|
|
headers["server"] = "unsloth-studio"
|
|
await send(message)
|
|
|
|
await self.app(scope, receive, send_wrapper)
|
|
|
|
|
|
app.add_middleware(SecurityHeadersMiddleware)
|
|
|
|
|
|
# Cap request bodies on protected POSTs. Upload routes get explicit multipart
|
|
# headroom; non-upload routes keep the default body cap.
|
|
import json as _json_for_413 # noqa: E402
|
|
from utils.upload_limits import ( # noqa: E402
|
|
UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES,
|
|
default_request_body_limit_bytes,
|
|
upload_request_limit_bytes,
|
|
)
|
|
|
|
_BODY_PROTECTED_PREFIXES = (
|
|
"/v1/chat/completions",
|
|
"/v1/completions",
|
|
"/p/",
|
|
"/api/inference",
|
|
"/api/data-recipe",
|
|
"/api/datasets",
|
|
"/api/hub",
|
|
"/api/chat",
|
|
"/api/settings",
|
|
"/api/train",
|
|
"/api/export",
|
|
)
|
|
_DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload"
|
|
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = (
|
|
"/api/data-recipe/seed/upload-unstructured-file"
|
|
)
|
|
_BODY_UPLOAD_PASSTHROUGH_PREFIXES = (
|
|
_DATASET_UPLOAD_PASSTHROUGH_PREFIX,
|
|
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX,
|
|
)
|
|
|
|
|
|
def _get_upload_passthrough_request_max_bytes(path: str) -> int:
|
|
if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX):
|
|
return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES)
|
|
if path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX):
|
|
return upload_request_limit_bytes()
|
|
return default_request_body_limit_bytes()
|
|
|
|
|
|
async def _send_411(send) -> None:
|
|
payload = _json_for_413.dumps(
|
|
{"detail": "Content-Length required for upload requests."},
|
|
).encode("utf-8")
|
|
await send(
|
|
{
|
|
"type": "http.response.start",
|
|
"status": 411,
|
|
"headers": [
|
|
(b"content-type", b"application/json"),
|
|
(b"content-length", str(len(payload)).encode("ascii")),
|
|
],
|
|
}
|
|
)
|
|
await send({"type": "http.response.body", "body": payload, "more_body": False})
|
|
|
|
|
|
async def _send_413(send, total_bytes: int, max_bytes: int) -> None:
|
|
payload = _json_for_413.dumps(
|
|
{"detail": (f"Request body too large ({total_bytes:,} bytes; max {max_bytes:,}).")},
|
|
).encode("utf-8")
|
|
await send(
|
|
{
|
|
"type": "http.response.start",
|
|
"status": 413,
|
|
"headers": [
|
|
(b"content-type", b"application/json"),
|
|
(b"content-length", str(len(payload)).encode("ascii")),
|
|
],
|
|
}
|
|
)
|
|
await send({"type": "http.response.body", "body": payload, "more_body": False})
|
|
|
|
|
|
class MaxBodyMiddleware:
|
|
"""Reject oversized bodies on protected POST/PUT/PATCH; raw ASGI so chunked uploads cannot bypass the cap."""
|
|
|
|
def __init__(
|
|
self,
|
|
app,
|
|
max_bytes_getter,
|
|
protected_prefixes: tuple,
|
|
upload_passthrough_prefixes: tuple = (),
|
|
upload_passthrough_max_bytes_getter = None,
|
|
):
|
|
self.app = app
|
|
self.max_bytes_getter = max_bytes_getter
|
|
self.protected_prefixes = protected_prefixes
|
|
self.upload_passthrough_prefixes = upload_passthrough_prefixes
|
|
self.upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter
|
|
|
|
def _upload_passthrough_max_bytes(self, path: str) -> int:
|
|
if self.upload_passthrough_max_bytes_getter is None:
|
|
return int(self.max_bytes_getter())
|
|
try:
|
|
return int(self.upload_passthrough_max_bytes_getter(path))
|
|
except TypeError:
|
|
try:
|
|
return int(self.upload_passthrough_max_bytes_getter())
|
|
except Exception:
|
|
return int(self.max_bytes_getter())
|
|
except Exception:
|
|
return int(self.max_bytes_getter())
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
method = scope.get("method", "").upper()
|
|
path = scope.get("path", "")
|
|
if method not in ("POST", "PUT", "PATCH") or not any(
|
|
path.startswith(p) for p in self.protected_prefixes
|
|
):
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
max_bytes = int(self.max_bytes_getter())
|
|
declared = None
|
|
for name, value in scope.get("headers", []):
|
|
if name == b"content-length":
|
|
try:
|
|
declared = int(value.decode("latin-1"))
|
|
except (ValueError, UnicodeDecodeError):
|
|
declared = None
|
|
break
|
|
|
|
if any(path.startswith(p) for p in self.upload_passthrough_prefixes):
|
|
upload_max_bytes = self._upload_passthrough_max_bytes(path)
|
|
if declared is None:
|
|
await _send_411(send)
|
|
return
|
|
if declared > upload_max_bytes:
|
|
await _send_413(send, declared, upload_max_bytes)
|
|
return
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
if declared is not None and declared > max_bytes:
|
|
await _send_413(send, declared, max_bytes)
|
|
return
|
|
|
|
chunks: list = []
|
|
total = 0
|
|
while True:
|
|
msg = await receive()
|
|
mtype = msg.get("type")
|
|
if mtype == "http.disconnect":
|
|
return
|
|
if mtype != "http.request":
|
|
# Mid-stream unexpected frame: forwarding would corrupt downstream
|
|
return
|
|
body = msg.get("body", b"") or b""
|
|
if body:
|
|
total += len(body)
|
|
if total > max_bytes:
|
|
await _send_413(send, total, max_bytes)
|
|
return
|
|
chunks.append(body)
|
|
if not msg.get("more_body", False):
|
|
break
|
|
|
|
replayed = {"sent": False}
|
|
|
|
async def replay_receive():
|
|
if not replayed["sent"]:
|
|
replayed["sent"] = True
|
|
return {
|
|
"type": "http.request",
|
|
"body": b"".join(chunks),
|
|
"more_body": False,
|
|
}
|
|
# After replay, fall through so http.disconnect still propagates.
|
|
return await receive()
|
|
|
|
await self.app(scope, replay_receive, send)
|
|
|
|
|
|
app.add_middleware(
|
|
MaxBodyMiddleware,
|
|
max_bytes_getter = default_request_body_limit_bytes,
|
|
protected_prefixes = _BODY_PROTECTED_PREFIXES,
|
|
upload_passthrough_prefixes = _BODY_UPLOAD_PASSTHROUGH_PREFIXES,
|
|
upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes,
|
|
)
|
|
|
|
# Tracks in-flight inference requests for idle auto-unload; off -> passthrough.
|
|
from core.inference.llama_keepwarm import LlamaKeepWarmMiddleware # noqa: E402
|
|
|
|
app.add_middleware(LlamaKeepWarmMiddleware)
|
|
|
|
|
|
from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402
|
|
|
|
|
|
@app.get("/recipes", include_in_schema = False)
|
|
@app.get("/recipes/{rest:path}", include_in_schema = False)
|
|
async def _recipes_redirect(rest: str = ""):
|
|
target = "/data-recipes" + (("/" + rest) if rest else "")
|
|
return _RedirectResponse(url = target, status_code = 308)
|
|
|
|
|
|
from utils.host_policy import cors_origins_for_mode # noqa: E402
|
|
|
|
_cors_origins = cors_origins_for_mode(
|
|
api_only = os.environ.get("UNSLOTH_API_ONLY") == "1",
|
|
secure = os.environ.get("UNSLOTH_SECURE") == "1",
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins = _cors_origins,
|
|
allow_credentials = True,
|
|
allow_methods = ["*"],
|
|
allow_headers = ["*"],
|
|
)
|
|
|
|
|
|
# ============ Register API Routes ============
|
|
|
|
# Register routers
|
|
app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"])
|
|
app.include_router(training_router, prefix = "/api/train", tags = ["training"])
|
|
app.include_router(models_router, prefix = "/api/models", tags = ["models"])
|
|
app.include_router(chat_history_router, prefix = "/api/chat", tags = ["chat"])
|
|
app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
|
|
# Studio-only inference endpoints (cancel, etc.) are NOT exposed on the /v1
|
|
# OpenAI-compat prefix below.
|
|
app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"])
|
|
|
|
# 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"])
|
|
app.include_router(prompts_router, prefix = "/api/prompts", tags = ["prompts"])
|
|
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
|
|
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
|
|
app.include_router(llama_router, prefix = "/api/llama", tags = ["llama"])
|
|
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
|
|
app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
|
|
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])
|
|
app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"])
|
|
app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"])
|
|
|
|
# Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic
|
|
# error envelopes; non-/v1 paths keep FastAPI's default {"detail": ...} shape.
|
|
install_api_error_handlers(app)
|
|
|
|
|
|
# ============ Health and System Endpoints ============
|
|
|
|
|
|
@app.get("/api/liveness")
|
|
async def liveness_check():
|
|
"""Cheap process liveness for desktop port validation."""
|
|
return {
|
|
"status": "alive",
|
|
"service": "Unsloth UI Backend",
|
|
"desktop_protocol_version": 1,
|
|
"desktop_manageability_version": 1,
|
|
"supports_desktop_auth": True,
|
|
"supports_desktop_backend_ownership": True,
|
|
"studio_root_id": _studio_root_id(),
|
|
**({"desktop_owner": owner} if (owner := _desktop_owner()) else {}),
|
|
}
|
|
|
|
|
|
@app.get("/api/health")
|
|
async def health_check(request: Request):
|
|
"""Liveness plus launcher capability bits; host fingerprint gated on a bearer.
|
|
|
|
Unauthenticated callers get non-sensitive fields (service, studio_root_id,
|
|
chat_only, desktop_*, native_path_leases_supported) to re-adopt a sibling
|
|
backend and gate UI before a token exists. version / studio_version /
|
|
device_type require a bearer since they fingerprint the host.
|
|
"""
|
|
base = {
|
|
"status": "healthy",
|
|
"timestamp": datetime.now().isoformat(),
|
|
"service": "Unsloth UI Backend",
|
|
"chat_only": _hw_module.CHAT_ONLY,
|
|
"desktop_protocol_version": 1,
|
|
"desktop_manageability_version": 1,
|
|
"supports_desktop_auth": True,
|
|
"supports_desktop_backend_ownership": True,
|
|
# Opaque per-install id; launchers reject sibling Studios on the same port.
|
|
"studio_root_id": _studio_root_id(),
|
|
"native_path_leases_supported": native_path_leases_supported(),
|
|
**({"desktop_owner": owner} if (owner := _desktop_owner()) else {}),
|
|
}
|
|
auth = request.headers.get("authorization", "")
|
|
if not auth.lower().startswith("bearer "):
|
|
return base
|
|
try:
|
|
from auth.authentication import get_current_subject as _gcs
|
|
from fastapi.security import HTTPAuthorizationCredentials
|
|
|
|
creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = auth.split(" ", 1)[1])
|
|
# Must await: a bare coroutine is truthy and would skip the auth check
|
|
subject = await _gcs(creds)
|
|
except HTTPException:
|
|
return base
|
|
except Exception:
|
|
return base
|
|
if not subject:
|
|
return base
|
|
|
|
platform_map = {"darwin": "mac", "win32": "windows", "linux": "linux"}
|
|
device_type = platform_map.get(sys.platform, sys.platform)
|
|
return {
|
|
**base,
|
|
# Why chat_only is set. This fingerprints the host, so keep it authed.
|
|
"chat_only_reason": getattr(_hw_module, "CHAT_ONLY_REASON", None),
|
|
"version": UNSLOTH_VERSION,
|
|
"studio_version": STUDIO_VERSION,
|
|
"device_type": device_type,
|
|
# API-screen fields (authed-only; they fingerprint how the host is exposed).
|
|
"cloudflare_url": getattr(request.app.state, "cloudflare_url", None),
|
|
"server_url": getattr(request.app.state, "server_url", None),
|
|
"secure": bool(getattr(request.app.state, "secure", False)),
|
|
}
|
|
|
|
|
|
@app.get("/api/studio/install-source")
|
|
def studio_install_source(_current_subject: str = Depends(get_current_subject)):
|
|
"""Return source-aware install metadata without remote update checks."""
|
|
return get_studio_install_source_status(UNSLOTH_VERSION)
|
|
|
|
|
|
@app.get("/api/studio/update-status")
|
|
def studio_update_status(_current_subject: str = Depends(get_current_subject)):
|
|
"""Return source-aware manual update status for browser-served Studio."""
|
|
return get_studio_update_status(UNSLOTH_VERSION)
|
|
|
|
|
|
@app.get(
|
|
"/api/studio/download-transport-capabilities",
|
|
response_model = TransportCapabilities,
|
|
)
|
|
def studio_download_transport_capabilities(_current_subject: str = Depends(get_current_subject)):
|
|
return asdict(get_download_transport_capabilities())
|
|
|
|
|
|
@app.post("/api/shutdown")
|
|
async def shutdown_server(request: Request, current_subject: str = Depends(get_current_subject)):
|
|
"""Gracefully shut down the Unsloth Studio server.
|
|
|
|
Called by the frontend quit dialog so users can stop the server from the UI
|
|
without the CLI or killing the process manually.
|
|
"""
|
|
|
|
async def _delayed_shutdown():
|
|
await asyncio.sleep(0.2) # Let the HTTP response return first
|
|
trigger = getattr(request.app.state, "trigger_shutdown", None)
|
|
if trigger is not None:
|
|
trigger()
|
|
else:
|
|
# Fallback when not launched via run_server() (e.g. direct uvicorn)
|
|
import signal
|
|
import os
|
|
os.kill(os.getpid(), signal.SIGTERM)
|
|
|
|
request.app.state._shutdown_task = asyncio.create_task(_delayed_shutdown())
|
|
return {"status": "shutting_down"}
|
|
|
|
|
|
@app.get("/api/system")
|
|
async def get_system_info(current_subject: str = Depends(get_current_subject)):
|
|
"""Get system information.
|
|
|
|
Auth-gated: the response (platform, Python/GPU, memory, ML packages) can
|
|
fingerprint a host, which matters in -H 0.0.0.0 / Colab / Tauri-relayed
|
|
setups where remote callers can reach /api/system.
|
|
"""
|
|
import platform
|
|
import psutil
|
|
from utils.hardware import get_device
|
|
from utils.hardware.hardware import _backend_label
|
|
|
|
visibility_info = get_backend_visible_gpu_info()
|
|
gpu_info = {
|
|
"available": visibility_info["available"],
|
|
"devices": visibility_info["devices"],
|
|
}
|
|
|
|
# CPU & Memory
|
|
memory = psutil.virtual_memory()
|
|
|
|
return {
|
|
"platform": platform.platform(),
|
|
"python_version": platform.python_version(),
|
|
# _backend_label so /api/system reports "rocm" (not "cuda") on AMD,
|
|
# matching /api/hardware and /api/gpu-visibility.
|
|
"device_backend": _backend_label(get_device()),
|
|
"cpu_count": psutil.cpu_count(),
|
|
"memory": {
|
|
"total_gb": round(memory.total / 1e9, 2),
|
|
"available_gb": round(memory.available / 1e9, 2),
|
|
"percent_used": memory.percent,
|
|
},
|
|
"gpu": gpu_info,
|
|
}
|
|
|
|
|
|
@app.get("/api/system/gpu-visibility")
|
|
async def get_gpu_visibility(current_subject: str = Depends(get_current_subject)):
|
|
return get_backend_visible_gpu_info()
|
|
|
|
|
|
@app.get("/api/system/hardware")
|
|
def get_hardware_info(
|
|
include_details: bool = Query(False), current_subject: str = Depends(get_current_subject)
|
|
):
|
|
"""Return GPU name, total VRAM, and key ML package versions.
|
|
|
|
Gated behind auth alongside /api/system -- same fingerprinting concern.
|
|
/api/system/gpu-visibility is also auth-gated.
|
|
|
|
``include_details`` is for About/diagnostics. The default response stays
|
|
cheap for callers that only need the primary GPU summary, like training
|
|
method auto-selection. Sync def (not async): hardware/detail probes can
|
|
shell out, and FastAPI runs sync endpoints in a threadpool.
|
|
"""
|
|
from utils.hardware import get_gpu_summary, get_package_versions
|
|
|
|
body = {
|
|
"gpu": get_gpu_summary(),
|
|
"versions": get_package_versions(),
|
|
}
|
|
if include_details:
|
|
from utils.llama_cpp_update import get_installed_llama_version
|
|
|
|
# All backend-visible GPUs (respects CUDA_VISIBLE_DEVICES), so multi-GPU
|
|
# hosts list every device -- get_gpu_summary alone reports only the primary.
|
|
# Sort by visible_ordinal: the nvidia-smi path returns rows in physical order,
|
|
# so under a reordering CUDA_VISIBLE_DEVICES (e.g. "5,3") labeling by array
|
|
# index would otherwise disagree with the GPU 0/1 the backend actually sees.
|
|
devices = get_backend_visible_gpu_info().get("devices", [])
|
|
body["gpus"] = [
|
|
{"name": d.get("name"), "vram_total_gb": d.get("memory_total_gb")}
|
|
for d in sorted(devices, key = lambda d: d.get("visible_ordinal", 0))
|
|
]
|
|
body["llama_cpp"] = get_installed_llama_version()
|
|
return body
|
|
|
|
|
|
# ============ Serve Frontend (Optional) ============
|
|
|
|
|
|
def _strip_crossorigin(html_bytes: bytes) -> bytes:
|
|
"""Remove ``crossorigin`` attributes from script/link tags.
|
|
|
|
Vite's default ``crossorigin`` forces CORS mode on font loads, which
|
|
Firefox HTTPS-Only Mode breaks over plain HTTP; stripping it makes them
|
|
same-origin fetches that work on any protocol.
|
|
"""
|
|
html = html_bytes.decode("utf-8")
|
|
html = _re.sub(r'\s+crossorigin(?:="[^"]*")?', "", html)
|
|
return html.encode("utf-8")
|
|
|
|
|
|
def _inject_bootstrap(html_bytes: bytes, app: FastAPI):
|
|
"""Inject bootstrap credentials when password change is pending.
|
|
Returns ``(html_bytes, script_nonce_or_None)``; callers forward the nonce
|
|
via ``_CSP_SCRIPT_NONCE_HEADER`` so CSP allows the inline script.
|
|
"""
|
|
import json as _json
|
|
import secrets as _secrets
|
|
|
|
if not storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME):
|
|
return html_bytes, None
|
|
|
|
bootstrap_pw = getattr(app.state, "bootstrap_password", None)
|
|
if not bootstrap_pw:
|
|
return html_bytes, None
|
|
|
|
payload = _json.dumps(
|
|
{
|
|
"username": storage.DEFAULT_ADMIN_USERNAME,
|
|
"password": bootstrap_pw,
|
|
}
|
|
)
|
|
nonce = _secrets.token_urlsafe(16)
|
|
tag = f'<script nonce="{nonce}">window.__UNSLOTH_BOOTSTRAP__={payload}</script>'
|
|
html = html_bytes.decode("utf-8")
|
|
html = html.replace("</head>", f"{tag}</head>", 1)
|
|
return html.encode("utf-8"), nonce
|
|
|
|
|
|
_DEFAULT_PORTS = {"http": 80, "https": 443, "ws": 80, "wss": 443}
|
|
|
|
|
|
def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]]:
|
|
"""Canonicalise an Origin to ``(scheme, host, port)`` for equality.
|
|
Browsers strip default ports (RFC 6454 sec 6.1) and scheme/host are
|
|
case-insensitive (RFC 3986), so a bare string compare misclassifies
|
|
same-origin requests as cross-origin. Returns ``None`` on unparseable input
|
|
so callers fall to the safer cross-origin default.
|
|
"""
|
|
scheme = (scheme or "").strip().lower()
|
|
if not scheme or not netloc:
|
|
return None
|
|
# Strip userinfo (RFC 3986); Origin never carries credentials.
|
|
if "@" in netloc:
|
|
netloc = netloc.rsplit("@", 1)[1]
|
|
# IPv6 hosts use brackets (RFC 3986 sec 3.2.2): ``[::1]:8902``. Bare
|
|
# ``partition(":")`` mis-parses these, breaking ``unsloth studio -H ::1``.
|
|
if netloc.startswith("["):
|
|
close = netloc.find("]")
|
|
if close == -1:
|
|
return None
|
|
host = netloc[1:close]
|
|
rest = netloc[close + 1 :]
|
|
if rest.startswith(":"):
|
|
port_str = rest[1:]
|
|
elif rest == "":
|
|
port_str = ""
|
|
else:
|
|
return None
|
|
else:
|
|
host, _, port_str = netloc.partition(":")
|
|
host = host.strip().lower()
|
|
if not host:
|
|
return None
|
|
if port_str:
|
|
try:
|
|
port = int(port_str)
|
|
except ValueError:
|
|
return None
|
|
else:
|
|
port = _DEFAULT_PORTS.get(scheme, 0)
|
|
return (scheme, host, port)
|
|
|
|
|
|
def _is_same_origin_request(request: Request) -> bool:
|
|
"""True when Origin is missing or matches request's scheme://host:port.
|
|
|
|
Missing Origin counts as same-origin (top-level GETs omit it). Both sides
|
|
are canonicalised via :func:`_canonical_origin`; callers must emit
|
|
``Vary: Origin``.
|
|
"""
|
|
origin = request.headers.get("origin")
|
|
if origin is None:
|
|
# Missing header: top-level same-document GETs omit Origin.
|
|
return True
|
|
# Empty string is not a valid serialised origin (RFC 6454 sec 6.1).
|
|
if not origin:
|
|
return False
|
|
# "null" token (sandboxed iframes, file:// pages) is never same-origin.
|
|
if origin == "null":
|
|
return False
|
|
# ``urlparse`` raises ``ValueError`` on malformed IPv6 brackets; swallow
|
|
# so a garbage Origin doesn't 500 the SPA handler.
|
|
try:
|
|
parsed = urlparse(origin)
|
|
except ValueError:
|
|
return False
|
|
origin_canon = _canonical_origin(parsed.scheme, parsed.netloc)
|
|
if origin_canon is None:
|
|
return False
|
|
try:
|
|
self_canon = _canonical_origin(request.url.scheme, request.url.netloc)
|
|
except ValueError:
|
|
return False
|
|
if self_canon is None:
|
|
return False
|
|
return origin_canon == self_canon
|
|
|
|
|
|
def setup_frontend(app: FastAPI, build_path: Path):
|
|
"""Mount frontend static files (optional)"""
|
|
if not build_path.exists():
|
|
return False
|
|
|
|
assets_dir = build_path / "assets"
|
|
if assets_dir.exists():
|
|
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
|
|
|
|
def _build_index_response(request: Request) -> Response:
|
|
content = (build_path / "index.html").read_bytes()
|
|
content = _strip_crossorigin(content)
|
|
# Bootstrap pw is same-origin only; Vary: Origin keeps caches honest.
|
|
if _is_same_origin_request(request):
|
|
content, nonce = _inject_bootstrap(content, app)
|
|
else:
|
|
nonce = None
|
|
headers = {
|
|
"Cache-Control": "no-cache, no-store, must-revalidate",
|
|
"Vary": "Origin",
|
|
}
|
|
if nonce:
|
|
headers[_CSP_SCRIPT_NONCE_HEADER] = nonce
|
|
return Response(
|
|
content = content,
|
|
media_type = "text/html",
|
|
headers = headers,
|
|
)
|
|
|
|
@app.get("/")
|
|
async def serve_root(request: Request):
|
|
return _build_index_response(request)
|
|
|
|
@app.get("/{full_path:path}")
|
|
async def serve_frontend(request: Request, full_path: str):
|
|
# Unknown API paths: raise a real 404 so the api_errors handlers can
|
|
# render the correct envelope for /v1/* (and {"detail":...} for /api/*).
|
|
# This handler only sees paths NOT matched by a real route. The full
|
|
# request path is "/" + full_path.
|
|
if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")):
|
|
raise HTTPException(status_code = 404, detail = "API endpoint not found")
|
|
|
|
file_path = (build_path / full_path).resolve()
|
|
|
|
# Block path traversal — resolved path must stay inside build_path
|
|
if not file_path.is_relative_to(build_path.resolve()):
|
|
return Response(status_code = 403)
|
|
|
|
if file_path.is_file():
|
|
return FileResponse(file_path)
|
|
|
|
# Serve index.html as bytes — avoids Content-Length mismatch
|
|
return _build_index_response(request)
|
|
|
|
return True
|