unsloth/studio/backend/main.py
Daniel Han 5211b506e1
Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm (#6392)
* 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>
2026-07-01 06:42:23 -07:00

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