ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm (#6392)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-07-01 06:42:23 -07:00 • committed by GitHub
commit 5211b506e1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 5243 additions and 109 deletions
Download patch file Download diff file Expand all files Collapse all files

11
studio/backend/auth/authentication.py
View file

@ -143,6 +143,17 @@ async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depend
)
async def authenticated_via_api_key(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> bool:
"""True when the caller used an sk-unsloth API key, not a UI session JWT.
Lets routes treat programmatic API callers differently from the Studio UI
(e.g. refuse a teardown the UI would allow).
"""
return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX))
async def get_current_subject_allow_password_change(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:

298
studio/backend/core/inference/llama_keepwarm.py Normal file
View file

@ -0,0 +1,298 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Opt-in idle auto-unload (TTL keep-warm) for the local llama.cpp model.
Off by default (idle seconds = 0). When enabled, a background loop unloads the
loaded GGUF once it has been idle for the configured TTL, freeing VRAM. A
pure-ASGI middleware tracks in-flight inference requests so a long stream that
outlives the TTL is never unloaded mid-response.
"""
from __future__ import annotations
import asyncio
import contextlib
import threading
import time
from loggers import get_logger
logger = get_logger(__name__)
_lock = threading.Lock()
_inflight = 0
# Requests blocked on the unload gate but not yet counted in _inflight: the idle
# loop must not unload while one is waiting (it would unload out from under it).
_pending = 0
_last_active = time.monotonic()
# The (id, quant) idle-unload last freed, so an alias/unknown request that would
# otherwise 503 against an empty backend can reload it (set on unload, cleared on
# reload). Storing the quant means the reload restores the exact freed variant.
_last_unloaded_model = None
# Guards inflight bumps against the idle-check-then-unload race, and blocks new
# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is
# shared across every event loop in the process, so a per-loop gate would let a
# request on loop B start inference while a swap on loop A tears the model down.
_lifecycle_lock = threading.Lock()
@contextlib.asynccontextmanager
async def _unload_gate():
# Acquire off the loop: non-blocking first (the common uncontended case), else
# poll a non-blocking acquire off a short sleep. Polling keeps the wait off this
# loop AND cancellation-safe -- a cancel lands during the sleep, when the gate is
# not held, so it never leaks (mirrors the auto-switch swap gate).
while not _lifecycle_lock.acquire(blocking = False):
await asyncio.sleep(0.02)
try:
yield
finally:
_lifecycle_lock.release()
_INFERENCE_PREFIXES = ("/v1/", "/api/inference/")
_INFERENCE_SUFFIXES = (
"/chat/completions",
"/completions",
"/messages",
"/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages
"/embeddings",
"/responses",
"/generate/stream", # Studio's own streaming route on the same llama-server
"/audio/generate", # direct GGUF TTS; can outlive the idle TTL
)
def _is_inference_path(path: str) -> bool:
if path.startswith(_INFERENCE_PREFIXES) and path.endswith(_INFERENCE_SUFFIXES):
return True
# Public checkpoint preview (/p/{run}/v1/chat/completions) delegates to the
# chat handler and streams from the same backend, so protect it from idle unload.
return path.startswith("/p/") and path.endswith("/v1/chat/completions")
def _note_pending() -> None:
global _pending
with _lock:
_pending += 1
def _note_unpending() -> None:
global _pending
with _lock:
_pending = max(0, _pending - 1)
def _note_start() -> None:
# Do not stamp _last_active here: while _inflight > 0 the model is already
# protected (see _is_idle), and stamping on start lets an external-provider
# request that is later untracked still reset the local idle timer.
global _inflight, _pending
with _lock:
_pending = max(0, _pending - 1)
_inflight += 1
def _note_end() -> None:
global _inflight, _last_active
with _lock:
_inflight = max(0, _inflight - 1)
_last_active = time.monotonic()
def _note_untracked_end() -> None:
# Drop a request that never used the local GGUF without stamping local
# activity, so periodic external-provider traffic can't keep the model warm.
global _inflight
with _lock:
_inflight = max(0, _inflight - 1)
def _is_idle(ttl_seconds: float) -> bool:
with _lock:
return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds
def _note_activity() -> None:
"""Stamp activity, e.g. on a (re)load, so the model survives at least one TTL."""
global _last_active
with _lock:
_last_active = time.monotonic()
def other_inference_request_count(
current_request_counted: bool = True, *, include_pending: bool = True
) -> int:
"""Tracked inference requests other than the current route call.
The middleware counts OpenAI-compatible requests before route code runs, so
the caller is excluded by default. Idle-unload counts pending waiters too (a
swap holding the gate would unload out from under them). The swap guard passes
include_pending=False: a pending request is blocked in the middleware and has
not started inference, so it can't be the request a swap would interrupt.
"""
with _lock:
active = _inflight
if current_request_counted and active > 0:
active -= 1
return max(0, active) + (_pending if include_pending else 0)
# Set on the ASGI scope by a route that proved this request won't touch
# llama.cpp (e.g. it proxied to an external provider), so the keep-warm count
# excludes it and the middleware skips its own end-decrement.
_UNTRACKED_SCOPE_KEY = "_unsloth_keepwarm_untracked"
def untrack_current_request(scope) -> None:
"""Drop this request from the in-flight count once the route knows it won't
use the local GGUF, so unrelated external-provider traffic can't trip the
swap busy guard. Idempotent; the middleware then skips its end-decrement."""
if not isinstance(scope, dict) or scope.get(_UNTRACKED_SCOPE_KEY):
return
scope[_UNTRACKED_SCOPE_KEY] = True
_note_untracked_end()
def inference_lifecycle_gate():
"""The gate a model swap holds so new inference can't start mid-load. Process-
wide, so a swap on one loop blocks inference starting on any other loop."""
return _unload_gate()
def note_model_loaded() -> None:
"""Record a successful GGUF load: stamp activity and drop any reload stash so
a manual load clears it synchronously, not only on the next idle poll."""
_note_activity()
_set_last_unloaded(None)
def note_model_unloaded() -> None:
"""Record a deliberate (user/API) unload: drop any idle reload stash so the next
request can't resurrect the just-unloaded model. The idle loop unloads via the
backend directly and then stashes the freed model for an alias reload; an
explicit unload instead means "stay unloaded", so it must not stamp activity."""
_set_last_unloaded(None)
def get_last_unloaded_model():
with _lock:
return _last_unloaded_model
def _set_last_unloaded(value) -> None:
global _last_unloaded_model
with _lock:
_last_unloaded_model = value
class LlamaKeepWarmMiddleware:
"""Pure ASGI: count in-flight inference requests and stamp activity on completion."""
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
# Inference endpoints are all POST; skipping non-POST avoids counting CORS
# preflight (OPTIONS). ``or ""`` guards an explicit None path.
if (
scope.get("type") != "http"
or scope.get("method") != "POST"
or not _is_inference_path(scope.get("path") or "")
):
await self.app(scope, receive, send)
return
# Always track in-flight on inference paths, even when the feature is off,
# so a stream that starts before idle-unload is enabled can't be unloaded
# mid-response if the operator turns it on during that stream. Counting is
# cheap and invisible to clients (the response is proxied unchanged).
# Mark pending before the gate so the idle loop (which holds the gate while
# unloading) can't free the model while this request is waiting to start.
_note_pending()
started = False
try:
async with _unload_gate():
_note_start()
started = True
finally:
if not started:
_note_unpending()
ended = {"done": False}
status = {"code": None}
def _finish() -> None:
# A route that untracked itself already decremented; don't double-count.
if ended["done"]:
return
ended["done"] = True
if scope.get(_UNTRACKED_SCOPE_KEY):
return
# This middleware runs before FastAPI auth, so a 401/403 reaches here
# without ever touching llama.cpp. Decrement the in-flight count (to
# balance _note_start) but do NOT stamp activity, or repeated
# unauthenticated probes on an exposed server would keep the model warm
# and never let idle-unload free VRAM.
if status["code"] in (401, 403):
_note_untracked_end()
else:
_note_end()
async def send_wrapper(message):
if message.get("type") == "http.response.start":
status["code"] = message.get("status")
# Final body frame marks the end of a (possibly streaming) response.
elif message.get("type") == "http.response.body" and not message.get(
"more_body", False
):
_finish()
await send(message)
try:
await self.app(scope, receive, send_wrapper)
finally:
_finish()
def _loaded_identity(backend):
if not backend.is_loaded or not backend.model_identifier:
return None
# Third slot is the advertised id (repo id) an auto-switch load sets on the
# backend; it's the override key, so an idle stash keyed by the concrete load
# path doesn't drop the user's saved launch flags on the alias reload.
advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier
return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised)
async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
"""Unload the loaded GGUF once idle past the configured TTL. Inert when off."""
from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds
seen_model = None
while True:
await asyncio.sleep(poll_seconds)
try:
ttl = get_auto_unload_idle_seconds()
if ttl <= 0:
continue
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
# Track by (id, variant): a (re)loaded model -- including the same repo
# at a different quant -- counts as activity so it survives one TTL
# before its first request (loads bypass the activity middleware).
current = _loaded_identity(backend)
if current != seen_model:
seen_model = current
if current is not None:
_note_activity()
_set_last_unloaded(None) # a model is loaded; drop stale stash
async with _unload_gate():
if backend.is_loaded and _is_idle(ttl):
freed = _loaded_identity(backend)
await asyncio.to_thread(backend.unload_model)
_set_last_unloaded(freed) # let an alias request reload it
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)
seen_model = None
except Exception as exc:
logger.debug("idle_unload_loop iteration failed: %s", exc)

269
studio/backend/core/inference/local_model_resolver.py Normal file
View file

@ -0,0 +1,269 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Resolve an OpenAI-request ``model`` string to a downloaded local GGUF.
Used by the opt-in auto-switch path. The match is conservative: only names
that map to an already-downloaded local GGUF (and a quant that is actually on
disk) are eligible, so an arbitrary OpenAI model string still falls through to
the loaded model (drop-in compat) and no surprise multi-GB download is ever
triggered. The local-model scan is cached for a few seconds since auto-switch
consults it per request.
"""
from __future__ import annotations
import threading
import time
from dataclasses import dataclass
from typing import Optional
from core.inference.model_ids import public_model_id
from loggers import get_logger
logger = get_logger(__name__)
@dataclass(frozen = True)
class _LocalGgufEntry:
loader_id: str # advertised id (repo id / folder name), also the override key
load_path: str # concrete on-disk dir/file passed to /load so it never downloads
variants: tuple[str, ...] # local quant labels; () for a standalone .gguf
_CACHE_TTL_S = 5.0
_lock = threading.Lock()
_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {})
def _is_abs_path_id(value: str) -> bool:
"""True when an id is an absolute filesystem path (the ./models and LM Studio
scanners use the on-disk path as the id) rather than a repo id like org/name."""
from pathlib import Path
try:
return Path(value).is_absolute()
except Exception:
return False
def _advertised_loader_id(info) -> Optional[str]:
"""The id to advertise for a scanned model: prefer a client-facing alias over
an absolute filesystem path so /v1/models and the override key never expose a
host path (the ./models and LM Studio scanners report the path as info.id)."""
raw_id = getattr(info, "id", None)
if not raw_id or not _is_abs_path_id(raw_id):
return raw_id
for alt in (getattr(info, "model_id", None), getattr(info, "display_name", None)):
if alt and not _is_abs_path_id(alt):
return alt
# No clean alias: strip to a path-free public id so a host path is never advertised.
return public_model_id(raw_id) or raw_id
def _resolve_load_dir(p):
"""The concrete dir holding the GGUFs. For an HF cache repo (``models--*``
with ``snapshots/``) this is the latest snapshot dir, so /load takes the
local branch instead of the download-capable repo-id branch."""
from pathlib import Path
try:
if (p / "snapshots").is_dir():
from routes.models import _resolve_hf_cache_realpath
real = _resolve_hf_cache_realpath(p)
if real:
return Path(real)
except Exception:
pass
return p
def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]:
"""Build an entry only when GGUF quants are on disk (not Transformers/
safetensors), listing only on-disk quants. ``load_path`` is a concrete local
path so /load resolves the variant locally and never fetches a remote one."""
from pathlib import Path
from utils.models.model_config import _is_mmproj, list_local_gguf_variants
path = getattr(info, "path", None)
if not isinstance(path, str):
return None
p = Path(path)
try:
if p.is_file():
# A standalone .gguf loads by its own path; no quant sub-selection. An
# mmproj companion (vision/audio projector) is not a servable model on
# its own: _scan_models_dir's standalone-file pass does not filter it
# the way the directory scan does, so reject it here or /v1/models would
# advertise a projector and a switch could load it instead of the weights,
# evicting the loaded model. The directory branch below is already mmproj
# free (list_local_gguf_variants drops mmproj quants).
if p.suffix.lower() != ".gguf" or _is_mmproj(p.name):
return None
return _LocalGgufEntry(loader_id, str(p), ())
load_dir = _resolve_load_dir(p)
variants, _ = list_local_gguf_variants(str(load_dir))
quants = tuple(v.quant for v in variants if getattr(v, "quant", None))
return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None
except Exception:
return None
def info_has_local_gguf(info) -> bool:
"""True when *info* (a LocalModelInfo) points to on-disk GGUF weights the
auto-switch path can load. Read from the files, not ``info.model_format``: the
HF-cache scanner leaves model_format unset for GGUF snapshots, so a
model_format filter would drop every cached GGUF. Lets /v1/models advertise
exactly what /v1 can serve."""
from pathlib import Path
path = getattr(info, "path", None)
# Ollama-link entries come from a scanner _build_index intentionally skips (it
# creates symlinks on the request path), so their advertised ids never resolve.
# Don't report them as servable, or /v1/models would list unswitchable models.
if isinstance(path, str) and any(
seg in (".studio_links", "ollama_links") for seg in Path(path).parts
):
return False
return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None
def _build_index() -> dict[str, _LocalGgufEntry]:
"""Map normalized id/model_id/display_name -> local GGUF entry.
Scans the same roots Studio's model picker lists (./models, the active plus
legacy/default HF caches, LM Studio dirs, and user scan folders) so a named
local model is never missed and silently served as the loaded one. Ollama's
scanner is skipped: it creates symlinks as a side effect and this runs on the
request path.
"""
# Lazy import: routes.models imports core.inference, so import at call time.
from pathlib import Path
from routes.models import (
_scan_models_dir,
_scan_hf_cache,
_scan_lmstudio_dir,
_resolve_hf_cache_dir,
_is_hidden_model,
)
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs
index: dict[str, _LocalGgufEntry] = {}
seen_hf: set[str] = set()
def _scan_hf_once(directory) -> list:
if directory is None:
return []
try:
d = Path(directory)
if not d.is_dir():
return []
rp = str(d.resolve())
if rp in seen_hf:
return []
seen_hf.add(rp)
return _scan_hf_cache(directory)
except Exception as exc: # a missing/malformed root must skip, never crash the index
logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc)
return []
# Each source is guarded on its own so one bad root (a permission error, a
# malformed cache) drops only that source, not the whole index.
found: list = []
try:
found += _scan_models_dir(Path("./models").resolve())
except Exception as exc:
logger.debug("auto-switch: ./models scan failed: %s", exc)
try:
for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()):
found += _scan_hf_once(hf_dir)
except Exception as exc:
logger.debug("auto-switch: HF cache scan failed: %s", exc)
try:
for lm_dir in lmstudio_model_dirs():
found += _scan_lmstudio_dir(lm_dir)
except Exception as exc:
logger.debug("auto-switch: LM Studio scan failed: %s", exc)
try:
from storage.studio_db import list_scan_folders
for folder in list_scan_folders():
try:
fp = Path(folder["path"])
found += (
_scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp)
)
except Exception as exc:
logger.debug("auto-switch: scan folder %r failed: %s", folder, exc)
except Exception as exc:
logger.debug("auto-switch: scan folders enumerate failed: %s", exc)
for info in found:
raw_id = getattr(info, "id", None)
if not raw_id:
continue
# Skip what Studio hides from its pickers (validation probe, RAG embed
# weights): not chat models, so never an auto-switch target.
if _is_hidden_model(raw_id, getattr(info, "path", None)):
continue
# Advertise a client-facing alias, not an absolute filesystem path.
loader_id = _advertised_loader_id(info)
entry = _local_gguf_entry(loader_id, info)
if entry is None:
continue
# Index every alias (including the path) so a client can resolve by any of
# them, even though only the non-path loader_id is advertised.
for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)):
if key:
index.setdefault(key.strip().lower(), entry)
return index
def _index() -> dict[str, _LocalGgufEntry]:
global _scan
# Build under the lock so concurrent callers with an expired cache don't all
# run the (multi-dir) scan at once; the rest wait and reuse the fresh result.
with _lock:
now = time.monotonic()
ts, cached = _scan
if now - ts < _CACHE_TTL_S:
return cached
fresh = _build_index()
# Stamp AFTER the scan, not with the pre-scan ``now``: a multi-root scan on
# an install with many local models can itself exceed the TTL, which would
# store the cache already expired and make every request rebuild the index.
_scan = (time.monotonic(), fresh)
return fresh
def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]:
"""Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None.
``load_path`` is the concrete on-disk path to hand /load (so it never fetches
a remote), ``loader_id`` is the advertised id used as the launch-override key.
``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first
(so ids containing a colon still resolve); else the last ``:VARIANT`` is split
off and resolves only when that quant is on disk.
"""
if not isinstance(requested, str) or not requested.strip():
return None
requested = requested.strip()
try:
index = _index()
entry = index.get(requested.lower())
if entry is not None:
variant = entry.variants[0] if entry.variants else None
return entry.load_path, variant, entry.loader_id
base, sep, variant = requested.rpartition(":")
if not sep:
return None
entry = index.get(base.strip().lower())
if entry is None:
return None
wanted = variant.strip().lower()
for v in entry.variants:
if v.lower() == wanted:
return entry.load_path, v, entry.loader_id
return None
except Exception:
# Best-effort: any resolver failure falls through to the loaded model,
# so a malformed name can never turn a servable request into a 500.
return None

18
studio/backend/main.py
View file

@ -532,6 +532,11 @@ async def lifespan(app: FastAPI):
_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
@ -561,6 +566,14 @@ async def lifespan(app: FastAPI):
)
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()
@ -883,6 +896,11 @@ app.add_middleware(
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

930
studio/backend/routes/inference.py
View file

File diff suppressed because it is too large Load diff

9
studio/backend/routes/preview.py
View file

@ -17,7 +17,11 @@ from loggers import get_logger
from auth.authentication import get_current_subject
from auth.storage import DEFAULT_ADMIN_USERNAME
from models.inference import ChatCompletionRequest, LoadRequest
from routes.inference import load_model, openai_chat_completions
from routes.inference import (
disable_openai_auto_switch_for_request,
load_model,
openai_chat_completions,
)
from state.tool_policy import tools_force_disabled
from utils.client_ip import client_ip
from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint
@ -155,6 +159,9 @@ async def _serve_chat(
path = _resolve_or_4xx(run, checkpoint)
is_lora = (path / "adapter_config.json").exists()
payload = _sanitize_preview_payload(payload, is_lora)
# Preview always serves the pinned checkpoint it loads below; a public caller's
# `model` field must never trigger an OpenAI auto-switch to another GGUF.
disable_openai_auto_switch_for_request(getattr(request, "scope", None))
await _preview_lock.acquire()
keep_locked = False
try:

101
studio/backend/routes/settings.py
View file

@ -32,6 +32,16 @@ from utils.helper_precache_settings import (
helper_model_disabled_by_env,
set_helper_precache_enabled,
)
from utils.openai_auto_switch_settings import (
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS,
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
get_auto_unload_idle_seconds,
get_model_overrides,
get_openai_auto_switch_enabled,
get_stored_auto_unload_idle_seconds,
set_model_override,
set_openai_auto_switch,
)
from utils.preview_sharing_settings import (
DEFAULT_PREVIEW_SHARING_ENABLED,
get_preview_sharing_enabled,
@ -66,6 +76,33 @@ class HelperPrecacheResponse(BaseModel):
disabled_by_env: bool
class OpenAIAutoSwitchPayload(BaseModel):
enabled: bool
auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0)
class OpenAIAutoSwitchResponse(BaseModel):
enabled: bool
auto_unload_idle_seconds: int
default_enabled: bool = DEFAULT_OPENAI_AUTO_SWITCH_ENABLED
# True when the idle-unload loop will actually unload (effective TTL > 0). With
# UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled
# is false, so the UI can show idle-unload as active instead of "needs enable".
idle_unload_active: bool = False
class ModelOverridePayload(BaseModel):
model_id: str = Field(..., min_length = 1)
llama_extra_args: list[str] = Field(default_factory = list)
# ge=1: 0 is not a valid sequence length, and the setter drops a falsy value,
# so reject it at the boundary instead of accepting then silently discarding it.
max_seq_length: Optional[int] = Field(default = None, ge = 1, le = 1048576)
class ModelOverridesResponse(BaseModel):
overrides: dict[str, dict]
def _upload_limit_response(limit_mb: int) -> UploadLimitResponse:
return UploadLimitResponse(
max_upload_size_mb = limit_mb,
@ -128,6 +165,70 @@ def update_helper_precache(
return _helper_precache_response(enabled)
@router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse)
def get_openai_auto_switch(
current_subject: str = Depends(get_current_subject),
) -> OpenAIAutoSwitchResponse:
return OpenAIAutoSwitchResponse(
enabled = get_openai_auto_switch_enabled(),
auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(),
idle_unload_active = get_auto_unload_idle_seconds() > 0,
)
@router.put("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse)
def update_openai_auto_switch(
payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject)
) -> OpenAIAutoSwitchResponse:
try:
enabled, idle_seconds = set_openai_auto_switch(
payload.enabled, payload.auto_unload_idle_seconds
)
except ValueError as exc:
raise log_and_http_error(
exc,
400,
safe_error_detail(exc, fallback = "Invalid OpenAI auto-switch setting."),
event = "settings.update_openai_auto_switch_failed",
log = logger,
) from exc
return OpenAIAutoSwitchResponse(
enabled = enabled,
auto_unload_idle_seconds = idle_seconds,
idle_unload_active = get_auto_unload_idle_seconds() > 0,
)
@router.get("/openai-auto-switch/overrides", response_model = ModelOverridesResponse)
def get_openai_auto_switch_overrides(
current_subject: str = Depends(get_current_subject),
) -> ModelOverridesResponse:
return ModelOverridesResponse(overrides = get_model_overrides())
@router.put("/openai-auto-switch/overrides", response_model = ModelOverridesResponse)
def update_openai_auto_switch_override(
payload: ModelOverridePayload, current_subject: str = Depends(get_current_subject)
) -> ModelOverridesResponse:
from core.inference.llama_server_args import validate_extra_args
try:
extra_args = validate_extra_args(payload.llama_extra_args)
set_model_override(
payload.model_id,
llama_extra_args = extra_args,
max_seq_length = payload.max_seq_length,
)
except ValueError as exc:
raise log_and_http_error(
exc,
400,
safe_error_detail(exc, fallback = "Invalid model launch override."),
event = "settings.update_model_override_failed",
log = logger,
) from exc
return ModelOverridesResponse(overrides = get_model_overrides())
class PreviewLinkRotateResponse(BaseModel):
rotated: bool = True

22
studio/backend/routes/training.py
View file

@ -47,7 +47,7 @@ except ImportError:
from utils.paths import resolve_dataset_path
# Auth
from auth.authentication import get_current_subject
from auth.authentication import authenticated_via_api_key, get_current_subject
from utils.utils import log_and_http_error
@ -114,7 +114,9 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu
@router.post("/start")
async def start_training(
request: TrainingStartRequest, current_subject: str = Depends(get_current_subject)
request: TrainingStartRequest,
current_subject: str = Depends(get_current_subject),
via_api_key: bool = Depends(authenticated_via_api_key),
):
"""
Start a training job.
@ -125,6 +127,22 @@ async def start_training(
try:
logger.info(f"Starting training job with model: {request.model_name}")
# When Studio is driven as an inference API (API-key auth), refuse to start
# training while a request is in flight: training frees VRAM by unloading
# the chat model, which 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.)
if via_api_key is True:
from core.inference.llama_keepwarm import other_inference_request_count
if other_inference_request_count(current_request_counted = False) > 0:
raise HTTPException(
status_code = 409,
detail = (
"Cannot start training over the API while an inference request is in "
"progress. Wait for it to finish, or start training from the Studio UI."
),
)
# No in-process ensure_transformers_version(): the subprocess
# (worker.py) activates the correct version before importing ML libs.

37
studio/backend/storage/studio_db.py
View file

@ -1689,6 +1689,43 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]:
conn.close()
def upsert_app_setting_map_entry(
key: str, entry_key: str, entry_value: dict[str, Any] | None
) -> dict[str, Any]:
"""Set (or delete, when entry_value is falsy) one sub-entry of a dict-valued
app setting, atomically under BEGIN IMMEDIATE so concurrent writers to other
sub-entries cannot drop each other's updates."""
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone()
current = _json_loads(row["value_json"], {}) if row else {}
if not isinstance(current, dict):
current = {}
if entry_value:
current[entry_key] = entry_value
else:
current.pop(entry_key, None)
now = datetime.now(timezone.utc).isoformat()
conn.execute(
"""
INSERT INTO app_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
(key, json.dumps(current), now),
)
conn.commit()
return current
except Exception:
conn.rollback()
raise
finally:
conn.close()
def list_chat_settings() -> dict[str, Any]:
conn = get_connection()
try:

3039
studio/backend/tests/test_openai_auto_switch.py Normal file
View file

File diff suppressed because it is too large Load diff

30
studio/backend/tests/test_openai_catalog.py
View file

@ -13,6 +13,7 @@ if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import routes.inference as inf # noqa: E402
from core.inference import local_model_resolver as resolver # noqa: E402
class _Info:
@ -21,10 +22,12 @@ class _Info:
id,
display_name,
model_id = None,
is_gguf = True,
):
self.id = id
self.display_name = display_name
self.model_id = model_id
self.is_gguf = is_gguf # drives the files-based GGUF check in the test
class _FakeLlama:
@ -53,10 +56,16 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
return [
_Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup
_Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded
_Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id
# HF-cache GGUF: model_format is unset for these, so a files-based check
# (not model_format) must still list it.
_Info("models--org--Foo", "Foo", model_id = "org/Foo"),
# Non-GGUF (safetensors) can't be served via /v1: must NOT be advertised.
_Info("/data/models/Mistral-7B", "Mistral-7B", is_gguf = False),
]
monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog)
# GGUF-ness is read from the on-disk files; drive it off each info's flag here.
monkeypatch.setattr(resolver, "info_has_local_gguf", lambda info: info.is_gguf)
data = asyncio.run(inf._openai_catalog_objects())
ids = {m["id"]: m for m in data}
@ -64,9 +73,12 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
# Loaded model is present, marked loaded, and keeps context fields.
assert ids["Qwen3-Q4"]["loaded"] is True
assert ids["Qwen3-Q4"]["context_length"] == 4096
# Available-but-not-loaded models are listed too.
# Available-but-not-loaded GGUF models are listed too.
assert ids["Llama-8B-Q8"]["loaded"] is False
# The HF-cache GGUF is listed despite model_format being unset.
assert ids["org/Foo"]["loaded"] is False
# The non-GGUF model is filtered out (/v1 can never serve it).
assert "Mistral-7B" not in ids
# The loaded gguf and the on-disk copy collapse to one clean id.
assert [m["id"] for m in data].count("Qwen3-Q4") == 1
# No absolute paths or .gguf suffixes leak anywhere.
@ -76,6 +88,20 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
assert "/data/" not in blob
def test_catalog_lock_is_per_loop():
# Codex P2: a module-level asyncio.Lock ties its waiters to the loop that first
# awaited it, so a second event loop awaiting it in a multi-loop process can
# hang. The catalog lock must be per-loop (distinct lock per running loop), and
# the old shared _CATALOG_LOCK must be gone so it can't be reintroduced.
async def _get():
return inf._catalog_lock()
a = asyncio.run(_get())
b = asyncio.run(_get()) # a fresh event loop
assert a is not b
assert not hasattr(inf, "_CATALOG_LOCK")
def test_empty_and_errored_scans_are_cached(monkeypatch):
# Cache validity is keyed on the timestamp, not list contents, so an empty
# (fresh install / no local models) or errored scan is still cached for the

180
studio/backend/utils/openai_auto_switch_settings.py Normal file
View file

@ -0,0 +1,180 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Persisted opt-in controls for OpenAI-compatible model auto-switching.
Two settings, both off by default so existing API behavior is unchanged:
- ``openai_api_auto_switch_model``: when on, a ``/v1`` request whose ``model``
names a downloaded local GGUF different from the loaded one transparently
loads it before serving (llama-swap-style). Unknown names pass through.
- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is
unloaded after this many idle seconds to free VRAM.
The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env
var. Unlike the stored setting (which stays 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.
Reads are cached for a short window because these are consulted on the
per-request hot path; writes invalidate the cache.
"""
from __future__ import annotations
import os
import threading
import time
from typing import Any, Optional
OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model"
AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds"
MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides"
MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
_CACHE_TTL_S = 2.0
_cache_lock = threading.Lock()
_cache: dict[str, tuple[float, Any]] = {}
def _coerce_bool(value: Any) -> bool | None:
if isinstance(value, bool):
return value
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in {"1", "true", "yes", "on"}:
return True
if normalized in {"0", "false", "no", "off", ""}:
return False
return None
def _coerce_int(value: Any) -> int | None:
try:
return max(0, int(value))
except (TypeError, ValueError):
return None
def _cached_setting(key: str, default: Any) -> Any:
"""Read an app setting, memoized for _CACHE_TTL_S to spare the hot path."""
now = time.monotonic()
with _cache_lock:
hit = _cache.get(key)
if hit is not None and now - hit[0] < _CACHE_TTL_S:
return hit[1]
try:
from storage.studio_db import get_app_setting
stored = get_app_setting(key, None)
except Exception:
stored = None
value = default if stored is None else stored
with _cache_lock:
_cache[key] = (now, value)
return value
def _invalidate(key: str) -> None:
with _cache_lock:
_cache.pop(key, None)
def get_openai_auto_switch_enabled() -> bool:
parsed = _coerce_bool(_cached_setting(OPENAI_AUTO_SWITCH_SETTING_KEY, None))
return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED
def _stored_idle_seconds() -> Optional[int]:
"""The persisted idle TTL as an int, or None when never set."""
return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None))
def _env_idle_seconds() -> Optional[int]:
"""UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid."""
raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR)
if raw is None or not raw.strip():
return None
return _coerce_int(raw)
def get_stored_auto_unload_idle_seconds() -> int:
"""The persisted idle-unload TTL, independent of whether auto-switch is on.
The settings UI reads this so it can display and round-trip the saved value;
toggling auto-switch off must not erase it. Falls back to the env override so
the UI shows the startup default. The idle loop uses the gated reader below.
"""
stored = _stored_idle_seconds()
if stored is not None:
return stored
env = _env_idle_seconds()
return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS
def get_auto_unload_idle_seconds() -> int:
"""Effective idle TTL the idle loop runs on (0 = never unload)."""
stored = _stored_idle_seconds()
if stored is not None:
# An explicit UI/API value stays gated on auto-switch: off reports 0 so the
# off state is identical to pre-feature.
return stored if get_openai_auto_switch_enabled() else 0
# No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that
# enables idle-unload even with auto-switch off (headless/container deploys).
env = _env_idle_seconds()
return env if env is not None else 0
def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]:
"""Set both auto-switch flags in one transaction so a settings PUT can't leave
one key updated and the other stale. Both values are coerced before any write,
so an invalid value raises without persisting either."""
parsed_enabled = _coerce_bool(enabled)
if parsed_enabled is None:
raise ValueError("OpenAI auto-switch must be true or false.")
parsed_idle = _coerce_int(idle_seconds)
if parsed_idle is None:
raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
from storage.studio_db import upsert_app_settings
upsert_app_settings(
{OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle}
)
_invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY)
_invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
return parsed_enabled, parsed_idle
def get_model_overrides() -> dict[str, dict]:
"""Per-model launch overrides keyed by model id ({llama_extra_args, max_seq_length})."""
raw = _cached_setting(MODEL_OVERRIDES_SETTING_KEY, None)
return raw if isinstance(raw, dict) else {}
def get_model_override(model_id: str) -> dict:
"""The launch override applied when auto-switch loads ``model_id`` (or empty)."""
override = get_model_overrides().get(model_id)
return override if isinstance(override, dict) else {}
def set_model_override(
model_id: str,
llama_extra_args: Optional[list[str]] = None,
max_seq_length: Optional[int] = None,
) -> dict:
"""Upsert one model's launch override; an override with no fields removes it."""
if not model_id or not model_id.strip():
raise ValueError("model_id is required.")
entry: dict[str, Any] = {}
if llama_extra_args:
entry["llama_extra_args"] = [str(arg) for arg in llama_extra_args]
if max_seq_length:
entry["max_seq_length"] = max(0, int(max_seq_length))
from storage.studio_db import upsert_app_setting_map_entry
# Atomic per-entry merge so two PUTs for different models can't drop each other.
upsert_app_setting_map_entry(MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None)
_invalidate(MODEL_OVERRIDES_SETTING_KEY)
return entry

89
studio/frontend/src/features/settings/api/openai-auto-switch.ts Normal file
View file

@ -0,0 +1,89 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { readFastApiError } from "@/lib/format-fastapi-error";
export type OpenAIAutoSwitchSettings = {
enabled: boolean;
autoUnloadIdleSeconds: number;
defaultEnabled: boolean;
// True when the idle-unload loop will actually unload (e.g. enabled via the
// UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off).
idleUnloadActive: boolean;
};
type ApiOpenAIAutoSwitchSettings = {
enabled: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_idle_seconds: number;
// biome-ignore lint/style/useNamingConvention: API schema
default_enabled: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
idle_unload_active?: boolean;
};
let cachedSettings: OpenAIAutoSwitchSettings | null = null;
let inFlightSettings: Promise<OpenAIAutoSwitchSettings> | null = null;
function fromApi(
settings: ApiOpenAIAutoSwitchSettings,
): OpenAIAutoSwitchSettings {
return {
enabled: settings.enabled,
autoUnloadIdleSeconds: settings.auto_unload_idle_seconds,
defaultEnabled: settings.default_enabled,
idleUnloadActive: settings.idle_unload_active ?? false,
};
}
async function fetchOpenAIAutoSwitchSettings(): Promise<OpenAIAutoSwitchSettings> {
const res = await authFetch("/api/settings/openai-auto-switch");
if (!res.ok) {
throw new Error(
await readFastApiError(res, "Failed to load model auto-switch settings"),
);
}
return fromApi(await res.json());
}
function cacheSettings(settings: OpenAIAutoSwitchSettings) {
cachedSettings = settings;
return settings;
}
export async function loadOpenAIAutoSwitchSettings() {
if (cachedSettings) {
return cachedSettings;
}
inFlightSettings ??= fetchOpenAIAutoSwitchSettings()
.then(cacheSettings)
.finally(() => {
inFlightSettings = null;
});
return inFlightSettings;
}
export async function updateOpenAIAutoSwitchSettings(
enabled: boolean,
autoUnloadIdleSeconds: number,
): Promise<OpenAIAutoSwitchSettings> {
const res = await authFetch("/api/settings/openai-auto-switch", {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
enabled,
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_idle_seconds: autoUnloadIdleSeconds,
}),
});
if (!res.ok) {
throw new Error(
await readFastApiError(
res,
"Failed to update model auto-switch settings",
),
);
}
return cacheSettings(fromApi(await res.json()));
}

165
studio/frontend/src/features/settings/components/model-auto-switch-section.tsx Normal file
View file

@ -0,0 +1,165 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Switch } from "@/components/ui/switch";
import { useT } from "@/i18n";
import { useEffect, useState } from "react";
import {
type OpenAIAutoSwitchSettings,
loadOpenAIAutoSwitchSettings,
updateOpenAIAutoSwitchSettings,
} from "../api/openai-auto-switch";
import { SettingsRow } from "./settings-row";
import { SettingsSection } from "./settings-section";
export function ModelAutoSwitchSection() {
const t = useT();
const [settings, setSettings] = useState<OpenAIAutoSwitchSettings | null>(
null,
);
const [draftIdleSeconds, setDraftIdleSeconds] = useState("0");
const [error, setError] = useState<string | null>(null);
const [isSaving, setIsSaving] = useState(false);
useEffect(() => {
let cancelled = false;
void loadOpenAIAutoSwitchSettings()
.then((loaded) => {
if (cancelled) return;
setSettings(loaded);
setDraftIdleSeconds(String(loaded.autoUnloadIdleSeconds));
setError(null);
})
.catch((loadError) => {
if (cancelled) return;
setError(
loadError instanceof Error
? loadError.message
: t("settings.general.modelAutoSwitch.loadError"),
);
});
return () => {
cancelled = true;
};
}, [t]);
// Parse the idle-seconds draft to a non-negative integer; empty/invalid -> null.
const parseIdleSeconds = (): number | null => {
if (!draftIdleSeconds.trim()) {
return null;
}
const parsed = Number(draftIdleSeconds);
return Number.isInteger(parsed) && parsed >= 0 ? parsed : null;
};
const persist = async (
enabled: boolean,
idleSeconds: number,
syncDraft = true,
) => {
setIsSaving(true);
setError(null);
try {
const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds);
setSettings(saved);
if (syncDraft) {
setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds));
}
} catch (saveError) {
setError(
saveError instanceof Error
? saveError.message
: t("settings.general.modelAutoSwitch.saveError"),
);
} finally {
setIsSaving(false);
}
};
// Idle-unload is tied to auto-switch (the freed model reloads via the swap).
// Toggling off preserves the saved seconds rather than zeroing them — the
// backend gates unloading on the enabled flag, so it never unloads while off.
// Enabling commits the drafted value, falling back to the last saved one so
// it can never get stuck.
const handleToggle = (enabled: boolean) => {
const savedIdleSeconds = settings?.autoUnloadIdleSeconds ?? 0;
if (!enabled) {
void persist(false, savedIdleSeconds, false);
return;
}
void persist(true, parseIdleSeconds() ?? savedIdleSeconds);
};
const handleSaveIdle = () => {
const idleSeconds = parseIdleSeconds();
if (idleSeconds === null) {
setError(t("settings.general.modelAutoSwitch.idleError"));
return;
}
void persist(true, idleSeconds);
};
return (
<SettingsSection title={t("settings.general.modelAutoSwitch.sectionTitle")}>
<SettingsRow
label={t("settings.general.modelAutoSwitch.enable")}
description={t("settings.general.modelAutoSwitch.enableDescription")}
>
<Switch
checked={settings?.enabled ?? false}
disabled={!settings || isSaving}
onCheckedChange={handleToggle}
/>
</SettingsRow>
<SettingsRow
label={t("settings.general.modelAutoSwitch.idleUnload")}
description={t(
"settings.general.modelAutoSwitch.idleUnloadDescription",
)}
>
<div className="flex flex-col items-end gap-1">
<div className="flex items-center gap-2">
<div className="relative w-28">
<Input
type="number"
min={0}
step={1}
value={draftIdleSeconds}
aria-label="Idle auto-unload seconds"
disabled={!settings?.enabled || isSaving}
onChange={(event) => setDraftIdleSeconds(event.target.value)}
className="h-8 w-full pr-8"
/>
<span className="pointer-events-none absolute inset-y-0 right-3 flex items-center text-xs font-medium text-muted-foreground">
s
</span>
</div>
<Button
variant="outline"
size="sm"
disabled={!settings?.enabled || isSaving}
onClick={handleSaveIdle}
>
{isSaving ? t("common.saving") : t("common.save")}
</Button>
</div>
{error ? (
<span className="max-w-[260px] text-right text-xs text-destructive">
{error}
</span>
) : settings && !settings.enabled && settings.idleUnloadActive ? (
<span className="max-w-[260px] text-right text-xs text-muted-foreground">
{t("settings.general.modelAutoSwitch.idleActiveViaEnv")}
</span>
) : settings && !settings.enabled ? (
<span className="max-w-[260px] text-right text-xs text-muted-foreground">
{t("settings.general.modelAutoSwitch.idleNeedsEnable")}
</span>
) : null}
</div>
</SettingsRow>
</SettingsSection>
);
}

135
studio/frontend/src/features/settings/components/usage-examples.tsx
View file

@ -27,6 +27,11 @@ import {
import { HugeiconsIcon } from "@hugeicons/react";
import { useEffect, useMemo, useState } from "react";
import { Streamdown } from "streamdown";
import {
type OpenAIAutoSwitchSettings,
loadOpenAIAutoSwitchSettings,
updateOpenAIAutoSwitchSettings,
} from "../api/openai-auto-switch";
// API call type; OS axis applies to curl only (Python is OS-identical).
type ExampleType =
@ -68,6 +73,13 @@ const OS_AWARE: Record<ExampleType, boolean> = {
const CURL_TYPES = new Set<ExampleType>(["curl", "curlTools", "curlAdvanced"]);
const PROMPT = "Can Unsloth Studio do API calling?";
// Auto-switch demo: a second call naming a different downloaded GGUF so the
// example shows that the model field selects which model serves.
// A placeholder the user replaces with one of their downloaded GGUFs. A fixed
// repo is usually not one they have, so the resolver would fall through and the
// demo would keep serving the current model instead of switching.
const SWITCH_MODEL = "your-other-downloaded-GGUF";
const SWITCH_PROMPT = "Now answer as a different model.";
// web_search + python + terminal are the reliable built-in tools.
const TOOLS = ["web_search", "python", "terminal"];
// Sampling/thinking knobs for the "+ advanced" examples.
@ -163,13 +175,19 @@ function winBody(model: string, variant: Variant): string {
return JSON.stringify(body);
}
// A leading comment (valid in both bash and PowerShell) noting the model field
// selects the served model when auto-switch is on.
const SWITCH_NOTE =
'# "Switch model by request" is on: set "model" to any downloaded GGUF to switch.\n';
function curlUnix(
base: string,
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
return `curl ${base}/v1/chat/completions \\
return `${autoSwitch ? SWITCH_NOTE : ""}curl ${base}/v1/chat/completions \\
-H "Authorization: Bearer ${key}" \\
-H "Content-Type: application/json" \\
-d '${shSingle(curlBodyPretty(model, variant))}'`;
@ -181,8 +199,9 @@ function curlWindows(
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
return `$body = '${psSingle(winBody(model, variant))}'
return `${autoSwitch ? SWITCH_NOTE : ""}$body = '${psSingle(winBody(model, variant))}'
Set-Content -Path body.json -Value $body -Encoding ascii
curl.exe ${base}/v1/chat/completions \`
-H "Authorization: Bearer ${key}" \`
@ -190,11 +209,29 @@ curl.exe ${base}/v1/chat/completions \`
-d "@body.json"`;
}
// A second OpenAI call naming a different downloaded GGUF: with auto-switch on,
// Studio loads it before serving, so the model field selects the served model.
function pythonSwitchDemo(): string {
return `
# "Switch model by request" is on: replace the model below with another GGUF you
# have downloaded and Studio loads it before serving. Unknown names keep serving
# the current model.
response = client.chat.completions.create(
model=${j(SWITCH_MODEL)},
messages=[{"role": "user", "content": ${j(SWITCH_PROMPT)}}],
stream=True,
)
for chunk in response:
print(chunk.choices[0].delta.content or "", end="")`;
}
function pythonSnippet(
base: string,
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
// Standard OpenAI args are named; Unsloth extensions go through extra_body.
const named =
@ -241,7 +278,7 @@ response = client.chat.completions.create(
messages=[{"role": "user", "content": ${j(PROMPT)}}],${named}${extraBody}
stream=True,
)
${loop}`;
${loop}${autoSwitch ? pythonSwitchDemo() : ""}`;
}
function buildSnippets(
@ -249,15 +286,16 @@ function buildSnippets(
key: string,
model: string,
os: Os,
autoSwitch: boolean,
): Record<ExampleType, string> {
const curl = os === "windows" ? curlWindows : curlUnix;
return {
curl: curl(base, key, model, "plain"),
python: pythonSnippet(base, key, model, "plain"),
curlTools: curl(base, key, model, "tools"),
pythonTools: pythonSnippet(base, key, model, "tools"),
curlAdvanced: curl(base, key, model, "advanced"),
pythonAdvanced: pythonSnippet(base, key, model, "advanced"),
curl: curl(base, key, model, "plain", autoSwitch),
python: pythonSnippet(base, key, model, "plain", autoSwitch),
curlTools: curl(base, key, model, "tools", autoSwitch),
pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch),
curlAdvanced: curl(base, key, model, "advanced", autoSwitch),
pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch),
};
}
@ -346,12 +384,31 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
const [copied, setCopied] = useState(false);
const [copiedUrl, setCopiedUrl] = useState(false);
const [useTunnel, setUseTunnel] = useState<boolean>(readUseTunnelPref);
// null while loading; the same setting the General tab exposes (shared cache).
const [autoSwitch, setAutoSwitch] = useState<OpenAIAutoSwitchSettings | null>(
null,
);
const [savingAutoSwitch, setSavingAutoSwitch] = useState(false);
// Tunnel may start after the first /api/health read; refresh so it surfaces here.
useEffect(() => {
void fetchDeviceType({ force: true });
}, []);
useEffect(() => {
let cancelled = false;
void loadOpenAIAutoSwitchSettings()
.then((s) => {
if (!cancelled) setAutoSwitch(s);
})
.catch(() => {
// Best-effort: leave the toggle off if the setting can't be read.
});
return () => {
cancelled = true;
};
}, []);
const model = useLoadedModelName();
// Real key while revealed (before "Done"); otherwise a placeholder.
const key = apiKey || KEY_PLACEHOLDER;
@ -361,9 +418,10 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
const base =
useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin);
const autoSwitchOn = autoSwitch?.enabled ?? false;
const snippets = useMemo(
() => buildSnippets(base, key, model, os),
[base, key, model, os],
() => buildSnippets(base, key, model, os, autoSwitchOn),
[base, key, model, os, autoSwitchOn],
);
const osAware = OS_AWARE[lang];
@ -385,6 +443,20 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
writeUseTunnelPref(next);
};
// Same setting as the General tab; persist optimistically and revert on failure
// so the examples reflect the live model-switch behavior.
const handleToggleAutoSwitch = (next: boolean) => {
const idle = autoSwitch?.autoUnloadIdleSeconds ?? 0;
setAutoSwitch((prev) => (prev ? { ...prev, enabled: next } : prev));
setSavingAutoSwitch(true);
void updateOpenAIAutoSwitchSettings(next, idle)
.then(setAutoSwitch)
.catch(() => {
setAutoSwitch((prev) => (prev ? { ...prev, enabled: !next } : prev));
})
.finally(() => setSavingAutoSwitch(false));
};
const handleCopyUrl = async () => {
if (cloudflareUrl && (await copyToClipboard(cloudflareUrl))) {
setCopiedUrl(true);
@ -398,6 +470,41 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
{t("settings.apiKeys.usageExamples")}
</h2>
<div className="min-w-0 max-w-full overflow-hidden rounded-lg border border-border bg-muted/20">
{/* Same setting as the General tab; surfaced here so the request `model`
actually switches the served model, which the examples below show. */}
<div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
<div className="flex shrink-0 items-center gap-1.5">
<Switch
size="sm"
checked={autoSwitchOn}
disabled={autoSwitch === null || savingAutoSwitch}
onCheckedChange={handleToggleAutoSwitch}
aria-label={t("settings.general.modelAutoSwitch.enable")}
/>
<span className="text-[11px] font-medium text-foreground">
{t("settings.general.modelAutoSwitch.enable")}
</span>
<Tooltip>
<TooltipTrigger asChild={true}>
<button
type="button"
className="flex items-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label={t(
"settings.general.modelAutoSwitch.enableDescription",
)}
>
<HugeiconsIcon
icon={InformationCircleIcon}
className="size-3.5"
/>
</button>
</TooltipTrigger>
<TooltipContent className="max-w-[260px] text-[11px] leading-snug">
{t("settings.general.modelAutoSwitch.enableDescription")}
</TooltipContent>
</Tooltip>
</div>
</div>
{cloudflareUrl ? (
<div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
<div className="flex shrink-0 items-center gap-1.5">
@ -412,9 +519,9 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
</span>
{/* Only when not launched with --secure: the raw 0.0.0.0 port is
still globally reachable, so point the user at --secure. */}
{!secure ? (
{secure ? null : (
<Tooltip>
<TooltipTrigger asChild>
<TooltipTrigger asChild={true}>
<button
type="button"
className="flex items-center rounded text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
@ -430,7 +537,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
{t("settings.apiKeys.secureHttpsHint")}
</TooltipContent>
</Tooltip>
) : null}
)}
</div>
{/* Always rendered (dimmed when off) so toggling never changes the
row height and shifts the code block below. */}

3
studio/frontend/src/features/settings/tabs/general-tab.tsx
View file

@ -48,6 +48,7 @@ import {
updateUploadLimitSettings,
} from "../api/upload-limit";
import { ChangePasswordDialog } from "../components/change-password-dialog";
import { ModelAutoSwitchSection } from "../components/model-auto-switch-section";
import { SettingsRow } from "../components/settings-row";
import { SettingsSection } from "../components/settings-section";
import { StudioVersionSection } from "../components/studio-version-section";
@ -528,6 +529,8 @@ export function GeneralTab() {
</SettingsRow>
</SettingsSection>
<ModelAutoSwitchSection />
<SettingsSection
title={t("settings.general.previewSharing.sectionTitle")}
>

16
studio/frontend/src/i18n/locales/en.ts
View file

@ -142,6 +142,22 @@ export const en = {
loadError: "Failed to load Helper LLM settings.",
saveError: "Failed to save Helper LLM settings.",
},
modelAutoSwitch: {
sectionTitle: "Model auto-switch (OpenAI API)",
enable: "Switch model by request",
enableDescription:
"When an OpenAI-compatible request names a different downloaded GGUF, load it before serving. Off by default; unknown names keep serving the loaded model.",
idleUnload: "Idle auto-unload",
idleUnloadDescription:
"Unload the model after this many idle seconds to free VRAM; the next request reloads it. 0 keeps it loaded.",
idleNeedsEnable:
"Turn on Switch model by request so an unloaded model reloads on next use.",
idleActiveViaEnv:
"Idle auto-unload is active via the UNSLOTH_MODEL_IDLE_TTL environment variable.",
loadError: "Failed to load model auto-switch settings.",
saveError: "Failed to save model auto-switch settings.",
idleError: "Enter a whole number of seconds (0 or more).",
},
previewSharing: {
sectionTitle: "Preview sharing",
enableLabel: "Public preview links",