Commit graph

64 commits

Author SHA1 Message Date
Tai An
63c6750532
fix(studio/mmproj): block cross-family projectors in flat local GGUF dirs (#5347) (#5350)
* fix(studio/mmproj): block cross-family projectors in flat local GGUF dirs (#5347)

When a flat local GGUF directory holds several unrelated models with their
own mmproj siblings, detect_mmproj_file() returned the first projector it
walked into. For the layout reported in #5347 (Qwen weights + a Gemma
mmproj in the same dir) that meant llama-server was launched with
--mmproj pointing at the Gemma projector, which fails to load and surfaces
as a confusing crash.

Disambiguation rules:
- Drop candidates whose family token (qwen/gemma/llama/mistral/phi/...)
  disagrees with the model's family. Candidates with no recognised
  family token (e.g. the HF-convention 'mmproj-F16.gguf') are kept.
- Among same-family candidates, prefer the one whose stem shares the
  longest prefix with the model (Qwen3.5-9B mmproj beats Qwen3.5-35B
  mmproj for a Qwen3.5-9B model).
- If every candidate is dropped, return None — better than attaching
  a wrong projector and getting a server-launch failure.

Tests cover the cross-family block, multi-candidate prefix tie-break,
HF-convention 'mmproj-F16.gguf', unrecognised families, and the
existing search_root walk.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio/mmproj: word-bounded family match, expanded token list, launcher guard

Tighten the family-token detector to match only on word boundaries so
substring collisions stop tagging false families: phi no longer matches
sapphire, yi no longer matches yip, mimo no longer matches mimosa, and
mistral does not bleed into ministral/magistral/devstral. Pick the token
whose first occurrence is leftmost in the filename rather than the first
hit in tuple order, so merge models disambiguate predictably (llama-phi
tags llama; phi-llama tags phi).

Expand _MODEL_FAMILY_TOKENS with the families an audit of the unsloth
HF org turned up that the previous list missed: devstral, ministral,
magistral (Mistral-derivative naming), nemotron, kimi, nanonets, cosmos,
mimo, apriel, lfm. Without these, a flat local GGUF directory containing
one of these weights plus an unrelated renamed projector still hit the
original #5347 failure.

Add mmproj_matches_model_family() and call it at the llama-server launch
site in core/inference/llama_cpp.py. detect_mmproj_file already drops
cross-family candidates at discovery time, but mmproj_path can also reach
the launcher via config injection or future overrides; this guard keeps
those paths from silently loading a known-wrong projector.

Tests: 12 new cases covering substring rejection, leftmost-position
selection, new family tokens, a new flat-dir Nemotron + Gemma rejection
case, and the launcher-level guard. All 21 detect_mmproj_file tests and
the existing 106 llama_cpp tests pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio/mmproj: pair via GGUF general.* metadata, not just filenames

Real Unsloth vision GGUFs carry rich identity metadata that has been
ignored by the discovery path. Every projector under the unsloth org
has general.type='mmproj' plus general.base_model.0.repo_url pointing
at the same upstream HF repo as its weight, and the equivalent
basename, base_model.0.name, and base_model.0.organization fields. A
flat-dir mismatch is therefore decidable from the headers alone, no
matter how the user has renamed the files.

Add utils/models/gguf_metadata.py with read_gguf_general_metadata():
a fast (~30 ms) header walk that pulls only the general.* string
fields and skips everything else, cached by (resolved path, mtime_ns,
size). Mirrors the parser shape already used by
LlamaCppBackend._read_gguf_metadata so the format handling is
consistent.

is_mmproj_by_metadata() returns True/False/None from general.type,
and pairing_score() returns 100 for an exact base_model URL match,
80 for basename plus organization match, 60 for basename only, -1
for definitive metadata disagreement, and 0 when neither side has
enough metadata to decide.

Rewire detect_mmproj_file() to a two-stage selector:
  1. Detect projectors via metadata (general.type) when present, else
     fall back to the filename substring heuristic. This recovers
     headerless projectors AND projectors whose name does not contain
     'mmproj' but whose header advertises one.
  2. Score each candidate against the weight via pairing_score. Drop
     candidates with score -1 (definitive metadata disagreement). For
     candidates with score 0 (no usable metadata) fall back to the
     existing filename family-token check, dropping recognised-family
     mismatches. Pick the survivor with the highest (score,
     longest_prefix, -len(stem)) tuple, so a metadata URL match
     always wins over a filename-prefix match.

Tests: 16 new cases. tests/test_gguf_metadata.py covers the parser
(missing file, non-GGUF, string extraction, walking past arrays and
uint32s, cache invalidation by mtime/size) and the score helpers.
tests/test_detect_mmproj_file.py adds end-to-end cases that synthesise
real on-disk GGUF headers: URL match wins over a longer-prefix
sibling, URL mismatch returns None even when filenames match, a
projector named 'vision-projector.gguf' is still discovered via
general.type, and a 100-score header match outranks a near-perfect
filename prefix on a headerless candidate.

All 75 tests across detect_mmproj_file, gguf_metadata, llama_cpp
load progress, cached gguf routes, trained model scan, and vision
cache pass.

* studio/mmproj: shorten comments and docstrings across the #5347 changes

Trim verbose explanations to one-line statements of intent. The
behaviour is unchanged: 161 tests across detect_mmproj_file,
gguf_metadata, llama_cpp_load_progress (+ matrix), llama_server_args,
llama_cpp_cache_aware_disk_check, trained_model_scan, and vision_cache
all pass.

* studio/mmproj: shorten remaining detect_mmproj_file body comments

Trim the docstring and the dir-walking block comments inside
detect_mmproj_file to one-liners. Behaviour unchanged; 44 mmproj +
gguf_metadata + llama_cpp_load_progress tests pass.

* studio/mmproj: cap gguf_metadata cache below ceiling on every insert

The eviction branch popped exactly one entry when len >= max, so the
cache size could only converge to the cap when entries were added
slowly enough for natural growth. After a sandbox sim that reduced
the cap mid-run, len stayed above the cap because each insert popped
one and added one. Switch to a while loop so we evict until len is
strictly below the cap before inserting. Steady-state behaviour at
the default 4096 ceiling is unchanged.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-14 20:31:20 -07:00
Roland Tannous
79adfd9c71
studio: skip flash-attn install on Blackwell GPUs (sm_100+) (#5420)
* studio: skip flash-attn install on Blackwell GPUs (sm_100+)

Dao-AILab does not publish prebuilt flash-attn wheels for sm_100, sm_120,
or sm_121, and the older-arch wheels fail to load on Blackwell. Add a
shared has_blackwell_gpu() helper and gate both the install-time
(install_python_stack._ensure_flash_attn) and runtime
(worker._ensure_flash_attn_for_long_context) paths on it. Detection uses
nvidia-smi --query-gpu=compute_cap, which works on Linux and Windows.

* test: stub has_blackwell_gpu in pre-existing runtime flash-attn tests

prefers_prebuilt_wheel and falls_back_to_pypi exercise the install
paths that the Blackwell guard now short-circuits. Make them explicit
about non-Blackwell so they pass on real Blackwell hosts.

* studio: cache has_blackwell_gpu, skip Blackwell warning under NO_TORCH

- Wrap has_blackwell_gpu in functools.lru_cache so repeated calls in a
  single process avoid redundant nvidia-smi spawns. Tests clear the
  cache via setup_method/teardown_method.
- In _ensure_flash_attn, run the NO_TORCH short-circuit before the
  Blackwell check so GGUF-only users (who never install torch anyway)
  do not see a Blackwell warning. Blackwell check still runs above the
  IS_WINDOWS / IS_MACOS gates so Blackwell-on-Windows users still see
  the explicit reason rather than a silent OS skip.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: add has_blackwell_gpu to mlx worker test wheel_utils stub

test_mlx_training_worker_config loads worker.py against a hand-rolled
utils.wheel_utils stub. Adding has_blackwell_gpu to the stub symbol
list so worker's import line resolves.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-14 18:13:50 +04:00
Daniel Han
4192fe6ebe
studio: drop unused max_grad_value schema + route plumbing (#5424)
* studio: drop unused max_grad_value schema + route plumbing

The MLX worker hardcodes max_grad_value to 5.0 after PR #5340. The
schema field, frontend payload type, route forwarder, and start_training
kwarg threading were all left in place as a transitional buffer for old
clients. The field is now genuinely unused everywhere except inside the
MLX worker, so the schema, route forwarder, and config-build entries can
go. Pydantic still tolerates older clients that send max_grad_value
because TrainingStartRequest's model_config defaults to extra=ignore.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-14 05:43:58 -07:00
DoubleMathew
a932294627
MLX training support for Studio on Apple Silicon (#5340)
* mlx fixes

* Fix studio integration, local dataset files, chat templates without the torch gpu imports

* pass grad norm in mlx worker

* fix(studio): pass MLX grad clipping settings

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* mlx: update grad value

* fix(mlx): address ci and clipping review

* fix backward compatibility and CI tests

* unsloth local is mlx function

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* dont reference runtime

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio mlx: hardcode value clipping, drop max_grad_value from frontend

Simplifies the MLX grad-clipping plumbing now that we are standardising on
elementwise value clipping at [-5, 5] for the compiled MLX path and norm
clipping disabled. The MLX worker no longer reads max_grad_norm /
max_grad_value from the request; both are pinned in one place. Frontend
stops sending the field at all, and the TypeScript request type drops it
to match. Non-MLX (CUDA/AMD/Intel) is untouched and continues to pick up
HF TrainingArguments' default max_grad_norm = 1.0.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-14 05:24:20 -07:00
Roland Tannous
9a0d6f80cb
studio: API external provider support for chat (OpenAI, Mistral, Gemini, Cohere, Anthropic, OpenRouter, DeepSeek, custom providers) (#4706)
* studio: add external provider support for chat inference

Adds the ability to connect to OpenAI, Mistral, Google, Cohere, Together,
Fireworks, and Perplexity from the Studio chat interface.

- Provider configs stored in SQLite (no API keys persisted)
- RSA-2048 key pair generated at startup for client-side key encryption
- httpx proxy client streams SSE responses in OpenAI-compatible format
- New /api/providers routes: registry, CRUD, test, models
- /v1/chat/completions routes to external provider when provider fields present
- Integration test suite covering CRUD, connection, model listing, and inference
- Frontend spec doc with full API contract

* remove frontend spec doc from branch

* fix auth fixture: handle forced password change on fresh install

* fix tests: default port 8000, allow 400 for no-model-loaded

* fix: update Cohere models to current (command-r retired Sept 2025)

* feat: add OpenRouter as 8th provider

* feat: add native Anthropic provider with Messages API translation

* fix: correct Anthropic base URL and drop top_p (conflicts with temperature)

* feat: add DeepSeek provider (deepseek-chat, deepseek-reasoner)

* feat: rename google -> gemini, refresh model list to 2.5 series

* feat: remove together, fireworks, perplexity providers

* feat: multimodal image support for external providers

- Add _build_external_messages() that preserves image_url parts for
  vision-capable providers instead of stripping them
- Update _proxy_to_external_provider() to use new helper
- Translate image_url content parts to Anthropic native image format
  in _stream_anthropic()
- Add TestVisionInference pytest class (1x1 PNG smoke test)

* test: use sloth photo URL for vision test, add Anthropic remote URL support

* fix: update Mistral model to mistral-small-2506

* update mistral default model to mistral-large-2512

* fix gemini vision test: download image as base64 data URI instead of remote URL

* add gemini-3-flash-preview as default gemini model

* fix gemini truncated reply (max_tokens 16->64) and suppress GeneratorExit on client disconnect

* increase vision test max_tokens to 215

* fix GeneratorExit: aclose stream generator before closing httpx client

* fix httpcore GeneratorExit: explicitly aclose aiter_lines before response closes

* fix duplicate [DONE] and suppress httpcore RuntimeError on Python 3.13 asyncgen cleanup

* fix: call response.aclose() before lines_gen.aclose() to prevent httpcore RuntimeError on Python 3.13

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Potential fix for code scanning alert no. 36: Clear-text logging of sensitive information

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* review: add comments for manual iteration rationale, mask password in test print, clarify Anthropic URL/models support

* perf: use shared module-level httpx client for connection pooling across requests

* studio: add API provider UI and integrate wiring (#4737)

* feat: expose external models in selector and chat settings

* feat(chat): wire external providers to backend + RSA key flow

- Fetch registry/configs; create/update/delete saved providers
- Encrypt API keys (Web Crypto RSA-OAEP) for test/models/chat
- External model selection + chat payload (provider_id/type, external_model, encrypted key, optional base URL)
- Local storage for keys + provider list; small UX/copy and guardrails

* add missing providers-api.ts file by Imagineer99

* fix: address PR review comments — system prompt visibility, retry loop, test logging

* feat(studio): encrypt external provider API keys at rest in localStorage

API keys for external providers (OpenAI, Mistral, etc.) were stored as
plaintext in localStorage, vulnerable to browser extensions and XSS.

Add password-derived AES-256-GCM encryption: on login the user's password
is used via PBKDF2 (100k iterations, SHA-256) to derive an in-memory
encryption key. API keys are encrypted before writing to localStorage and
decrypted on read. The derived key is never persisted — cleared on logout,
re-derived on next login.

Legacy plaintext keys are transparently migrated on first access. Password
changes re-encrypt all stored keys. No backend changes required — the
existing RSA-OAEP transit encryption is unaffected.

* fix: cast PBKDF2 salt to BufferSource for strict TypeScript lib types

* fix: persist session password in sessionStorage to survive page refreshes

* feat(studio): preserve image parts in external provider chat requests

toOpenAIMessage() now returns multimodal content arrays (OpenAI vision
format) when messages contain images, instead of always flattening to
plain text. This enables vision-capable external providers (OpenAI,
Gemini, Anthropic, etc.) to receive user images. The backend already
handles image_url content parts in _build_external_messages().

* studio: fix external models selectable in chat-only mode (#4779)

* fix: external models selectable in chat-only mode

* fix: model selector tabs default to active model kind

* Studio: API external provider registry + curated catalogs (HF/OpenRouter) and chat UX (#4787)

* fix: external models selectable in chat-only mode

* fix: model selector tabs default to active model kind

* feat(studio): expand provider registry, curated catalogs, and chat UX

- Add Hugging Face, Kimi, Qwen; remove Cohere; reorder registry
- model_list_mode curated for HF/OpenRouter; lightweight /models check
- API returns default models for curated providers; expose model_list_mode
- Frontend: provider logos in model picker, providerType on external models
- Chat providers dialog: curated vs remote flows, motion polish
- Thread: LayoutGroup + composer motion alignment with app easing

* fix(studio): disable Anthropic tool-calling flag and preselect curated defaults

* feat(studio): add external provider logos and ApiProviderLogo helper

* Studio: Polish API Providers dialog  (#4899)

* fix: lower verbage in API providers page

* fix: fix(studio): tune API Providers dialog width with rem-based responsive caps

* feat: add custom provider support (#4902)

* fix: replace crypto.subtle with node-forge for HTTP compatibility

crypto.subtle is only available in secure contexts (HTTPS/localhost),
which breaks provider API key encryption when Studio is accessed over
plain HTTP on remote GPU VMs. Switch to node-forge for RSA-OAEP and
AES-256-GCM operations — same algorithms, works on any origin.

* fix: store provider API keys as plaintext in localStorage

Drop AES-256-GCM at-rest encryption for provider API keys. The
session-password-derived encryption broke on auto-login via refresh
token (password never captured), causing keys to silently vanish.
API keys are still RSA-encrypted in transit via node-forge. At-rest
encryption in localStorage added no real security since the
decryption key also had to live client-side.

Removes crypto-storage.ts, session password plumbing, and
reEncryptAllKeys.

* fix: use max_completion_tokens for OpenAI provider

Newer OpenAI models (gpt-4o, gpt-5.x) reject the max_tokens param
and require max_completion_tokens instead. Other providers still use
max_tokens.

* fix: skip empty assistant messages in external provider requests

Some providers (Mistral) reject assistant messages with empty content.
Filter them out when building the message list for external providers.

* Update model-selector.tsx

* Update model-selector.tsx

* Update model-selector.tsx

* Update chat-adapter.ts

* Update chat-adapter.ts

* Update chat-page.tsx

* Update chat-settings-sheet.tsx

* Update chat-settings-sheet.tsx

* Update chat-settings-sheet.tsx

* Update chat-providers-dialog.tsx

* feat: polish providers settings form UI

* style: polish provider row icon sizing and alignment

* style: stabilize provider layout

* style: add provider API key visibility toggle

* fix: add provider render on empty list

* studio/frontend: sync package-lock.json with package.json

npm ci was failing because node-forge and @types/node-forge were
declared in package.json but missing from the lockfile. Ran
npm install to regenerate.

* studio/backend: fix backend CI failures for providers router

- test_desktop_auth: include providers_router in the routes stub so
  studio.backend.main imports cleanly under the monkeypatched module
- test_providers_api: skip the whole module when STUDIO_TEST_PASSWORD
  is unset (it is an integration test against a live Studio server,
  same shape as the already-ignored test_studio_api.py)

* studio/chat: drive ChatSettingsPanel from a per-provider capability map

Replace the binary isExternalModel toggle in the sampling section with a
provider-aware capability map. Each external provider type advertises
which of top_k / min_p / repetition_penalty / presence_penalty its
chat-completions API actually accepts, so the panel only renders the
knobs that map onto the active provider's request body.

Anthropic now exposes top_k; DeepSeek hides presence_penalty (deprecated
in their docs); OpenRouter and custom providers continue to show every
knob (OpenRouter drops unsupported server-side, custom assumes
OpenAI-compat or a permissive vLLM/Ollama backend). Local models are
unaffected — null capabilities means 'show everything'.

chat-adapter.ts now forwards top_k / presence_penalty to the external
proxy only when the active provider's capabilities permit it, so the
request body matches what the UI shows.

* studio/backend: forward top_k to Anthropic; filter OpenAI model list

Two paired changes so the frontend capability map has matching backend
behaviour:

1. ExternalProviderClient.stream_chat_completion now accepts top_k and
   forwards it to the Anthropic Messages body. OpenAI-compat providers
   (which all reject unknown sampling params) still receive only the
   fields they document. The proxy route in routes/inference.py passes
   payload.top_k through, so a UI request with top_k actually reaches
   Anthropic instead of being silently dropped at the boundary.

2. PROVIDER_REGISTRY['openai'] gains a model_id_allowlist regex that
   scopes the /models picker to current-gen ids (gpt-5.5 / gpt-5.4 /
   gpt-5.3 / gpt-4.5 / o3 families). The remote /v1/models listing
   otherwise returns dozens of historical snapshots, fine-tunes and
   non-chat models (embeddings, TTS, image, moderation) that we never
   want in the chat UI. default_models is refreshed to match.

* studio/chat: relax presence_penalty to optional on OpenAIChatCompletionsRequest

Followup to 1fbf445a — chat-adapter now omits presence_penalty for
providers that do not accept it (Anthropic / DeepSeek), but the
request type still required it as a non-optional number, breaking
tsc. The backend pydantic model already defaults presence_penalty
to 0, so making it optional client-side matches reality.

* studio/backend: route OpenAI traffic through /v1/responses

OpenAI's new flagship models (gpt-5.x) return 404 'This is not a chat
model' on /v1/chat/completions and are only reachable via /v1/responses.
Add a dedicated _stream_openai_responses path in ExternalProviderClient
that:

- Translates outbound messages into the Responses shape: system messages
  are folded into the top-level 'instructions' field, user/assistant
  messages become {role, content} items with input_text / input_image
  content parts (data URLs and https URLs both pass through).
- Drops presence_penalty / top_k / frequency_penalty, none of which the
  Responses contract accepts.
- Translates inbound SSE events back into OpenAI Chat Completions
  chunks so the frontend keeps a single SSE shape:
    response.output_text.delta  -> delta chunk with content
    response.completed          -> chunk with finish_reason='stop'
    response.incomplete         -> chunk with finish_reason='length'
    response.failed / error     -> propagated error SSE line
  Stream terminates with data: [DONE] (Responses emits this verbatim).

stream_chat_completion dispatches all provider_type='openai' calls to
this path; other OpenAI-compatible providers (mistral, gemini, etc.)
continue to use /v1/chat/completions.

Frontend provider-capabilities map updated to hide presence_penalty for
OpenAI in the chat settings panel, matching the new request contract.

Includes unit coverage in tests/test_openai_responses_translation.py
exercising the request body translation, image-part rewriting, and
SSE-to-chat-completions translation via httpx.MockTransport.

* studio/chat: clamp external max_tokens to 32k to stay within provider caps

The chat settings slider already capped maxTokens at 32768 for external
models, but a value persisted from a prior local-model session (where
the cap can be 128k+) was sent verbatim to the provider — Claude Opus
returns 'max_tokens: 131072 > 128000' on requests like that, and other
providers have stricter limits still.

Expose EXTERNAL_MAX_OUTPUT_TOKENS from provider-capabilities (32k) and
use it both for the slider max and as the clamp inside chat-adapter's
external-request body. 32k sits below the tightest declared output
limit across the providers we ship and well above what a typical chat
reply needs; the local-model path is unaffected.

* studio: drop temperature/top_p for OpenAI reasoning models

gpt-5.x / o3 / gpt-4.5 are reasoning-class models served via
/v1/responses, and reject temperature and top_p with
'Unsupported parameter' 400s. The OpenAI registry allowlist already
scopes the picker to those families, so neither knob ever applies on
this branch.

- external_provider._stream_openai_responses no longer puts
  temperature or top_p in the request body (kept on the method
  signature for API symmetry with the other stream methods).
- ProviderCapabilities gains temperature/topP flags; OpenAI sets both
  to false. ChatSettingsPanel hides the sliders for OpenAI so the user
  does not see inert controls.
- chat-adapter omits temperature/top_p from the external request body
  when the active provider does not advertise them.
- OpenAIChatCompletionsRequest type marks both as optional, matching
  the new chat-adapter shape.
- test_responses_request_body_uses_input_and_instructions: assertions
  flipped to confirm temperature / top_p are absent from the body.

* studio: stop forwarding top_k to Anthropic

Claude 4.x (Opus / Sonnet / Haiku 4.x) returns 400 'top_k is
deprecated for this model' on any request that includes top_k. It
was always optional on the older 3.x line, so dropping it
unconditionally for every Anthropic call is the simplest path —
no per-model gate to maintain.

- external_provider._stream_anthropic no longer adds top_k to the
  Messages body (kept on the method signature for API symmetry).
- provider-capabilities sets anthropic.topK = false so the chat
  settings panel hides the Top K slider for Anthropic providers
  and chat-adapter does not send top_k in the external request.

* studio: gate Anthropic top_k drop to Claude 4.7 only

Previous commit (b5aa6ffd) dropped top_k for every Anthropic call,
but only Claude 4.7 (Opus/Sonnet/Haiku) actually rejects it. 4.6, 4.5,
and the 3.x line still accept top_k and use it as documented.

Backend: _stream_anthropic matches the model id against
^claude-(opus|sonnet|haiku)-4-7(-|.|$) and only strips top_k when it
hits. Every other Claude generation continues to receive the value
from the chat settings panel.

Frontend: anthropic.topK is restored to true so the Top K slider is
visible again — the backend handles the per-model drop, and the
4.7 case is silent (request still succeeds without top_k).

* chore: hide dated openai models in provider select

* studio/providers: apply model_id_denylist when listing remote models

The OpenAI registry entry gained a model_id_denylist regex matching
dated snapshot ids (-YYYY-MM-DD) in 048d73bf, but the list-models
route was never consulting it, so the snapshots still showed up
alongside their canonical ids (gpt-5.5 and gpt-5.5-2026-04-23 both
listed). Apply the denylist with .search() right after the allowlist
filter so dated entries are dropped before the response is built.

* studio/chat: seed registry default_models for remote providers in picker

The Anthropic provider runs in remote model-list mode, so the picker
started with an empty availableModels until the user clicked
'Load Models'. If that /api/providers/models call fails (e.g. the
known transient decryption error during key rotation), the user sees
no models at all — claude-haiku-4-5 in particular was missing from
the dialog even though it is seeded in the registry.

Always pre-populate availableModels with the registry's default_models
when a provider type is selected (curated and remote alike), and have
loadModels() return the union of defaults + the live /models response
so registry-seeded ids are reachable regardless of what the provider's
endpoint returns or whether the call succeeds at all.

* studio/backend: diagnostic logging on provider key decryption

Decryption failures currently log just 'Failed to decrypt API key:
Decryption failed', which leaves no way to tell whether the cause is
a stale public key in the browser, a corrupted ciphertext, an
unexpected exception class, or a server-side keypair rotation. That's
the gap the next reproduction needs to close.

- key_exchange now publishes a short SHA256 fingerprint of the public
  key PEM. init_key_pair logs the fingerprint on generation and warns
  if it is ever called a second time (re-init silently invalidates
  every browser that cached the previous public key).
- decrypt_api_key wraps both the base64 decode and the RSA decrypt
  in dedicated try/excepts that log exception type, ciphertext byte
  length (RSA-2048 should be exactly 256), input string length, and
  the current public-key fingerprint.
- GET /api/providers/public-key returns the fingerprint alongside the
  PEM so the frontend can correlate a future encrypt-time fingerprint
  against the decrypt-time fingerprint and prove or rule out a
  keypair rotation as the cause.
- The /test and /models route-level decrypt warnings now include the
  exception class name (alongside the existing message).

* studio/providers: hide dated Anthropic snapshots from the model picker

Anthropic's /v1/models returns dated snapshot ids (e.g.
claude-3-5-sonnet-20241022, claude-3-5-haiku-20241022) alongside
the canonical names users actually want to pick. Same intent as
the OpenAI denylist added in 048d73bf, just a different date
format — Anthropic uses -YYYYMMDD (no dashes) while OpenAI uses
-YYYY-MM-DD.

- Add model_id_denylist = re.compile(r'-\d{8}$') to the anthropic
  registry entry. The /api/providers/models route already applies
  any denylist after fetching, so dated ids drop out automatically.
- Strip the dated 3.5 ids from default_models so the seeded picker
  no longer surfaces them; keep claude-opus-4-7 and the 4.5 family
  as the curated set.

Net effect: the picker shows opus-4-7 / opus-4-5 / sonnet-4-5 /
haiku-4-5 only, regardless of whether the remote /models call
succeeds or fails.

* fix: provider dialog and mistral short list

* style: fix provider dialog curated list styling

* fix: provider dialog curated model ids placeholder reference

* style: rename Providers to Cloud and tighten dialog header spacing

* UX: rename Providers to Cloud, remove header shortcut

* studio/chat: normalize structured delta.content from reasoning providers

Mistral's magistral (and similarly-shaped reasoning models) stream
chat-completion deltas where choices[0].delta.content is an array of
structured parts rather than a plain string, e.g.
  [{ type: 'text', text: '...' }, { type: 'thinking', thinking: '...' }]
The accumulator did 'cumulativeText += delta', which coerced each
part to '[object Object]' and produced output like
  '[object Object][object Object]...Hey there!'.

Add extractDeltaText() to normalize delta.content before append:
- string → returned as-is
- array of parts → text/output_text parts contribute their .text or
  .content; thinking/reasoning parts are re-wrapped inline as
  <think>...</think> so the downstream parseAssistantContent lifts
  them into a reasoning part the same way it does for providers that
  emit thinking inline. magistral keeps its thinking panel; no other
  provider's output shape changes.
- unknown shapes → dropped rather than stringified, so a stray field
  cannot pollute the rendered chat with '[object Object]'.

* Studio: restore Cloud icon shortcut in chat header

Brings back the header chip that opens Settings -> Cloud (external
providers) directly from the chat view. Same button as before the
bf24e604 removal: single-mode only, opens useSettingsDialogStore on
the 'connections' tab, tooltip 'API providers'.

* studio/chat: strip trailing template literal from external provider streams

Mistral's magistral occasionally appends a literal '${response}' token
after its actual answer — likely a training-format artifact, since it
keeps happening with an empty system prompt and only on that model.

Apply a tight strip in the chat-adapter SSE accumulator: when the
active provider is external, drop a trailing '${...}' template literal
(with optional whitespace) from cumulativeText after each chunk. The
regex anchors to end-of-string, so mid-stream fragments ('${re')
remain untouched and only collapse once the closing brace arrives.
Local-model output is unaffected.

* studio/providers: scope Kimi picker to kimi-k2.6 / kimi-k2.5

Mirror what the live Kimi docs surface as the current models
(https://platform.kimi.ai/docs/models). Everything else the
remote /v1/models call returns — moonshot-v1-* legacy ids and
dated k2 previews like kimi-k2-0711-preview — is filtered out.

- default_models: ['kimi-k2.6', 'kimi-k2.5'] (was four
  legacy moonshot-v1 ids plus the dated k2 preview)
- model_id_allowlist: ^kimi-k2\.[56]$ applied in the
  /api/providers/models route after the live fetch
- doc-link comments point at platform.kimi.ai overview /
  models / list-models for the next refresh

* studio: drop temperature/top_p for Kimi reasoning models

Kimi k2.5/k2.6 are reasoning-class. The API locks temperature and
top_p to fixed defaults and 400s on any other value with
'invalid temperature: only 1 is allowed for this model'.

The frontend capability map already gated these knobs out of the
external request body, but the OpenAI-compat path on the backend
unconditionally re-adds them from the pydantic ChatCompletionRequest
defaults (temperature=0.7 etc), so the gate was bypassed end-to-end.

Add a generic body_omit hook on the provider registry that
stream_chat_completion consults after building the body, and use it
to strip temperature/top_p for Kimi. Frontend provider-capabilities
flips kimi.temperature and kimi.topP to false so the sliders are
hidden in the chat settings panel as well.

* studio/providers: scope Gemini picker to current 3.x + *-latest aliases

Google's /v1beta/openai/models returns dozens of historical,
experimental, and non-chat ids that we never want in the chat UI.
Cap the picker to the current curated set:

- gemini-3.1-pro-preview
- gemini-3.1-flash-lite
- gemini-3-flash-preview
- gemini-pro-latest
- gemini-flash-latest
- gemini-flash-lite-latest

Default_models seeded with these, model_id_allowlist applied in
the /api/providers/models route to drop anything else the live
fetch returns.

* studio/providers: switch Hugging Face to remote model listing

Per the Inference Providers docs
(https://huggingface.co/docs/inference-providers/index),
GET https://router.huggingface.co/v1/models returns the full
chat-model catalog across all providers, including per-provider
metadata. The OpenAI-compatible endpoint we already use for
chat completions accepts the same Bearer token, so flipping
model_list_mode from 'curated' to 'remote' lets users discover
models via the existing list_models() path without any new
wiring.

- model_list_mode: 'remote' (was 'curated')
- default_models refreshed with current popular ids
  (gpt-oss-120b, DeepSeek-V3, Llama-3.3-70B, Qwen2.5-72B) so the
  picker still has a sensible seed if /v1/models fails
- notes updated to reference the docs page and clarify the
  endpoint is chat-only

* UX: chat cloud icon changed to model select signifier

* studio/providers: org allowlist + count cap for HF Inference picker

The HF /v1/models response is the full cross-provider catalog (hundreds
of ids — community fine-tunes, mirrors, fp8 variants, dated snapshots).
Scope the picker to the first-party org repos worth surfacing and cap
the post-filter list.

- model_id_allowlist matches the org prefixes openai/, deepseek-ai/,
  google/, meta-llama/, Qwen/, moonshotai/, mistralai/, zai-org/.
  Anything outside those orgs is dropped.
- model_id_limit (new registry field) caps the post-filter list. The
  list-models route now slices [:limit] after allowlist/denylist; set
  to 15 for HF Inference. Other providers leave it unset and behave
  exactly as before.
- default_models stays as the seed so the flagship ids users care
  about (gpt-oss-120b, DeepSeek-V3, Llama-3.3-70B, Qwen2.5-72B) are
  always reachable regardless of the API's response order.

Dedup is already handled in loadModels() via Set, so no additional
work needed there.

* style: adjust cloud icon right margin with rem spacing

* Studio: cloud openai reasoning level toggle (#5402)

* feat: cloud openai reasoning level toggle

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: honor enable_thinking=false

* fix: prevent local reasoning toggle regressions and align OpenAI effort levels

* fix: isolate external OpenAI reasoning toggle state

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>

* fix: clamp reasoning effort

* fix: align OpenAI reasoning effort

* fix: clear stale GGUF badge state

* ui: new badge on cloud setting

* fix: separate selected models from cached provider model list

* Studio: anthropic effort by model family (#5412)

* feat: external thinking control and Anthropic effort mapping

* fix: anthropic thinking constraints and 4.6 max effort mapping

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: harden Anthropic thinking params and effort mapping

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* studio/backend: drop top_p from Anthropic body when thinking is enabled

PR 5412 added body['top_p'] = max(0.95, min(top_p, 1.0)) inside the
thinking branch of _stream_anthropic, but Anthropic returns 400 on
extended/adaptive thinking when both temperature and top_p are set:

  invalid_request_error: temperature and top_p cannot both be
  specified for this model. Please use only one.

(Observed on Claude Opus 4.6.) The contract for thinking-enabled
requests is temperature=1 with neither top_p nor top_k allowed.

Replace the body['top_p'] = ... line with body.pop('top_p', None).
Defensive pop rather than a bare delete: the base body construction
above does not currently set top_p, but a future edit that adds it
would silently reintroduce the regression.

* studio/chat: force reasoningEnabled=true on local reasoning-effort models

Followup to PR 5402 / 5412. The model-status refresh path in
use-chat-model-runtime carried reasoningEnabled forward verbatim for
every reasoning-capable model. That left one observable edge case:

  1. user picks an external model that supports Off (gpt-5.x, Claude
     4.x), clicks Off — store sets reasoningEnabled=false
  2. user switches back to a local reasoning-effort model
     (gpt-oss / Harmony-style) which does NOT support Off
  3. composer's effectiveReasoningEnabled override paints the UI as
     'Think: <level>' (on)
  4. chat-adapter sees reasoningEnabled=false on the local branch
     and sends '{}', so the backend's _request_reasoning_kwargs
     returns None and the Harmony template falls back to its own
     default effort instead of the displayed level

Mirror the composer's override in the store on load: for local
reasoning-effort models (where supportsReasoningOff is false), force
reasoningEnabled=true so the store and the UI agree on every send.
Other reasoning styles still inherit prior state — only the
reasoning-effort family changes.

* studio/backend: align Anthropic thinking with the extended-thinking docs

Two compliance fixes against
https://platform.claude.com/docs/en/build-with-claude/extended-thinking

1. Adaptive-mode effort field shape
   The docs spell adaptive thinking as:
     {'thinking': {'type': 'adaptive'}, 'effort': {'type': '<level>'}}
   We had been sending the legacy 'output_config: {effort: <level>}'
   shape, which Anthropic appears to silently ignore — adaptive ran
   at the server default effort regardless of the user's selection.
   Rename to 'effort: {type: <level>}'.

2. thinking_delta event translation
   The Messages-API streams reasoning content as
   content_block_delta events with delta.type == 'thinking_delta',
   which our SSE loop was dropping entirely. On Claude 4.5/4.6 with
   display=summarized (the default), the user would see the answer
   text but never the reasoning panel. Wrap thinking_delta.thinking
   as inline <think>...</think> chunks (same pattern as the OpenAI
   Responses path) so the frontend's parseAssistantContent lifts it
   into the reasoning channel. The </think> closer fires on the
   first text_delta transition, on content_block_stop for the
   thinking block, on message_delta, and on message_stop —
   whichever arrives first — so no model path can leak an
   unclosed <think> into chat output.
   signature_delta events are left as no-ops; they carry
   verification metadata, not user-visible content.

Adds test_anthropic_thinking_translation.py with httpx.MockTransport
coverage of: effort shape on adaptive (Claude 4.6), budget_tokens
shape on manual (Claude 4.5), thinking_delta wrapping with signature
suppression, and thinking-only turns (display=omitted on Opus 4.7).

* studio/backend: revert Anthropic adaptive effort to output_config nesting

The previous commit (0a664df4) moved the adaptive-thinking effort
field to a top-level 'effort: {type: <level>}' based on a misread of
the docs page. The actual Messages API schema nests it under
output_config:

  thinking:       optional ThinkingConfigParam   ({type: 'adaptive'})
  output_config:  optional OutputConfig
    effort:       optional 'low' | 'medium' | 'high' | 'xhigh' | 'max'

Sending the top-level field produced:
  400 invalid_request_error: effort: Extra inputs are not permitted

Restore the body to:
  body['thinking'] = {'type': 'adaptive'}
  body['output_config'] = {'effort': effort}

This was the shape PR 5412 originally shipped (and the author
validated against live APIs). My 'compliance fix' was a regression.

The companion thinking_delta SSE translation added in 0a664df4 stays
— that part WAS missing from the previous shape and is unchanged
by this revert. Test pinning the body shape flipped to assert
output_config.effort, top-level effort is asserted absent.

* studio/backend: opt in to summarized thinking display on adaptive

Per the adaptive-thinking docs, the 'display' field on the thinking
config defaults to 'omitted' on Claude Opus 4.7 (and Mythos Preview).
With 'omitted' the API still emits a thinking content block, but its
'thinking' field is empty — only the signature_delta arrives.

Our SSE handler would then surface a stray '<think></think>' for the
empty block and the reasoning panel would stay blank for the entire
response. Set 'display': 'summarized' explicitly on the adaptive
thinking config so Opus 4.7 emits thinking_delta events the same way
Opus 4.6 / Sonnet 4.6 do (where 'summarized' is the default, making
the explicit setting a no-op there).

The manual-thinking branch (Claude 4.5) is unaffected — its default
is also 'summarized', and we have no reason to override it.

* studio/backend: log Anthropic SSE event counts for thinking diagnostics

Reports of 'no reasoning panel content on Anthropic' have two
distinct causes that produce the same symptom:

  1. Anthropic streamed thinking_delta events but our frontend
     dropped them somewhere on the rendering side.
  2. Anthropic did not emit thinking_delta at all (adaptive mode
     can skip thinking for simple prompts even with effort=high,
     and display=summarized only re-enables the *content* — it
     does not force thinking to happen).

Tally each event type for the duration of one stream and log the
counts in the finally branch, so the next 'no reasoning content'
report shows immediately whether thinking_delta was even on the
wire. Zero counts → upstream (model/effort/prompt choice).
Non-zero counts → triage moves to chat-adapter / parse-assistant
-content / the reasoning component.

* studio/backend: route external_provider logs through structlog

The studio backend wires structlog as the active logger (via
LogConfig.setup_logging at main.py:262), but external_provider.py
was using stdlib logging.getLogger(__name__) for every diagnostic.
The stdlib root logger defaults to WARNING with no handlers
attached, so plain logger.info('...') and logger.debug('...') from
this module were being silently dropped — including the
'Proxying chat completion to <url>' and the new
'Anthropic stream event counts' lines. Only WARNING/ERROR survived
(via the implicit fallthrough that the user actually observed
when an Anthropic call 400'd).

Switch the module-level logger to structlog.get_logger(__name__),
matching the routes/providers.py and routes/inference.py pattern.
All existing call sites use printf-style positional args, which
structlog accepts unchanged — no other edits needed.

* studio/backend: disable read timeout on SSE streams to external providers

Anthropic Opus 4.7 (adaptive thinking) and OpenAI gpt-5.x (/v1/responses)
can pause for tens of seconds between bytes while the model is
internally reasoning. httpx's read timeout is the *gap* between
successive reads, not a wall clock on the whole request — so the
shared 120s default was cutting streams mid-response:

  log: Anthropic stream event counts (... text_delta: 11)
       Read timeout from anthropic

(eleven text deltas in, no content_block_stop, no message_stop)

Add a separate _stream_timeout on ExternalProviderClient with
read = None (no gap timeout) and the same 10s / 120s connect/write/
pool bounds, then use it at the three SSE streaming call sites:
default OpenAI-compat chat completions, _stream_anthropic, and
_stream_openai_responses. Non-streaming call sites (chat_completion,
list_models, verify_models_endpoint_lightweight) keep self._timeout
because a stuck non-streaming response should still fail fast.

* studio/backend: log outbound Anthropic request shape for thinking debug

After bumping to Xhigh effort the user still saw zero thinking_delta
events and only one content_block_start, meaning Anthropic Opus 4.7
opened no thinking block at all. Per the effort docs that should be
impossible — Xhigh always thinks. Two open hypotheses:

  1. Our adaptive branch is not wiring output_config.effort onto the
     outbound body for this code path (regex miss, frontend never
     propagated reasoning_effort, etc).
  2. Anthropic is silently accepting output_config as an unknown
     field and falling back to high default effort regardless.

Add a single-line structlog INFO right before the stream POST that
echoes the keys actually present on the body (thinking, output_config,
temperature, presence of top_p / top_k, max_tokens). Messages are
deliberately excluded to keep PII out of the log. With this in place
the next 'no thinking on 4.7 at Xhigh' report shows immediately
whether we sent the effort knob — separating client bug from
provider behaviour.

* studio/chat: surface delta.reasoning_content from Kimi / DeepSeek thinking

Kimi (kimi-k2.6, kimi-k2-thinking) and DeepSeek's reasoner stream
their thinking content via a separate top-level field on the
chat-completion delta — choices[0].delta.reasoning_content — rather
than as a structured part inside delta.content. Per Kimi docs:

    In streaming output (stream=True), the reasoning_content field
    will always appear before the content field.

Our chat-adapter SSE loop only read delta.content (via
extractDeltaText), so the entire reasoning channel from these
providers was being silently dropped — kimi-k2.6 thinks by default
yet the chat UI showed no reasoning panel.

In the adapter:
- Read both delta.content and delta.reasoning_content per chunk
- When reasoning_content arrives, open a <think> block in
  cumulativeText (mirrors how the backend wraps Anthropic
  thinking_delta and OpenAI Responses reasoning summaries)
- When content arrives after reasoning, close </think> first
- On stream end, force-close any still-open <think> so
  parseAssistantContent can lift it into a reasoning part cleanly

Anthropic and OpenAI Responses paths are unaffected — they already
wrap as <think> on the backend and never set reasoning_content.

* studio: Kimi thinking toggle + 16k max_tokens floor

Two coordinated changes so Kimi's thinking is user-controllable and
the response budget meets the docs' floor.

Toggle (frontend + backend):
- getExternalReasoningCapabilities now handles provider=='kimi':
  kimi-k2.6 -> reasoning_style=enable_thinking, reasoningOff allowed
  kimi-k2-thinking -> always on (reasoningAlwaysOn=true, no off)
  kimi-k2.5 (and anything else) -> no reasoning controls
- chat-adapter already forwards enable_thinking on the
  enable_thinking-style branch, so the user toggle reaches the
  backend without additional wiring there.
- external_provider stream_chat_completion now translates the
  boolean into Kimi's wire shape on the default OAI-compat path:
    enable_thinking=True  -> body['thinking'] = {type: enabled, keep: all}
    enable_thinking=False -> body['thinking'] = {type: disabled}
  kimi-k2-thinking ignores the toggle so the API never gets a
  disabled value it would reject. Other providers on the same
  path are unaffected (gated on provider_type == 'kimi').

Max tokens floor:
- New EXTERNAL_MIN_OUTPUT_TOKENS_BY_PROVIDER table and
  getExternalMinOutputTokens helper. Kimi entry = 16000 per docs:
  'Set max_tokens >= 16,000 to ensure the full reasoning_content
  and final content can be returned without truncation.'
- chat-adapter clamps the outbound max_tokens to
  min(max(stored, providerMin), EXTERNAL_MAX_OUTPUT_TOKENS),
  so a stored value of 4096 still becomes 16000 when sending to
  Kimi (other providers unaffected, min stays effectively 64).
- chat-settings-sheet's Max Tokens slider min mirrors the same
  floor when an external Kimi model is selected, so the slider
  cannot show a value lower than what we'd actually send.
- chat-page threads activeExternalProviderType down to the panel.

* fix: stabilize external reasoning controls for Anthropic 4.6 and OpenAI o3

normalize Anthropic 4.6 reasoning effort handling by accepting max as an alias and mapping it to xhigh, while keeping Sonnet/Opus 4.6 in default model suggestions.
broaden reasoning effort typing across backend/frontend and migrate persisted max selections to xhigh for compatibility.
remove reasoning.summary=\"auto\" from OpenAI /v1/responses payloads to avoid o3 eligibility/gating errors.
tighten provider model filtering to hide retired gpt-5.3 IDs and add exact/prefix filtering support in provider routes.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: add openrouter/free + full reasoning passthrough on OpenRouter

Four-layer wire-up so the OpenRouter free-router model (which picks
a free model at random per request, filtered by needed capabilities)
shows up in the picker and its reasoning channel surfaces in the
chat UI.

Registry:
- providers.py: openrouter/free seeded at the top of openrouter
  default_models. Curated list, so picker shows it immediately.

Frontend capability map:
- provider-capabilities.ts: getExternalReasoningCapabilities now
  treats openrouter as enable_thinking style with off support. The
  Think dropdown appears for every OpenRouter model; the gateway
  silently no-ops the parameter for models that do not reason, so
  surfacing one toggle on every model is safe.

Backend reasoning passthrough:
- external_provider.py stream_chat_completion (default OAI-compat
  branch): for provider_type=='openrouter', translate the request:
    reasoning_effort in {low,medium,high} -> body['reasoning'] =
        {'effort': <level>}
    enable_thinking=True  -> body['reasoning'] = {'enabled': True}
    enable_thinking=False -> body['reasoning'] = {'enabled': False}
  Matches the documented shape at
  https://openrouter.ai/docs/guides/best-practices/reasoning-tokens
  with effort and max_tokens mutually exclusive.

Frontend SSE reader:
- chat-adapter.ts: OpenRouter streams reasoning as a third shape we
  did not handle yet: delta.reasoning_details is an array of parts
  like {type: 'reasoning.text', text: '...'}. Pull text from every
  part, merge with the existing delta.reasoning_content channel
  used by Kimi/DeepSeek, and feed the combined string through the
  same <think>...</think> wrap path so parseAssistantContent lifts
  it into the reasoning panel. Anthropic/OpenAI Responses paths
  already wrap on the backend, so they never set this field — no
  cross-provider interference.

* studio/backend: surface OpenRouter SSE errors and router-chosen model in logs

The frontend showed 'Provider returned error' for some openrouter/free
requests with nothing on the backend side to triage from — the
existing 4xx error log only fires when the upstream returns a non-200
status code, but OpenRouter (and most OAI-compat providers) return
200 OK and emit the actual failure as an SSE error event mid-stream,
which our default-path stream loop forwarded verbatim without
logging.

Best-effort diagnostics on the default OpenAI-compat stream path:
- Peek at every `data:` line in the inner forward loop, parse JSON
  best-effort (silently skip on failure so nothing is dropped).
- Count event types: delta / error / done.
- On any chunk containing an `error` field, emit a structlog WARNING
  with the provider type and the error payload — same trail the
  user would otherwise have to dig out of browser devtools.
- Latch the first non-empty `chunk.model` field. OpenRouter reports
  the router-picked underlying model there per request, so the
  finally-block summary log shows which free model handled the call.

In the finally block:

    'openrouter stream complete (model=openrouter/free,
     chosen=google/gemini-2.5-flash, events={delta: 47, done: 1})'

Zero overhead for non-error streams (a json.loads per chunk +
dict-key lookups). The structlog logger is already configured at
INFO; ERROR and WARNING surface in JSON logs without further setup.

Hoists `import json as _json` to module top so the default path can
reuse it; the existing in-function imports in _stream_anthropic and
_stream_openai_responses are now redundant but harmless.

* studio/chat: show router-picked model after 'openrouter/free:' in chip

When the user picks openrouter/free, the gateway routes each request
to a different underlying free model. Until now there was no way to
tell which one actually replied without reading the backend logs.

Surface the picked model in the active-model chip:

- chat-runtime-store gains lastOpenRouterChosenModel: string|null
  plus a setter. Reset on every model switch unless the user stays
  on openrouter/free.
- chat-adapter SSE loop latches chunk.model into the store on
  every chunk whose top-level model differs from
  openrouter/free, gated on the active checkpoint being
  openrouter/free under an OpenRouter provider.
- chat-page externalModels useMemo appends :<chosen> to the display
  name for the openrouter/free option when the store has a value,
  so ModelSelector renders e.g.
    'openrouter/free:google/gemini-2.5-flash'
  in the chip. Other models unaffected.
- Model-switch callback in chat-page clears the cached value when
  the user moves to any model other than openrouter/free, so the
  chip never shows a stale suffix from a previous session.

* studio/chat: shorten openrouter/free chip to openrouter:<short-chosen>

The full display name in use was:
  openrouter/free:inclusionai/ring-2.6-1t-20260508:free

The `:free` suffix on the underlying id already conveys 'free model',
which made the leading `/free` on the router id redundant, and the
`inclusionai/` org prefix was just noise crowding the chip.

Trim both. Now the chip renders as:
  openrouter:ring-2.6-1t-20260508:free

Strictly a display change in chat-page externalModels useMemo — the
backend wire id stays `openrouter/free`, the runtime store still
caches the full `inclusionai/...:free` value, and the model-switch
clearing logic is unchanged.

* studio/providers: switch OpenRouter to remote listing with org allowlist + cap

Same shape as Hugging Face Inference. The curated list had only four
entries; remote listing fetches OpenRouter's full ~300-model
catalog via /v1/models and the new allowlist + limit scope it back
down to a usable picker.

- model_list_mode: remote (was curated)
- model_id_allowlist matches the prefixes:
    openrouter | openai | anthropic | google | meta-llama | qwen
    | mistralai | deepseek | moonshotai | inclusionai | zai-org
    | z-ai
  Anything outside drops out.
- model_id_limit: 20 — first 20 post-filter matches from the live
  fetch; default_models stays seeded so the most useful canonical
  ids are always visible regardless of API response order.
- default_models seed extended from 4 to 6 (openrouter/free,
  openai/gpt-4o, anthropic/claude-sonnet-4-5, google/gemini-2.5-flash,
  mistralai/mistral-large-2411, deepseek/deepseek-r1).
  openrouter/free remains the first entry, so the dialog's
  loadModels() union-merge (registryDefaults first, then remote,
  deduped via Set) keeps it at the top of the picker.

* feat: external mistral thinking toggle

* studio/chat: fix TS2540 by replacing readonly ContentPart instead of mutating

The ContentPart type from @assistant-ui/react marks `text` as readonly,
so the coalesce-adjacent-same-type-part optimization in
parseAssistantContent failed the tsc build with:

  parse-assistant-content.ts(15,10): error TS2540: Cannot assign to
      'text' because it is a read-only property.
  parse-assistant-content.ts(25,10): error TS2540: ...

This broke npm run build, the Studio installer's `building frontend...`
step, and every downstream CI job that runs against an installed
Studio (Mac/Windows/Linux variants of Studio API CI, GGUF CI, UI CI,
Tauri CI, Wheel CI).

Replace the last element with a fresh merged object instead of
mutating its `text` field. Same allocation profile as the previous
path (one object swap per merge), type-safe under the readonly
declaration. Behaviour unchanged.

* studio/backend: restore summary='auto' on OpenAI Responses reasoning body

A recent refactor dropped the `summary: 'auto'` field from the
reasoning config we send to /v1/responses. Without it OpenAI does
not emit reasoning summary events on most reasoning models, which
means our SSE handler has no <think>…</think> to wrap and the chat
reasoning panel stays blank for any gpt-5.x / o3 response.

The expected wire shape is:
    body['reasoning'] = {'effort': '<level>', 'summary': 'auto'}

Two backend tests pin this:
- test_responses_reasoning_effort_included_when_requested (high)
- test_responses_reasoning_effort_xhigh_passthrough (xhigh)
Both were failing with AssertionError because the produced body
omitted `summary: auto`.

Restore the field. Skip it only for the explicit "off" case
(effort: 'none'), where summaries serve no purpose. The
enable_thinking=True fallback (no explicit effort) also pairs
medium effort with summary='auto' so that branch produces
reasoning text too.

* chat: external reasoning, OpenRouter curation, Think toggle fixes

* fix: opus and sonnet 4.6 xhigh --> max

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-14 16:13:59 +04:00
Daniel Han
0881a7a5d7
studio: security and hardening pass (auth rate-limit, sandbox, path containment, schema validation, headers) (#5375)
* studio: contain export and dataset paths under their configured roots

resolve_under_root and resolve_dataset_path previously returned absolute
paths unchanged, so an authenticated client could supply
save_directory="/tmp/escape" (or any other absolute path) and have the
exporter drop adapter files anywhere the server user could write. This
turned up during a recent audit pass where an authenticated POST to
/api/export/export/lora with save_directory="/tmp/lora_escape_test"
returned 200 and wrote adapter_model.safetensors, adapter_config.json,
and tokenizer files under /tmp.

The fix is two-layered:

storage_roots.py adds an _assert_contained(resolved, root) helper that
runs after path resolution and rejects any result whose realpath does
not sit under realpath(root). resolve_under_root now rejects '..'
segments and null bytes outright, and only accepts absolute inputs when
they are already inside the configured root (internal call sites that
re-resolve a stored absolute path stay idempotent;
worker.py:resolve_output_dir(output_dir) etc. continue to work).
resolve_dataset_path picks up the same containment rule, scoped to the
three dataset roots.

models/export.py adds field_validator("save_directory", mode="before")
to ExportCommonOptions and ExportGGUFRequest so bad input fails fast at
422 with a clear message rather than a 500 deep inside the resolver.
The validator rejects empty/whitespace, null bytes, control chars,
strings longer than 255 chars, absolute paths, and '..' segments.

routes/export.py:_export_details now returns os.path.relpath(output_path,
exports_root()) so the Export Complete dialog and /api/models/loras no
longer leak the absolute install prefix to the UI; the basename is
used as a last-resort fallback.

Verified end to end:
- POST /api/export/export/lora {"save_directory":"/tmp/foo"} -> 422
  "save_directory must be a name or relative path under the export
  root; absolute paths are rejected". /tmp/foo is not created.
- "../../etc/escape" -> 422 "may not contain '..' segments".
- save_directory="my_subdir" -> still accepted (400 only because the
  test had no checkpoint loaded yet, not because of validation).
- Internal idempotent re-resolve via resolve_export_dir(absolute path
  that is already under exports_root) returns the same path unchanged.

* studio/sandbox: harden bash + python tool execution

The sandboxed Bash and Python tool channels in Chat ran with a thin
preexec hook (PR_SET_NO_NEW_PRIVS + RLIMIT_FSIZE only). Bash had a
small word blocklist; Python had an AST safety pass aimed at
signal-tampering and shell-escape primitives. An audit pass showed
several gaps that a tool-calling model could trigger inadvertently:

- bash curl/wget/nc reached AWS IMDSv2 and returned live STS
  credentials for the instance role.
- python "import socket; s.connect((169.254.169.254, 80))"
  reached the same endpoint regardless of the bash blocklist.
- "cat /etc/passwd" was blocked at the bash side (because "passwd"
  is in the blocklist), but "open('/etc/passwd').read()" in Python
  happily returned its contents.
- "chr(115)+chr(117)+chr(100)+chr(111)" style dynamic-arg
  construction slipped through the AST shell-escape check.
- The supervisor used proc.kill() on timeout, which only signals
  the immediate pid; bash-backgrounded children survived. A fork
  bomb could spawn for the full 300s timeout window.
- Session work directories under ~/studio_sandbox/<id>/ were
  created with default umask (0o755), so any other UID on the host
  could enumerate them.
- session_id sanitisation used a one-shot str.replace("..",""),
  which is non-iterative and a small footgun.

This commit takes a conservative middle path: the sandbox still
runs as the Studio UID with no namespace tricks where the kernel
disallows them, but every chokepoint is tightened.

_sandbox_preexec now:
- calls os.setsid() so children share a process group; the
  supervisor uses os.killpg(SIGKILL) on timeout/cancel so
  backgrounded children die with the parent (new _kill_process_tree
  helper, wired into _cancel_watcher and both _bash_exec /
  _python_exec timeout branches).
- calls os.umask(0o077) so files the child writes default to 0o600.
- applies PR_SET_PDEATHSIG=SIGKILL so an orphaned child dies if
  Studio exits.
- best-effort unshare(CLONE_NEWNET) for a private network namespace
  (failure is logged and swallowed; defense-in-depth is still in
  place via the bash blocklist and the AST checker below).
- sets RLIMIT_NPROC=10000 (tunable via UNSLOTH_STUDIO_SANDBOX_NPROC),
  RLIMIT_AS=8GB, RLIMIT_CPU=300, RLIMIT_NOFILE=1024. The 10k NPROC
  figure is chosen to sit well above the ~500 LWPs a healthy Studio
  + llama-server combination already uses while still capping a
  runaway fork bomb. NPROC counts LWPs per real UID, so a lower
  figure (e.g. 256) starves legitimate bash forks
  ("bash: fork: retry: Resource temporarily unavailable").

_get_workdir:
- rejects session_id that doesn't match [A-Za-z0-9_-]{1,64};
  non-matching values bucket into a shared "_invalid" dir.
- chmod 0o700 on both the workdir and on ~/studio_sandbox/ so
  other UIDs cannot read another session's contents.

_BLOCKED_COMMANDS_COMMON gains: doas, pkexec, halt, poweroff, curl,
wget, nc, ncat, netcat, socat, ssh, scp, sftp, rsync, eval, source.
The intent is to keep general bash usage working (echo, ls, pipes,
loops, for, head, etc.) while denying the obvious egress and
escalation paths.

The AST checker (_check_signal_escape_patterns) is split into the
existing shell/signal/loop checks plus a new narrow IO denylist:
- Always flag non-literal args to anything in _SHELL_EXEC_FUNCS,
  not just _STRING_SHELL_FUNCS. Closes the dynamic-arg bypass.
- Reject calls to socket.create_connection, socket.socket().connect,
  urllib.request.urlopen, http.client.HTTP*Connection, requests.*,
  httpx.* whose literal host argument is in a cloud-metadata
  denylist (169.254.169.254 + 169.254.* + 100.64.*, plus the
  GCP/Alibaba/ECS metadata hostnames and IPv6 link-local). Public
  hosts (example.com, huggingface.co, ...) still work. Dynamic
  hosts cannot be statically blocked; mitigated by the bash
  blocklist + the netns where the kernel allows it.
- Reject literal open("/etc/passwd"), /etc/shadow, /etc/sudoers,
  /etc/ssh/*, and /proc/<pid>/environ. Other files
  (/etc/os-release, /etc/hostname, /tmp/*, user dirs) still work.

The _check_code_safety summariser is updated to include the new
network_calls and sensitive_file_reads buckets in its error string.

Regression-checked: echo, sleep, ls /tmp, for loops, piped helpers
(echo a | tr a A), urllib.request.urlopen("http://example.com"),
socket.getaddrinfo("example.com",80), open("/etc/os-release"),
open("/tmp/...","w") all still succeed. curl, wget, nc, ssh, rm,
socket.create_connection(("169.254.169.254",80)),
open("/etc/passwd"), open("/proc/self/environ") all correctly
blocked.

* studio: rate-limit login, rotate refresh tokens, add logout, security headers, gate bootstrap injection

A pass over the auth surface found a cluster of related issues that this
commit closes together.

Login (routes/auth.py):
- Add an in-memory per-IP login rate limiter. Five failed POSTs to
  /api/auth/login inside a 60s window produce 429 with Retry-After.
  A successful login clears the bucket. Previously 30 wrong passwords
  in under one second was accepted as 30x 401, which combined with
  the (now fixed) admin-username leak from /api/auth/status made
  brute-force trivial against a small password.

Logout (routes/auth.py):
- New POST /api/auth/logout returns 204 and calls
  storage.revoke_user_refresh_tokens(subject) so the refresh token
  is no longer valid. Previously POST /api/auth/logout returned 405
  and there was no way to invalidate refresh tokens short of
  changing the password. Frontend session.ts already calls
  clearAuthTokens() to drop localStorage; the new endpoint lets the
  client also tell the server to revoke server-side state.

Refresh-token rotation (routes/auth.py + auth/storage.py):
- New storage.consume_refresh_token(token) atomically validates +
  deletes a refresh token, returning (username, is_desktop). The
  /api/auth/refresh handler now mints both a new access AND a new
  refresh token; the supplied token becomes invalid. Replaying a
  consumed refresh returns 401 "Invalid or expired refresh token".
  The previous refresh_access_token helper is left in place for
  callers that intentionally want the non-rotating shape; nothing
  in the route layer uses it now.

/api/auth/status no longer leaks default_username (models/auth.py +
routes/auth.py):
- AuthStatusResponse.default_username becomes Optional[str] with a
  None default; the handler always returns None. The frontend already
  hardcodes HIDDEN_LOGIN_USERNAME = "unsloth" (auth-form.tsx:82), so
  no UI change is required.

window.__UNSLOTH_BOOTSTRAP__ no longer auto-injects (main.py):
- _inject_bootstrap is now opt-in via the
  UNSLOTH_STUDIO_INJECT_BOOTSTRAP env var. The previous default
  (inject whenever requires_password_change is true) embedded the
  plaintext bootstrap password into the first-boot HTML for any
  caller that hit /, /change-password, or any unknown SPA path.
  Browser extensions and any XSS payload on the page could read it
  trivially. With the new gate the bootstrap password lives only in
  the auth/.bootstrap_password file (mode 0o600) where it has always
  been; users typing it into a current-password field is the right
  UX. routes/auth.py:change_password also clears
  app.state.bootstrap_password defensively.

Security headers + server fingerprint (main.py + run.py):
- New SecurityHeadersMiddleware adds Content-Security-Policy,
  X-Frame-Options: DENY, X-Content-Type-Options: nosniff,
  Referrer-Policy: no-referrer,
  Permissions-Policy: camera=(), microphone=(), geolocation=(),
  interest-cohort=(), and stamps server: unsloth-studio so the
  generic uvicorn banner no longer fingerprints the stack. The
  uvicorn.Config gains server_header=False so it stops emitting its
  own Server header.

/api/health minimisation (main.py):
- Unauthenticated GET /api/health returns just
  {"status":"healthy","timestamp":...} so load-balancer liveness
  probes keep working without leaking version, device_type,
  chat_only, desktop_protocol_version, or studio_root_id to
  arbitrary callers. A request that presents a valid Bearer token
  still gets the full diagnostic payload so internal launchers and
  sibling-Studio detection (which compares studio_root_id) keep
  working.

Verification:
- 30 wrong-password POSTs to /api/auth/login -> first 5 = 401, 6th
  through 30th = 429.
- POST /api/auth/logout with a fresh token -> 204. The matching
  refresh token then fails 401.
- Login -> R1; /api/auth/refresh with R1 -> new access + R2 (R2 !=
  R1); /api/auth/refresh with R1 again -> 401; /api/auth/refresh
  with R2 -> still succeeds once and rotates again.
- curl /api/auth/status -> default_username: null.
- curl http://127.0.0.1/ does not contain __UNSLOTH_BOOTSTRAP__.
- curl -I / shows CSP, X-Frame-Options: DENY,
  X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer,
  Permissions-Policy, and server: unsloth-studio.
- curl /api/health unauthenticated -> {status, timestamp} only.
  curl with Authorization: Bearer <valid> -> full payload.
- Existing /api/system, /api/models/list, /api/train/status,
  /api/inference/status, /api/auth/api-keys, login flow, SPA root
  all still return 200 after the changes (regression smoke).

* studio: add SecurityHeadersMiddleware, MaxBodyMiddleware, /recipes redirect, gate _inject_bootstrap, minimise /api/health

This commit lands the main.py-side changes that share a single
middleware-registration spot. They are kept together because every
change here is either (a) a top-level middleware definition that has
to be added next to LoggingMiddleware, or (b) a route handler at the
same file-level.

SecurityHeadersMiddleware (Content-Security-Policy, X-Frame-Options:
DENY, X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer,
Permissions-Policy, server: unsloth-studio). The previous responses
emitted no CSP, no XFO, no Referrer-Policy and were stamped
server: uvicorn.

MaxBodyMiddleware rejects POST/PUT/PATCH on the inference / dataset /
data-recipe / train / export prefixes when Content-Length exceeds
UNSLOTH_STUDIO_MAX_BODY_MB (default 100). The audit hit this by
attaching a 50 MB plain-text file to a chat message and watching
Studio base64-encode it into the JSON body; uvicorn has no enforced
cap so the only previous guard was the per-file 50 MB ceiling that
data-recipe upload routes already enforce. The new middleware extends
that ceiling to the OpenAI-compat path that the Chat attachments
flow through. Verified: a 200 MB JSON POST to /v1/chat/completions
returns HTTP 413 "Request body too large (209,715,264 bytes; max
104,857,600)". A small valid request continues to reach the handler.

_inject_bootstrap is gated behind UNSLOTH_STUDIO_INJECT_BOOTSTRAP.
The previous default was to inline window.__UNSLOTH_BOOTSTRAP__ =
{username, password} into the first-boot HTML whenever
requires_password_change was true, which exposed the plaintext
bootstrap password to any browser extension, page script, or LAN
caller on -H 0.0.0.0. The bootstrap password remains in the on-disk
.bootstrap_password file (mode 0o600) where it has always lived;
users typing it into a current-password field is the right UX.

/api/health unauthenticated returns {"status":"healthy","timestamp":
...} only; the previous payload (version, device_type, chat_only,
desktop_protocol_version, supports_desktop_auth, studio_root_id,
native_path_leases_supported) is preserved for callers that present
a valid Bearer token, so internal launchers and sibling-Studio
detection (which compares studio_root_id) keep working.

/recipes -> /data-recipes 308 redirect. The Data Recipes page lives
at /data-recipes; users typing /recipes hit the SPA catch-all and
saw "Not Found". The redirect also preserves any tail path, so
/recipes/<rest> -> /data-recipes/<rest>.

Verified end to end with curl: CSP / XFO / X-Content-Type-Options /
Referrer-Policy / Permissions-Policy all present on /, server header
is now unsloth-studio (uvicorn's own banner is suppressed via
server_header=False in run.py from the auth-batch commit). Followed
the /recipes redirect lands on the SPA HTML.

* studio: bound TrainingStartRequest hyperparameters at the schema level

POST /api/train/start accepted any value for learning_rate, batch_size,
max_steps, max_seq_length, warmup_steps, warmup_ratio, num_epochs,
save_steps, weight_decay, gradient_accumulation_steps, lora_r,
lora_alpha and lora_dropout, including -1, 0, 1e9, and non-numeric
strings like 'abc' or 'two' (which silently coerce to 0 in the
trainer). Probing showed the API returning 200 to learning_rate=-1
and batch_size=0; only max_steps had any partial clamping.

This commit adds field_validator on every numeric hyperparameter.
Bounds are chosen wide enough to span realistic single-host
configurations (B200 with 180 GB of memory comfortably fits the
upper end) while rejecting the values that always produce broken
training:

- learning_rate: parses str/float, requires 0 < lr < 1.0. Non-numeric
  input raises with "learning_rate must be parseable as float (got
  'abc')" instead of silently coercing to 0.
- batch_size: [1, 1024].
- gradient_accumulation_steps: [1, 4096].
- num_epochs: [1, 1000].
- max_steps: [1, 1_000_000].
- max_seq_length: [1, 131072].
- warmup_steps: [0, max_steps].
- warmup_ratio: [0.0, 1.0].
- save_steps: [0, 1_000_000].
- weight_decay: [0, 10] (typical 0..0.1).
- lora_r: [1, 512].
- lora_alpha: [1, 1024].
- lora_dropout: [0.0, 1.0).

Each validator names the offending field in its ValueError message
so the 422 response body identifies which input is bad. The
learning_rate validator returns its result as str (the schema field
type is str("2e-4") for backwards compatibility) so existing call
sites that float() the value continue to work.

Verified:
- learning_rate=-1 -> 422 "learning_rate must be > 0 (got -1.0);
  typical range is 1e-6 .. 1e-3".
- learning_rate='abc' -> 422 "must be parseable as float".
- batch_size=-1 / 0 / 999999 -> 422 "batch_size must be in [1, 1024]".
- batch_size='two' -> 422 (pydantic int parser).
- max_steps=0 / -5 -> 422 "must be a positive int".
- max_seq_length=200000 -> 422 "must be in [1, 131072]".
- warmup_ratio=2.5 -> 422 "must be in [0.0, 1.0]".
- lora_dropout=1.5 -> 422 "must be in [0.0, 1.0)".
- Valid request with learning_rate='2e-4', batch_size=1, max_steps=5
  passes validation and the training run starts as normal.

* studio: redact image-decode errors, clean checkpoint dirs on cancel, tolerate Stop-button + tool-result message shapes

Three small fixes that fall under "do not let the audit findings
become user-visible papercuts".

routes/inference.py - image-decode error redaction (the audit hit
this with a 0-byte / malformed / wrong-extension image upload). The
three image-normalise sites previously raised HTTPException(400,
detail=f"Failed to process image: {e}"). When PIL raised
UnidentifiedImageError(io.BytesIO(raw)) the message string included
"<_io.BytesIO object at 0x7e40a5d7bf60>", leaking both the Python
class name (confirming the PIL/io stack) and a heap address (mildly
useful for ASLR-bypass chaining if another memory-corruption bug is
ever found). Each site now catches UnidentifiedImageError and
returns the generic "Unsupported or corrupt image format"; the
fall-through generic except returns "Failed to process image". No
exception-repr is interpolated into a response body anywhere along
these paths.

core/training/training.py - checkpoint cleanup on cancel. When a
user clicks Cancel Training, the trainer flips _cancel_requested=True
and the supervisor force-terminates the subprocess. The trainer
writes checkpoint-<step> directories under output_dir every
save_steps; previously these survived the cancel and accumulated on
disk (the audit recorded ~67 MB stuck after a 200-step cancel with
save_steps=20). New helper _cleanup_cancelled_checkpoints(output_dir)
globs checkpoint-<int> entries and removes them. It is gated by a
realpath containment check against outputs_root() so it cannot
accidentally rmtree anything outside the configured outputs root.
force_terminate() invokes the helper after the subprocess join when
_cancel_requested is true. Stop-and-Save runs are unaffected because
that path keeps _cancel_requested=False.

models/inference.py - chat message shape tolerance. Two related
frontend interactions used to crash the request validator:

- After the Stop button truncates a generation, the frontend
  retained {role:"assistant", content:""} in the conversation
  history and replayed it on the next send. ChatMessage previously
  required role="assistant" to have non-empty content or tool_calls,
  so the next message returned 422 and the thread was permanently
  broken. The validator now normalises empty assistant content to
  None so the request round-trips and the trailing empty turn can
  be ignored downstream.

- The frontend's second-round tool POST drops the streamed
  tool_call_id, hitting the strict-spec check "role=tool requires
  tool_call_id". The validator now synthesises an opaque id
  (call_<8 hex>) when missing, so the request reaches the handler
  and the model's final summarising response gets generated. The
  proper fix lives in the frontend (carry the streamed id through
  the second POST) and will follow.

Verified end to end with curl: HTTP 400 (model not loaded) on both
the empty-assistant history shape and the tool-result-without-id
shape, instead of HTTP 422 from the schema validator.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: tighten code comments from security-hardening pass

Trim verbose docstrings and inline finding references added in the
previous commits in this branch. Functionality unchanged.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: await get_current_subject in /api/health and make refresh-token consumption atomic

The /api/health auth probe called get_current_subject(creds) without
awaiting it. The coroutine object is truthy, so any caller presenting a
Bearer header (valid or not) received the full diagnostic payload
including version, device_type, studio_root_id, etc. Await the coroutine
and treat HTTPException as 'fall back to the minimal liveness payload'.

consume_refresh_token did SELECT then DELETE WHERE id under default
autocommit isolation. Two concurrent POST /api/auth/refresh requests
could both win the SELECT before either DELETE ran, defeating
single-use refresh-token rotation. Replace with a single
DELETE ... WHERE token_hash = ? AND expires_at >= ? RETURNING ...
statement so the validate-and-delete lands as one atomic op under
SQLite's write lock (3.45.1 supports RETURNING; min was 3.35).

* studio: enforce body cap on chunked uploads and drop unsafe-inline from script-src

MaxBodyMiddleware previously only inspected the declared Content-Length
header; clients omitting it or sending Transfer-Encoding: chunked
bypassed the cap and could still drive an OOM via the downstream
JSON / file readers on /v1/chat/completions, /api/inference, /api/data-recipe,
/api/datasets, /api/train, /api/export. Rewrite as a raw ASGI middleware
that drains and counts http.request frames, replies 413 once the running
total exceeds UNSLOTH_STUDIO_MAX_BODY_MB before invoking the FastAPI
handler, and replays the buffered body to downstream so route code that
calls request.json() / await request.body() works unchanged.

CSP previously included 'unsafe-inline' on script-src, which defeats the
main XSS protection. The frontend bundle does not need inline scripts;
the only inline <script> the backend ever emits is _inject_bootstrap,
which is opt-in via UNSLOTH_STUDIO_INJECT_BOOTSTRAP. Drop 'unsafe-inline'
from script-src by default; when _inject_bootstrap fires, generate a
per-response nonce, embed it on the inlined <script>, and have
SecurityHeadersMiddleware splice 'nonce-XXX' into the CSP for that one
response (the internal x-internal-script-nonce header is popped before
the response leaves the server). 'unsafe-inline' stays on style-src for
Vite-injected styles.

* studio: drop empty assistant sentinel before passthrough

ChatMessage._validate_role_shape normalises role="assistant", content=""
(the post-Stop sentinel emitted by the frontend) to content=None so the
in-process path can drop it via _extract_content_parts. The passthrough
path then ran m.model_dump(exclude_none=True), which strips the now-None
content key entirely, sending {"role":"assistant"} to llama-server / the
OpenAI-compat backend. That fails upstream and leaves the user without a
recoverable Stop->resume.

Add _drop_empty_assistant_sentinels and call it at both passthrough
message origins: _openai_messages_for_passthrough (covers
/v1/chat/completions and the Responses API which routes through it) and
the anthropic_messages_to_openai output before
_anthropic_passthrough_*. Assistant messages that carry only tool_calls
(no content) are preserved.

* studio/tests: cover audit-fix surfaces and rebase pre-existing tests

Adds and updates pytest coverage for the four bot-flagged audit fixes
landed earlier in this branch and rebases two pre-existing tests that
were broken by the relaxed-validator and /api/health auth-gate changes.

studio/backend/tests/test_middleware.py (new)
  MaxBodyMiddleware: small protected, large declared, unprotected
  passthrough, chunked-upload-over-cap rejection (the regression for
  the original Content-Length-only gap), and chunked-under-cap replay.
  SecurityHeadersMiddleware: script-src no longer carries
  'unsafe-inline', style-src still does, default headers
  (XFO/XCTO/Referrer-Policy/Permissions-Policy/server), and the
  internal x-internal-script-nonce header is consumed by the
  middleware and converted to 'nonce-XXX' in the CSP.
  /api/health: no auth -> minimal, invalid Bearer -> minimal
  (the await regression), valid Bearer -> full diagnostic payload.

studio/backend/tests/test_desktop_auth.py
  consume_refresh_token: second-call returns None, expired returns
  None, and a 64-thread concurrent pile-up against the same hash
  produces exactly one successful consumer (regression for the
  SELECT-then-DELETE race).
  test_health_response_reports_desktop_capability_fields: rebase
  against the new health_check(request) signature by going through
  TestClient with a real bearer instead of asyncio.run-ing the
  handler directly.

studio/backend/tests/test_openai_tool_passthrough.py
  Pin the new ChatMessage tolerance: assistant without content or
  tool_calls is tolerated (normalises content -> None), empty-string
  and empty-list assistant content normalise to None, and a missing
  / empty tool_call_id on role='tool' is synthesised as call_<hex>
  rather than raising. Tests for _drop_empty_assistant_sentinels
  cover the three drop shapes (empty string, empty list, missing
  content key), preservation of assistant text and tool_calls-only
  messages, and end-to-end through
  _openai_messages_for_passthrough.

studio/backend/main.py
  SecurityHeadersMiddleware.dispatch used response.headers.pop(...)
  for the nonce-header handoff; Starlette's MutableHeaders has no
  pop. Read-then-del so the internal handoff header is still
  stripped before the response leaves the server.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio/tests: rebase three more pre-existing CI tests against this branch

CI on PR #5375 was red on three tests that were tuned for behaviour
predating this branch. Updates each so the assertions match what the
audit fixes intentionally changed; no production code touched.

studio/backend/tests/test_trained_model_scan.py
  test_scan_trained_models_includes_lora_and_full_finetune_outputs
  passed an absolute tmp_path through scan_trained_models, which now
  runs resolve_output_dir / _assert_contained against outputs_root().
  Repoint outputs_root() at tmp_path via monkeypatch so the fixture
  dirs land under the configured root and the realpath containment
  check passes.

tests/test_studio_install_workspace_guard.py
  test_health_endpoint_exposes_studio_root_id_not_raw_path read
  the first 1500 bytes after @app.get("/api/health") and asserted on
  the studio_root_id literal. The handler grew (unauth short-circuit
  + await dependency gate) and the literal slid past the byte window.
  Replace the fixed window with a slice up to the next top-level
  @app.* decorator so the test surveys the whole handler regardless
  of size.

tests/studio/studio_api_smoke.py
  The "login burst (5x wrong pw) -> 401 each" assertion was tagged
  "When/if we add one, this assertion updates in the same PR." We
  added the per-IP rate-limit in routes/auth.py
  (_LOGIN_MAX_FAILS=5/60s) but missed the assertion update. Rewrite
  the burst probe to observe the new invariant: at least one 401,
  eventual transition to 429, and Retry-After present on the 429.
  Adds a small _login_with_headers helper since the existing login()
  helper drops response headers.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* ci(studio-ui): set UNSLOTH_STUDIO_INJECT_BOOTSTRAP=1 for Playwright Studios

The Chat UI Playwright test drives the first-boot change-password
form, which (per playwright_chat_ui.py step "1. Change-password
through the UI") pre-seeds the hidden current_password field from
window.__UNSLOTH_BOOTSTRAP__. That global is only emitted when the
backend's _inject_bootstrap path fires, which since the security
pass on this branch is gated behind UNSLOTH_STUDIO_INJECT_BOOTSTRAP
and defaults to off. Without the global, the React form's
current_password validator never satisfies, the submit button stays
disabled, and the composer.wait_for() probe times out on
/change-password.

Re-enable injection only for the CI Studios that drive the chat UI
across linux/mac/windows. Production deployments are unaffected: the
env var has to be explicitly opted into, and the on-disk
auth/.bootstrap_password remains the source of truth for human users
typing the password in by hand.

Covers all eight Studio launch sites: the primary chat-ui boot and
the "extra UI tests" boot for each of the three OSes, plus the
pipeTransport JSON-crash retry relaunches in the macOS workflow that
re-spawn Studio mid-job.

A follow-up frontend PR will add a visible current_password input so
the form satisfies its own validator without needing the bootstrap
auto-fill at all; once that lands this CI knob can come back out.

* studio/sandbox: drop unshare(CLONE_NEWNET); add trusted-host allowlist; block sandbox file uploads; raise CPU rlimit default to 600 s

CLONE_NEWNET inside _sandbox_preexec silently killed every outbound
HTTP request from sandboxed Python whenever the kernel allowed
unprivileged user namespaces. requests.get('https://huggingface.co'),
urllib.request.urlopen('https://en.wikipedia.org/wiki/...'),
socket.connect(('arxiv.org', 443)) all failed despite the AST visitor
intending to allow them. The bash blocklist (curl / wget / nc / ssh /
scp / sftp / rsync / socat / eval / source) plus the AST-level
metadata-host denylist still carry the network policy after this
change; CLONE_NEWNET was redundant with both.

Add _TRUSTED_PUBLIC_HOST_LITERALS + _TRUSTED_PUBLIC_HOST_SUFFIXES
(~100 informational hosts: Wikipedia language subdomains, Wikimedia,
Wikidata, Google search, Bing, DuckDuckGo, HuggingFace, GitHub,
raw.githubusercontent.com, arXiv, StackOverflow / Stack Exchange,
MDN, docs.python.org, PyTorch / TensorFlow / NumPy / pandas docs,
pypi / files.pythonhosted.org / npmjs / crates.io, ReadTheDocs,
arXiv, Britannica, BBC / Reuters / Nature / Science, NASA / CDC /
NIH / WHO open data, api.weather.gov). The visitor now blocks
literal hosts that are neither metadata nor trusted with a short
LLM-readable string so the model can retry with an allowed source
instead of choking on a multi-line error.

Block upload-shape calls regardless of host: requests.post / put /
patch / delete / request with files= or data=open(...) /
data=bytes_literal; httpx equivalents; urllib.request.urlopen /
Request with data=...; HuggingFace upload_file / upload_folder /
upload_large_folder / create_commit (module-level FQ paths AND
method-name match on any receiver). Message: "Blocked: file upload
disallowed in sandbox".

Bump UNSLOTH_STUDIO_SANDBOX_CPU_S default 300 -> 600 s so long
agentic chains that span multiple tool calls don't get SIGXCPU'd
mid-stride. Env-var override path is unchanged.

Host normalisation now strips trailing dot, userinfo @, and explicit
port before allowlist / denylist comparison so trailing-DNS-dot,
userinfo-smuggling, and explicit-:443 URLs are decided correctly.

* studio: raise default request-body cap from 100 MB to 500 MB

UNSLOTH_STUDIO_MAX_BODY_MB default goes 100 -> 500 to comfortably
cover vision + audio + multi-recipe-batch JSON payloads. The
MaxBodyMiddleware stream-counting logic from this branch's earlier
06ec088 already handles chunked bodies up to the new cap; env-var
override path is unchanged for callers that want a tighter limit.

* studio/auth: restore /api/auth/status.default_username to 'unsloth'

This branch's earlier b39e9a4 changed default_username to None on the
public /api/auth/status endpoint so the username field didn't leak to
unauthenticated callers. In practice this regressed third-party
clients (and the in-tree React login form's pre-fill UX) without
adding meaningful security: the bootstrap password is the actual
secret, and the username 'unsloth' is the documented default.

Pin default_username to storage.DEFAULT_ADMIN_USERNAME ('unsloth')
and tighten the response model so the field is required rather than
Optional. Anyone who needs anonymisation can still reach for an
allow-list deployment with auth disabled.

* studio/training: raise max_seq_length / batch_size / lora_r / lora_alpha caps

This branch's 7102815 introduced field validators with conservative
caps. The follow-up loosens them so long-context experiments and
high-rank LoRA exploration aren't gated at the schema layer:

  _MAX_BATCH_SIZE   1024     -> 4096
  _MAX_SEQ_LENGTH   131_072  -> 2_000_000   (2M tokens)
  lora_r cap        512      -> 16_384      (_MAX_LORA_R)
  lora_alpha cap    1024     -> 32_768      (_MAX_LORA_ALPHA)

_MAX_GRAD_ACCUM / _MAX_STEPS / _MAX_EPOCHS / lora_dropout /
warmup_ratio / weight_decay are unchanged. Hardware (VRAM, host
RAM, kernel launch latency) is now the binding constraint at the
new caps, which is the correct ordering -- the validator stays a
sanity check on -1 / 0 / 'abc' style garbage, not a usability gate.

* studio/tests: cover sandbox allowlist + upload block + raised training caps

studio/backend/tests/test_sandbox_tools.py (new):
  TestMetadataHostDenylist     -- short "Blocked: cloud-metadata host"
                                  message on AWS IMDS, GCP metadata,
                                  Alibaba ECS, AWS IPv6 IMDS, 169.254/16.
  TestTrustedHostAllowlist     -- Wikipedia (any language subdomain),
                                  Google, DuckDuckGo, HF, raw GitHub,
                                  arXiv, StackOverflow / family,
                                  MDN, docs.python.org, pypi, BBC,
                                  api.weather.gov, NumPy / PyTorch docs.
  TestUntrustedHostBlock       -- example.com / random unlisted host
                                  rejected with the short "Blocked: host
                                  not in sandbox allowlist; use an
                                  allowed informational source" message.
                                  Dynamic URLs (computed var) still pass
                                  -- documented limit of static analysis.
  TestHostNormalization        -- trailing dot, explicit :443, uppercase,
                                  userinfo-@-smuggle all decided
                                  correctly without false-block /
                                  false-pass.
  TestUploadDenylist           -- requests / httpx / urllib.urlopen with
                                  files= / data=open / data=bytes,
                                  HfApi().upload_file / upload_folder /
                                  create_commit, module-level
                                  huggingface_hub.upload_folder. POST
                                  json= to trusted host still passes.
  TestSandboxCpuRlimitDefault  -- pin UNSLOTH_STUDIO_SANDBOX_CPU_S=600
                                  default and confirm CLONE_NEWNET
                                  source line is gone.
  TestMaxBodyDefault           -- pin UNSLOTH_STUDIO_MAX_BODY_MB=500
                                  default.

studio/backend/tests/test_studio_train_validation.py (new):
  Pin at-cap-accepts / over-cap-rejects boundaries for
  max_seq_length=2_000_000, batch_size=4_096, lora_r=16_384,
  lora_alpha=32_768 so a future regression that tightens them back
  without explicit user opt-in is caught.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: tighten code comments across the security-hardening pass

* studio: always inject bootstrap credentials on first boot

The UNSLOTH_STUDIO_INJECT_BOOTSTRAP gate added an extra
terminal-to-browser copy-paste on every fresh install. In practice
the LAN credential leak it guarded against is narrow: the password
is one-time, the user rotates it on the very next click, the
default Studio bind is 127.0.0.1, and -H 0.0.0.0 already exposes
the entire API surface. Drop the gate so the inject fires whenever
a bootstrap password is still pending. The CSP nonce wiring stays
in place; the inline script remains the only inline script the
backend ever emits.

The three Playwright UI smoke workflows lose their
UNSLOTH_STUDIO_INJECT_BOOTSTRAP=1 lines along with the explanatory
comment blocks since the inject now happens by default.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
2026-05-13 06:12:18 -07:00
Daniel Han
5205bc0ed6
Studio: pin GPU at 95% headroom and warn on silent CPU fallback (#5323)
* Studio: pin GPU at 95% headroom and warn on silent CPU fallback

Two related runtime-side fixes for unslothai/unsloth#5106 ("model
loaded fully on RAM instead of VRAM"):

1. GPU pin threshold bump 0.90 -> 0.95
-------------------------------------

``_select_gpus`` and the auto-ctx pin loop in ``start_llama_server``
used a ``pool * 0.90`` threshold to decide whether the model fits on
GPU. Models that needed 91-94% of free VRAM were classified as "does
not fit", so Studio set ``gpu_indices = None`` and shipped
``--fit on`` to llama-server without ``-ngl``. The unsloth
llama.cpp fork's ``--fit on`` then ran with its default
``--fit-target 1024`` (1 GiB margin per device, an upstream default
inherited from ggml-org#18679). On a tight fit where compute
buffers + CUDA context push the projected free below the 1 GiB
target, the fork's fit logic shaves layer weights off the GPU --
slow inference for users whose models would have loaded comfortably
with ``-ngl -1``.

The classic reproducer from #5106 (noahterbest's log):

    GGUF size: 20.8 GB, est. KV cache: 0.1 GB, context: 4096,
    GPUs free: [(0, 22805)], selected: None, fit: True

20.8 GiB on a 22.27 GiB free RTX 4090 is 94% utilization. The model
fits (1.4 GiB headroom), but the 0.90 threshold kicks it to fit
mode. Bumping to 0.95 keeps these in the fits-on-GPU branch and
emits ``-ngl -1`` directly. The fork's ``--fit on`` still serves as
the safety net for the genuinely-too-large case.

The auto-ctx fallback also re-checks fit at 4096 before handing off
to ``--fit on``: a 20.8 GiB model with a 131072 native context fails
the auto loop at native ctx, falls back to ``min(4096, ctx)``, but
its weights + 4096 KV pin to the GPU comfortably. Without the
re-check we still emitted ``--fit on``.

``_fit_context_to_vram``'s 0.90 budget for context binary search is
intentionally left tighter than the pin fraction. That routine
chooses the slider value, where over-promising would OOM at runtime.
``_select_gpus`` decides whether to pin at all, where being
conservative pushes layers to CPU.

2. Belt-and-suspenders: warn on silent CPU fallback
---------------------------------------------------

After ``_wait_for_health`` succeeds, scan llama-server's stdout for
``model buffer size`` lines. If Studio detected GPUs and intended
GPU use but only CPU buffers were allocated, log a structured
warning citing #5106. Markers cover CUDA / ROCm / Metal / Vulkan /
OpenCL / SYCL backends. New ``_gpu_offload_active: Optional[bool]``
field surfaces the result for any future API consumer.

This catches runtime-load failures the install-time fix cannot
cover (cudart bundle pairing PR #5322 is the install-side
companion): user overriding ``--fit-target``, uncommon driver +
toolkit configurations, future regressions in the install path.

Tests: 10 new cases in studio/backend/tests/test_llama_cpp_context_fit.py:
* TestTightFitPinsToGPU x3: noahterbest's exact reproducer (auto and
  explicit ctx pins to GPU at 94%); guard against threshold over-
  broadening (genuine overflow still falls back to ``--fit on``).
* TestClassifyGpuOffload x7: CUDA / ROCm / Metal buffer markers
  return True; CPU-only buffer lines return False; absent buffer
  lines or no GPUs detected return None (no warning).

25 context-fit tests pass (15 baseline + 10 new). 511 tests total
across the affected test files. No regressions.

Refs #5106

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Trim comments to be more succinct

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-13 04:48:15 -07:00
Daniel Han
379f5a5aa6
Studio: add torch's pip nvidia DLL dirs to PATH on Windows (#5324)
* Studio: add torch's pip nvidia DLL dirs to PATH on Windows

Studio's install_python_stack bundles torch with matching CUDA
wheels (nvidia-cuda-runtime-cu13, nvidia-cublas-cu13, etc.) which
ship cudart64_X.dll, cublas64_X.dll, and cublasLt64_X.dll under
the prefix's Lib/site-packages/nvidia/<pkg>/(bin|Library/bin)/
tree. The Linux runtime env block in start_llama_server already
pulls the equivalent nvidia/cu*/lib paths into LD_LIBRARY_PATH,
but the Windows block did not do this, so the prebuilt
llama-server.exe could not resolve cudart64_X.dll at runtime
unless the user had a matching system CUDA toolkit on PATH. That
is the root cause of the Windows reports in
unslothai/unsloth#5106 ("GPU detected but model loaded entirely
on RAM/CPU"), and matches Roland's repeated workaround in that
issue: install matching CUDA toolkit version.

Brings the Windows env block in line with the Linux pattern:

* New LlamaCppBackend._windows_pip_nvidia_dll_dirs resolver
  globs <prefix>/Lib/site-packages/nvidia/<pkg>/bin and
  <prefix>/Lib/site-packages/nvidia/<pkg>/Library/bin. Both
  layouts are seen in the wild across cuda_runtime / cublas /
  cudnn / nvjitlink wheels.

* The Windows env block now extends path_dirs with the
  resolver's output before falling back to CUDA_PATH/bin, so
  pip-installed wheels are the canonical source (mirroring the
  Linux LD_LIBRARY_PATH ordering). System CUDA toolkit remains a
  valid fallback.

Tests: 7 new cases in
studio/backend/tests/test_llama_cpp_windows_nvidia_path.py:

* empty resolver when no nvidia wheels installed
* nvidia/<pkg>/bin layout resolved
* nvidia/<pkg>/Library/bin layout resolved
* mixed bin and Library/bin layouts both resolved
* unrelated site-packages contents not walked
* non-directory entries skipped
* missing prefix does not raise

110 backend tests pass. No regressions.

Refs #5106

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: also scan torch/lib in Windows pip nvidia DLL resolver

PyTorch's Windows CUDA wheels frequently bundle cudart64_X.dll and
cublas64_X.dll directly under Lib/site-packages/torch/lib/ instead of
shipping separate nvidia-cuda-runtime-cuXX / nvidia-cublas-cuXX wheels.
On those installs _windows_pip_nvidia_dll_dirs previously returned
nothing useful, and llama-server.exe fell back to needing a system CUDA
toolkit on PATH -- the original #5106 failure mode.

The install-side equivalent python_runtime_dirs in
install_llama_prebuilt.py already treats torch/lib as a Python runtime
DLL source for the same reason. Bring the runtime resolver in parity
so torch-bundled-CUDA installs find their cudart at llama-server start.

Updates the existing test that codified the bug (asserted torch/lib was
excluded), and adds three new cases: pickup, combined-with-nvidia, and
the must-be-a-directory guard.

* Studio: cover cu13 bin/x86_64 layout in Windows DLL resolver

Three follow-ups from a 12-reviewer batch over c1c8a074 (PR #5324):

1. The current nvidia-cuda-runtime (unsuffixed) 13.2.75 and
   nvidia-cublas 13.4.0.1 Windows wheels on PyPI ship under
   nvidia/cu13/bin/x86_64/cudart64_13.dll etc, not under
   nvidia/PKG/bin/. The previous resolver matched only one
   directory level past nvidia/PKG/ and silently missed the
   actual cu13 DLL location, leaving CUDA 13 users on the same
   failure mode as before #5106. Verified against:
       pip download nvidia-cuda-runtime --platform win_amd64
   which produces nvidia/cu13/bin/x86_64/cudart64_13.dll.

2. glob.glob over sys.prefix interprets [ and ] as a
   character class. Valid Windows usernames / install paths can
   contain those characters (for example C:\Users\alice[work]\studio),
   so the previous resolver silently returned an empty list for such
   prefixes even when DLL dirs were present.

3. The resolver only ever returned nvidia/PKG/bin -- if both
   bin and bin/x86_64 exist (current wheels do), Windows
   DLL search should land on the arch-specific subdir first so the
   explicit cudart64_X.dll location wins.

Rewritten as a pathlib.Path.iterdir walk to fix all three:
no glob escaping needed, arch-specific subdirs added explicitly,
and ordering puts bin/x86_64 before bin. Conda-style
Library/bin/x86_64 and Library/bin/x64 are also covered for
parity. A seen set dedupes when wheels happen to expose the
same directory through multiple layouts.

New tests:
 - test_picks_up_cu13_bin_x86_64_layout (the actual real-world cu13 case)
 - test_picks_up_bin_x64_layout
 - test_mixed_cu12_and_cu13_layouts
 - test_glob_meta_in_prefix_is_safe (bracket repro)
 - test_arch_subdir_listed_before_parent_bin (ordering)

Verified empirically against PyPI:
       nvidia-cuda-runtime 13.2.75 -> nvidia/cu13/bin/x86_64/cudart64_13.dll
       nvidia-cublas       13.4.0.1 -> nvidia/cu13/bin/x86_64/cublas64_13.dll
                                       nvidia/cu13/bin/x86_64/cublasLt64_13.dll
       nvidia-cudnn-cu13   9.22.0.52 -> nvidia/cudnn/bin/cudnn64_9.dll (already covered)

Refs #5106

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-11 05:42:09 -07:00
Roland Tannous
c57a97958a
Studio: stop truncating long log lines as suspected base64 (#5335)
* Studio: stop truncating long log lines as suspected base64

filter_sensitive_data carried a heuristic from the original Studio
import that truncated any string >100 chars containing ',' or '/'
to value[:20] + '...'. The block was dormant until #5246 wired
filter_sensitive_data into the structlog processor chain to redact
native-path leases. Once active, the heuristic ate normal log lines
- llama_cpp_backend's GGUF size summary, mmproj selection, the full
llama-server command line, and any traceback containing a path -
all rendered as a 20-char prefix, defeating debugging of llama-server
exceptions and GPU selection.

Drop the base64 truncation. No call site in the codebase logs raw
base64; if one ever does, it should truncate at the source rather
than in a global filter. Native-path lease redaction added by #5246
is preserved.

* Studio: regression test for filter_sensitive_data truncation

Pins two properties in studio/backend/loggers/handlers.py:

1. Long log messages with ',' or '/' (the GGUF size summary, mmproj
   selection, full llama-server command, exception tracebacks) flow
   through filter_sensitive_data unchanged. Exercises the exact call
   sites that regressed when #5246 wired the processor in.

2. Native-path lease redaction still fires for both the inline
   native_path_lease=... regex form and the nativePathLease dict-key
   form, so a future cleanup of the truncation logic can't quietly
   strip #5246's redaction along with it.
2026-05-08 13:07:18 +04:00
Lee Jackson
b65a7450ca
Studio: Dark theme refactor, right sidebar redesign, and chat UI polish (#5150)
* Dark theme refactor, right sidebar redesign, and chat UI polish

- Dark theme refactor
- Redesign right sidebar
- Further left sidebar adjustments
- Wider chat and content area; layout tweaks for chat content
- Rounded corners across elements for consistency
- Show chat message menu icons on menu-area hover, not only on message hover
- Assistant message menu icons now always visible; user messages keep on-hover
- Redesigned copy icon used consistently across chat blocks and messages
- Redesigned trash icon, applied consistently
- Unified icon sizing and style with the sidebar
- Adjusted icon colors across chat
- Fix on-hover background design for chat icons
- Fix tooltip from 'more' button staying visible after clicking elsewhere
- Adjust position and design of generation speed info text below messages
- Adjust design of token speed info popup
- Adjust sidebar scrollbar to cover recent chats only

* Recents sidebar rename, UI/theme refactor, layout and chat polish

UI & Theme:
- Dark theme refactor
- Consistent rounded corners across elements
- CSS polish and cleanup
- Remove unused logo image assets

Recents sidebar:
- Add 'more' button for options menu
- Support renaming conversations and training runs
- Confirmation dialog before deleting chats
- Add optional display_name column to training_runs (idempotent ALTER TABLE) so renaming doesn't lose model_name/dataset_name from the run config
- New PATCH /api/train/runs/{run_id} endpoint accepts { display_name: string | null }; empty/whitespace clears the override
- Sidebar shows display_name ?? model_name and exposes Rename in the row's More menu, mirroring the chat rename flow
- Cache last list response in localStorage and hydrate from it on mount, so recents paint instantly on F5 / route revisit; cached items are shape-validated and dropped if malformed
- Optimistic updates on rename and delete (apply locally + cache before background refresh)
- Visible toast on rename/delete failure instead of swallowed errors

Layout:
- Redesigned right sidebar
- Further left sidebar adjustments
- Updated chat content layout; chat and content area slightly widened
- Sidebar scrollbar covers recent chats only

Icons:
- Redesigned copy icon, unified across chat blocks and messages
- Redesigned trash icon to match
- Consistent icon sizing and style across chat and sidebar
- Adjusted icon colors across chat
- Fix icon on-hover background design

Chat messages:
- Menu icons now appear on hover over the menu area, not just the message
- Assistant message menu icons always visible; user messages keep on-hover (next/previous response stays visible for edited prompts)
- Repositioned and restyled generation speed info text below messages
- Restyled token generation speed popup

Tooltips:
- Removed tooltip on hover for previous/next assistant response icons
- Unified tooltip design across sidebars and chat
- Removed tooltip animations (also fixes related lag)

Model & Chat Template config:
- Merged Chat Template config into Model Configuration section
- Added revert-to-original for chat template
- Fix Chat Template config disappearing on page refresh until model reload

Performance & scroll:
- Removed chatbox movement animations across pages/navigation (fixes related UI lag)
- Fix scroll flicker at end of streaming when a code block is the final element
- Additional chat scroll improvements

Bug fixes:
- Fix 'more' button tooltip remaining visible after clicking elsewhere

* Remove sidebar localStorage cache and optimistic updates

Drops the localStorage hydration and optimistic rename/delete logic from the recents sidebar; reverts to fetching fresh on mount.

* Fix missing cn import in shared-composer (regression from merge)

* chore(sidebar): import sidebar deps from feature indexes

Re-export deleteChatItem / renameChatItem / useChatSidebarItems / SidebarItem / useChatSearchStore / ChatSearchDialog from @/features/chat, and removeTrainingUnloadGuard from @/features/training. Switch app-sidebar.tsx to consume them via the public feature indexes instead of deep paths, clearing the no-restricted-imports eslint errors. No behavior or UX change.

* fix(studio/frontend): reload training Recents sidebar after F5 refresh

The Recents sidebar showed empty after a hard refresh. The hook's inFlightRef dedup guard collided with React StrictMode's double-mount in dev: the second mount's fetch returned silently with no error, no retry, and no toast — leaving the sidebar empty until navigation.

Replace skip-if-busy dedup with abort-previous via a hook-level AbortController. This also fixes a latent race where a slow poll could resurrect a just-deleted row by clobbering the optimistic update.

Changes (all in use-training-history-sidebar.ts):
- fetchRuns aborts any in-flight request before starting a new one; post-await signal.aborted check drops stale responses.
- Optimistic helpers (applyRunUpdate, removeRun) abort in-flight fetches so they don't depend on caller discipline to invalidate stale data.
- Initial load gets bounded retry-with-backoff (500ms / 1.5s / 3.5s) and surfaces a sonner toast with a Retry action on final failure.
- Failure toast auto-dismisses on any successful load (initial retry, Retry click, or polling recovery).
- Polling pauses while the tab is hidden and catches up on visible, avoiding wasted requests during long training runs.
- Both effects own their teardown explicitly (abort + clear timer).

* Apply unified tooltip design and behavior across remaining pages for consistency

* UI polish: spacing, tooltip on source icons, letter spacing, smaller icons, consistent edit icon

- Adjust tiny spacing between elements around the UI for subtle polish
- Redesign tooltip on source icons for web search / tool use, consistent with the new design
- Adjust chat text letter spacing
- Smaller icon sizes
- Replace 'edit message' icon in chat with the new Rename icon used in Recents for consistency

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Adjust CSS for right sidebar

* Fix scrollbar UI compatibility across browsers

* fix: preserve chat preset settings on model load

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio): remove duplicate chat template status field

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* chore: remove creative preset assumption

* fix(studio): align speculative decoding default

* fix(studio/chat): snap numeric param inputs to step grid

- Type a value in any param input (Temperature, Top K, Max Tokens, etc.)
  now clamps to [min, max] and snaps to the slider's step grid, killing
  off-grid values like 1.051234 and FP residue from slider drags.
- Branch picker chevrons share the action bar's 32px height + 10px radius
  via a new .aui-branch-chevron-btn utility; hover area aligns visually
  while staying narrower than the sibling icon buttons.

* fix(studio/chat): keep training-run polls converging and drop dead preset code

- Keep training-run polls converging when responses outrun the 5s interval
  (don't unconditionally abort prior in-flight; skip if one is still pending,
  mutation race still guarded).
- Drop dead Creative/Precise preset code paths (remove 'builtin-fixed' source
  variant + unreachable branches).

* fix(studio): training-run cards show custom name + model + dataset

- Training-run cards now display custom display_name + model + dataset,
  with cross-view sync on rename/delete.
- Enhance clarity of borders and colors in dark theme on export etc.

* fix(studio): match active state green to unsloth brand color

* fix(studio): preserve can_resume on training rename

* fix(studio): keep GGUF chat template override distinct

* fix(studio): treat audio input models as multimodal

* fix(studio): cancel numeric draft on Escape

* fix(studio): use default speculative mode on toggle

* fix(studio): detect GGUF audio VLM input models

* fix(studio): address final PR review findings

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio): refresh sidebar/history when a new training run starts so it appears without a manual reload

* fix: API and svg

* fix(studio/sidebar): align run rename dirty check with displayed baseline

* fix(studio/sidebar): use leading-tight on account block to prevent descender clipping with truncate

---------

Co-authored-by: sneakr <hauzin@hotmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: shine1i <wasimysdev@gmail.com>
2026-05-07 14:33:31 +04:00
Avaya Aggarwal
0c803242ef
feat(studio): add Continued Pretraining (CPT) as a training method (#4677)
* feat(studio): add Continued Pretraining (CPT) support

Implements CPT as a first-class training method in Unsloth Studio,
resolving feature request #4565.

Changes:
- frontend/src/types/training.ts: add 'cpt' to TrainingMethod union
- frontend/src/lib/vram.ts: add 'cpt' to VramTrainingMethod (fp16 footprint)
- frontend/src/features/export/constants.ts: add CPT to METHOD_LABELS
- frontend/src/features/training/api/mappers.ts: map 'cpt' -> 'Continued Pretraining',
  force packing=true and train_on_completions=false for CPT payloads
- frontend/src/features/studio/sections/model-section.tsx: add 'Continued Pretraining'
  option (purple dot) to Method selector; update tooltip
- frontend/src/features/onboarding/.../model-selection-step.tsx: add CPT to
  onboarding wizard method dropdown
- backend/models/training.py: update training_type field description
- backend/core/training/worker.py: detect is_cpt flag, force packing=True,
  train_on_completions=False, pass is_cpt to _train_worker
- backend/core/training/trainer.py: _train_worker reads is_cpt kwarg, forces
  packing on, skips train_on_responses_only for raw-text pretraining

CPT behaviour:
- Full model weights (no LoRA adapters), same as Full Finetuning
- Sequence packing always enabled for GPU efficiency
- Trains on every token (no chat-format masking)
- VRAM estimated at fp16 (2.0 bytes/param)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update mappers.ts

* Add CPT raw dataset support and UI fixes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add missing training methods module

* Handle invalid raw-text rows and expose raw in onboarding

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: Etherll <mrmrmidessam@gmail.com>
2026-05-06 13:38:35 +04:00
Manan Shah
d65149795b
feat(studio): MLX training tab on Apple Silicon (LoRA / full FT, VLM, export) (#5265)
* Add Apple Silicon MLX routing

Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports
Extract original GPU init to _gpu_init.py (unchanged)
MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code
GPU path unchanged: from ._gpu_init import *

* Add Apple Silicon MLX routing

- Rewrite __init__.py: detect MLX on macOS arm64 before any torch imports
- Extract original GPU init to _gpu_init.py (unchanged)
- MLX path imports FastMLXModel from unsloth_zoo, skips all GPU code
- GPU path unchanged: from ._gpu_init import *

* mlx with studio

* mlx with studio

* updating temporary install.sh

* updating temporary install.sh

* adding t_v5 path

* adding t_v5 path

* fixing vision training

* fixing vision training

* adding chat

* adding chat

* minor

* minor

* Adding export and fixing training issues, inference with lora adaptors

* Adding export and fixing training issues, inference with lora adaptors

* fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM

* fix: MLX worker pass load_in_4bit, override is_vlm based on dataset, streaming for VLM

* Merge mlx-apple-silicon into main

* update install.sh to point to main branch

* update install.sh to point to main branch

* fix: export returns 3 values (success, message, output_path) matching upstream worker

* fix: export returns 3 values (success, message, output_path) matching upstream worker

* fix(mlx): show training-process peak memory in Studio UI, not system-wide

Studio UI was showing ~95 GB during MLX training because get_gpu_utilization
read "In use system memory" from IORegistry's AGXAccelerator — system-wide
GPU memory across all processes (training + backend + browser + Display).

Now the trainer's mx.get_peak_memory value is forwarded through the
progress event and surfaced via /api/train/hardware while training is
active. Falls back to the system-wide reading when training is not running.

* fix(mlx): show training-process peak memory in Studio UI, not system-wide

Studio UI was showing ~95 GB during MLX training because get_gpu_utilization
read "In use system memory" from IORegistry's AGXAccelerator — system-wide
GPU memory across all processes (training + backend + browser + Display).

Now the trainer's mx.get_peak_memory() value is forwarded through the
progress event and surfaced via /api/train/hardware while training is
active. Falls back to the system-wide reading when training is not running.

* fix(mlx): make is_bfloat16_supported detect M1/M2 (no native bf16)

M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70%
slower prefill compared to native fp16. M3+ have native bf16 (macOS
Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware
detection via mx.device_info.

* fix(mlx): make is_bfloat16_supported() detect M1/M2 (no native bf16)

M1 and M2 chips emulate bf16 in software on the GPU, causing 40-70%
slower prefill compared to native fp16. M3+ have native bf16 (macOS
Sonoma+ MPSGraph). Replaces the always-True stub with chip-aware
detection via mx.device_info().

* feat(mlx): wire training_type="Full Finetuning" through MLX worker

Compute use_lora from the UI's training_type before loading the model,
pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and
let the existing 'if use_lora' branch skip get_peft_model. Matches the
GPU worker's flow.

* feat(mlx): wire training_type="Full Finetuning" through MLX worker

Compute use_lora from the UI's training_type before loading the model,
pass full_finetuning=not use_lora to FastMLXModel.from_pretrained, and
let the existing 'if use_lora' branch skip get_peft_model. Matches the
GPU worker's flow.

* fix(mlx): pass save_method='merged_16bit' from Studio's export page

Previously the MLX path called save_pretrained_merged with no
save_method, which fell through to a no-op that didn't actually fuse
LoRA into the base. Now Studio's "Merged Model" export properly
fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU
behavior for the same UI option.

* fix(mlx): pass save_method='merged_16bit' from Studio's export page

Previously the MLX path called save_pretrained_merged() with no
save_method, which fell through to a no-op that didn't actually fuse
LoRA into the base. Now Studio's "Merged Model" export properly
fuses LoRA + dequantizes any 4-bit base to bf16, matching the GPU
behavior for the same UI option.

* fix(studio): pass private to MLX push, return 3-tuples consistently

MLX push_to_hub branch now forwards private=private (matches GPU)
Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model
needed') were tripping the route's 3-tuple unpack. Added a None
output_path so the unpack always succeeds.

* fix(studio): pass private to MLX push, return 3-tuples consistently

- MLX push_to_hub branch now forwards private=private (matches GPU)
- Existing 2-tuple early-returns ('repo_id+token required', 'PEFT model
  needed') were tripping the route's 3-tuple unpack. Added a None
  output_path so the unpack always succeeds.

* studio wirings

* studio wirings

* Merge pull request #5 from Manan17/feat/quant_config

studio wirings

* fix(mlx): wire train_on_completions for VLM via per-template lookup

Mirror the GPU worker: stop excluding VLMs and stop hardcoding
template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and
fetch the per-template instruction/response markers from
TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables
train_on_completions for vision+image and audio cases, so backend
just trusts the flag.

* fix(mlx): wire train_on_completions for VLM via per-template lookup

Mirror the GPU worker: stop excluding VLMs and stop hardcoding
template detection. Look up the model in MODEL_TO_TEMPLATE_MAPPER and
fetch the per-template instruction/response markers from
TEMPLATE_TO_RESPONSES_MAPPER. The frontend already force-disables
train_on_completions for vision+image and audio cases, so backend
just trusts the flag.

* wire in lora rslora, init lora weights, random_state

* wire in lora rslora, init lora weights, random_state

* loftq studio error message fix

* loftq studio error message fix

* handle unknown optim and lr scheduler

* handle unknown optim and lr scheduler

* Merge pull request #6 from Manan17/update/peftkwargs

Update/peftkwargs

* feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel

Studio's four UI checkboxes now actually flow through to MLX get_peft_model
(which was just updated in unsloth-zoo to honor them). Also drops the
incorrect train_projector wiring that tied projector LoRA to the
attn/mlp flags — those are language-side toggles, not projector toggles.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx): pass finetune_language/attention/mlp/vision flags to FastMLXModel

Studio's four UI checkboxes now actually flow through to MLX get_peft_model
(which was just updated in unsloth-zoo to honor them). Also drops the
incorrect train_projector wiring that tied projector LoRA to the
attn/mlp flags — those are language-side toggles, not projector toggles.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp

UI guardrail. The four checkboxes (vision/language/attention/MLP) carry
"scope × module-type" semantics that aren't obvious — picking just
"Attention modules" + "MLP modules" without "Language layers" naturally
reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune
attn/mlp modules in *no* tower" → empty target_modules → zero
trainable params → crash inside value_and_grad.

If user selected attn or mlp module types but no layer scope, default
to language scope. Power users can still explicitly choose
language=False, vision=True if they want vision-only fine-tuning of
attn/mlp.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx,ux): auto-imply finetune_language_layers when user picks attn/mlp

UI guardrail. The four checkboxes (vision/language/attention/MLP) carry
"scope × module-type" semantics that aren't obvious — picking just
"Attention modules" + "MLP modules" without "Language layers" naturally
reads as "fine-tune attn/mlp" but our backend reads it as "fine-tune
attn/mlp modules in *no* tower" → empty target_modules → zero
trainable params → crash inside value_and_grad.

If user selected attn or mlp module types but no layer scope, default
to language scope. Power users can still explicitly choose
language=False, vision=True if they want vision-only fine-tuning of
attn/mlp.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm

Inference UI sliders for top_k and repetition_penalty had no effect on
MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing
bug: mlx_vlm.generate_step expects temperature= (long form), but we
were passing temp= which silently fell into **kwargs — every VLM chat
was effectively greedy regardless of the temperature slider.

Text path (_generate_text):
make_sampler now receives top_k in addition to temp/top_p
make_logits_processors built and forwarded when repetition_penalty is
non-trivial (skip when 0.0/1.0 to avoid pointless overhead)

VLM path (_generate_vlm):
Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate
forwards them to generate_step's sampler/logits_processor builders)
Rename temp= → temperature= so it's actually consumed

Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and
Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10,
rep_pen=1.5} now produces a distinct output, proving the parameters
reach the sampler.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* fix(mlx): wire top_k, repetition_penalty, and VLM top_p through to mlx-lm/mlx-vlm

Inference UI sliders for top_k and repetition_penalty had no effect on
MLX, and VLM top_p was also silently dropped. Plus a latent pre-existing
bug: mlx_vlm.generate_step expects temperature= (long form), but we
were passing temp= which silently fell into **kwargs — every VLM chat
was effectively greedy regardless of the temperature slider.

Text path (_generate_text):
- make_sampler now receives top_k in addition to temp/top_p
- make_logits_processors built and forwarded when repetition_penalty is
  non-trivial (skip when 0.0/1.0 to avoid pointless overhead)

VLM path (_generate_vlm):
- Pass top_p, top_k, repetition_penalty as kwargs (mlx_vlm.stream_generate
  forwards them to generate_step's sampler/logits_processor builders)
- Rename temp= → temperature= so it's actually consumed

Verified end-to-end with a smoke test on Qwen2.5-0.5B-Instruct (text) and
Qwen2.5-VL-3B-Instruct (VLM): each of {greedy, top_p=0.5, top_k=10,
rep_pen=1.5} now produces a distinct output, proving the parameters
reach the sampler.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push

export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit"
(was hardcoded merged_16bit, ignoring the UI choice).
Both export_merged_model and export_base_model now pass save_directory=
to push_to_hub_merged so it reuses the just-written local folder
instead of re-saving under a relative "username/model" directory.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* feat(mlx): map format_type to MLX save_method, reuse local save dir for hub push

- export_merged_model: format_type="4-bit (FP4)" → save_method="merged_4bit"
  (was hardcoded merged_16bit, ignoring the UI choice).
- Both export_merged_model and export_base_model now pass save_directory=
  to push_to_hub_merged so it reuses the just-written local folder
  instead of re-saving under a relative "username/model" directory.

Co-Authored-By: Manan17 <shahmanan170602@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* restore install

* restore install

* fix(mlx): restore FastVisionModel as a distinct class

unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel`
right after defining `class FastVisionModel(FastLanguageModel)` with a
`for_training` static method. The alias erased the class binding, so
the documented `FastVisionModel.for_training(model)` call from upstream
Unsloth's VLM notebooks raised `AttributeError` on MLX.

Remove the offending alias. `FastVisionModel` is now a real subclass of
`FastLanguageModel` again — inherits `from_pretrained` /
`get_peft_model` / `for_inference`, exposes `for_training` as a no-op
pass-through (no-op because MLX doesn't have a train/eval mode flag;
the call exists purely for GPU/MLX notebook parity).

Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via
FastVisionModel.from_pretrained → get_peft_model → for_training →
MLXTrainer.train runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs,
peak 5.89 GB).

Studio's path (FastLanguageModel.from_pretrained for any repo,
auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8.

* fix(mlx): restore FastVisionModel as a distinct class

unsloth/__init__.py was assigning `FastVisionModel = FastLanguageModel`
right after defining `class FastVisionModel(FastLanguageModel)` with a
`for_training` static method. The alias erased the class binding, so
the documented `FastVisionModel.for_training(model)` call from upstream
Unsloth's VLM notebooks raised `AttributeError` on MLX.

Remove the offending alias. `FastVisionModel` is now a real subclass of
`FastLanguageModel` again — inherits `from_pretrained` /
`get_peft_model` / `for_inference`, exposes `for_training` as a no-op
pass-through (no-op because MLX doesn't have a train/eval mode flag;
the call exists purely for GPU/MLX notebook parity).

Verified end-to-end: Qwen3-VL-2B + LaTeX_OCR LoRA + vision LoRA via
FastVisionModel.from_pretrained → get_peft_model → for_training →
MLXTrainer.train() runs 10 steps cleanly (loss 1.10 → 0.12, no NaNs,
peak 5.89 GB).

Studio's path (FastLanguageModel.from_pretrained for any repo,
auto-detect VLM in the loader) is unaffected. Tier-1 review finding #8.

* Studio: harden MLX training and export, restore GPU init guards

Studio export
Restore Tuple[bool, str, Optional[str]] contract on export_merged_model,
export_base_model, export_gguf, and export_lora_adapter, populating
output_path on successful local saves so routes/worker/CLI/frontend
details.output_path is non-empty again.
Lift the GPU save_method assignment out of the local-save branch so
Hub-only merged exports (save_directory='', push_to_hub=True) no longer
hit UnboundLocalError on the push branch.
For MLX merged and base hub-only export, stage to a tempfile.TemporaryDirectory
before push_to_hub_merged instead of passing save_directory=''.
Source _IS_MLX from unsloth instead of recomputing the platform check
(single source of truth, also enforces mlx-package availability).

Studio MLX training/inference
Pass token=hf_token into FastMLXModel.from_pretrained for gated/private
models, matching the inference path.
Strip hf_token and wandb_token from wandb.init(config=...) so secrets
do not leak into the W&B run config.
Replace load_from_disk(local_datasets[0]) with the existing
UnslothTrainer._resolve_local_files / _loader_for_files helpers so
uploaded JSON/JSONL/CSV/Parquet files train through the normal datasets
loader (load_from_disk still used for HF save_to_disk directories).
Make the dataset slice helper inclusive at the end and treat 0 as a real
index instead of "unset", matching the GPU and embedding paths.
Add a status_message -> message alias inside _send so the existing parent
pump (training.py) renders MLX status updates instead of blanks.
Forward min_p through generate_chat_response into _generate_text /
_generate_vlm and into make_sampler / vlm_kwargs so the sampling control
is no longer a no-op on MLX.
Wrap unsloth_zoo.mlx_loader / mlx_trainer imports with a clearer
ImportError pointing users at install.sh for Apple Silicon.
Exit the MLX stop-polling thread on EOFError/OSError instead of
busy-looping when the queue/pipe is permanently closed (one-line
why-safe rationale inline).

Studio frontend
ParamsSection subscribes to platform deviceType via the Zustand hook so
the gradient checkpointing dropdown re-renders after the async device
fetch completes.

Studio hardware
get_gpu_utilization MLX branch now reads _read_apple_gpu_stats once and
derives VRAM totals from psutil, removing the second ioreg subprocess
per utilization poll.

Unsloth core
Restore the os.geteuid == 0 guard around the CUDA ldconfig recovery
that was lost when GPU initialization moved into _gpu_init.py, plus the
non-root manual-fix warning branch. Non-root CUDA users no longer shell
out to ldconfig at import time.
Load dataprep/raw_text via importlib so the MLX import path no longer
pulls torch in through dataprep/__init__.py -> synthetic.py.
FastVisionModel.from_pretrained overrides the inherited delegator only
to inject text_only=False; this is an extension, not a duplication, and
is needed so VLM checkpoint loads keep the vision tower.
Wrap the MLX-branch unsloth_zoo import with a clearer ImportError.

* Studio: regression tests for MLX training/export and GPU init ldconfig guard

tests/python/test_gpu_init_ldconfig_guard.py asserts the geteuid root
check still wraps the ldconfig recovery and the non-root branch warns
bnb users; AST + source-text inspection so the test runs without torch.
tests/studio/test_export_output_path_contract.py covers the
Tuple[bool, str, Optional[str]] return contract on every export method,
the output_path assignment after successful local save, the Hub-only
GPU save_method binding fix, the MLX hub-only TemporaryDirectory
staging, and the single-source `_IS_MLX` import from unsloth.
tests/studio/test_mlx_training_worker_behaviors.py covers token
forwarding to FastMLXModel.from_pretrained, wandb config secret
stripping, file-aware local dataset loading, status_message ->
message aliasing, inclusive slice semantics, EOFError/OSError stop
thread exit, and the friendly mlx_loader / mlx_trainer ImportError.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(mlx): cap inference memory + release wired on unload + tame worker pre-pin

Three memory-hardening fixes for Studio's MLX path:

1. Inference applies the same Metal caps as the trainer.
   load_model previously only called set_wired_limit(100% of recommended)
   with no upper memory_limit, leaving large VLM checkpoints unbounded
   during the loader allocation. Add _configure_memory_limits() that sets
   memory_limit to 85% of recommended and wired_limit to min(recommended,
   memory_limit) — matching MLXTrainer's defaults so behavior is the same
   whether the user trains or just runs inference.

2. unload_model releases pinned memory back to the OS — but only when
   the cache is empty. Without this, pinned wired bytes stayed allocated
   to MLX after the model was gone, starving other apps. The release is
   guarded on `not self.models` so unloading one of several cached
   models doesn't un-pin weights still in use.

3. Worker pre-cap is conservative instead of aggressive.
   The previous pre-pin set_wired_limit(100% of recommended) competed
   with MLXTrainer's later more conservative cap. Replace with the same
   85%-memory / min(rec, memory) pair that the trainer applies later
   (idempotent re-apply). Bounds the model load + LoRA setup window
   without over-pinning.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tests/studio: regression tests for the _IS_MLX dispatch gate

Two gates drive every MLX-vs-CUDA dispatch decision in Studio:

  1. unsloth._IS_MLX in unsloth/__init__.py — evaluated once at import
     time, read by Studio worker code to choose the GPU vs MLX trainer
     and inference paths. Defined as
        Darwin AND arm64 AND find_spec("mlx") is not None.

  2. utils.hardware.detect_hardware() — runtime probe with priority
     CUDA > XPU > MLX > CPU. The MLX branch is reached only when both
     CUDA and XPU are unavailable and the host is Apple Silicon and
     mlx is importable.

Neither gate had a direct test. Adds tests/studio/test_is_mlx_dispatch_gate.py
with six tests:

  test_is_mlx_gate_uses_three_required_predicates
      AST-walks unsloth/__init__.py and asserts the _IS_MLX assignment
      is a BoolOp(And) of platform.system()=="Darwin",
      platform.machine()=="arm64", and find_spec("mlx") is not None.
      Catches accidental rewrites that drop a predicate.

  test_is_mlx_gate_true_on_apple_silicon_with_mlx_present
      Spoofs platform to Darwin/arm64, injects a fake mlx module so
      find_spec returns a real ModuleSpec, re-evaluates the gate
      expression. Verifies it flips True under the exact conditions
      Studio expects.

  test_is_mlx_gate_false_when_mlx_missing
      Spoofs Apple Silicon but with mlx absent. Verifies the gate stays
      False (so a Mac without mlx installed does not pretend to have
      MLX support).

  test_is_mlx_gate_false_on_non_apple_silicon
      Canary on the actual Linux+CUDA / AMD / Intel test host: the gate
      must remain False regardless of whether mlx happens to be
      importable. Protects existing GPU users from accidental MLX
      hijack when MLX support evolves.

  test_detect_hardware_picks_mlx_when_only_apple_silicon_available
      Forces torch.cuda and torch.xpu off, spoofs Apple Silicon, injects
      fake mlx and mlx.core. detect_hardware() must return DeviceType.MLX.

  test_detect_hardware_picks_cuda_on_real_host
      Canary: on a real CUDA host detect_hardware() must return
      DeviceType.CUDA. Protects against the MLX branch shadowing CUDA
      dispatch on NVIDIA / AMD ROCm hosts.

Uses the same monkeypatch.setitem(sys.modules, ...) fake-mlx pattern as
the existing test_mlx_inference_backend.py — no new test infrastructure,
no real mlx install required.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add AGPL-3.0 SPDX header to Studio MLX regression tests

Four Studio MLX test files shipped without an SPDX-License-Identifier:

  studio/backend/tests/test_mlx_training_worker_config.py
  tests/studio/test_mlx_training_worker_behaviors.py
  tests/studio/test_export_output_path_contract.py
  tests/studio/test_is_mlx_dispatch_gate.py

They sit in or alongside studio/backend/, which is governed by
studio/LICENSE.AGPL-3.0, and exercise AGPL Studio code. Add the same
"# SPDX-License-Identifier: AGPL-3.0-only" header that's already on
test_mlx_inference_backend.py so the license declaration matches
the code under test rather than defaulting to the repo-root
Apache-2.0.

* Wrap MLX submodule imports with friendly install hint

The _IS_MLX block at the top of unsloth/__init__.py already catches the
missing-package case with a friendly install hint, but the follow-up
"from unsloth_zoo.mlx_trainer import ..." and "from unsloth_zoo.mlx_loader import ..."
lines run unguarded. An Apple Silicon user who has unsloth-zoo installed
but on an older version (e.g. the current PyPI release, before the MLX
modules ship) sees a raw ImportError on the submodule rather than the
hint that points at install.sh.

Wrap the two submodule imports in the same try/except shape so the
friendly install message fires whether the package is missing entirely
or just predates the MLX submodules. No-op once both packages release
together; smooths the transitional window where unsloth/main has merged
but unsloth-zoo on PyPI has not.

---------

Co-authored-by: DoubleMathew <mmathew23@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-05 23:54:58 -07:00
Datta Nimmaturi
09505fcc6e
Update VRAM estimator to cater to broader model configs (#5175)
* Update VRAM estimator to cater to broader model configs

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix attn backend check, better support for MoE etc

* Studio: tighten VRAM estimator structured-shape and attention paths

- Conservative attention fallback: when resolve_attention_implementation
  fails, charge the quadratic non-flash activation path instead of
  silently keeping the optimistic flash_attention_2 default.
- Resolve attention on a shallow config copy so _set_attn_impl does not
  mutate the cached config returned by _load_config_for_gpu_estimate.
- Use getattr for AutoModelForCausalLM._model_mapping to avoid raising
  on private-attribute renames in transformers.
- Treat sdpa as O(n) linear attention; PyTorch SDPA dispatches to flash
  or memory-efficient backends, only eager needs the quadratic term.
- Per-layer activation accounting: structured archs (head_dim,
  layer_types, attention_k_eq_v, num_kv_shared_layers, double-wide MLP)
  now flow into compute_activation_bytes via _text_linear_dims, instead
  of using the legacy hidden_size//num_attention_heads KV/MLP shape.
- Exclude MLA configs (q_lora_rank set) from the structured-shape path
  so q_lora low-rank projection formulas keep applying when head_dim is
  also present.
- _build_text_module_elements emits a single MLA self_attn aggregate
  using _compute_attn_elements when q_lora_rank is set, avoiding the
  ~10% overcount that fed into _compute_skipped_quantizable_elements.
- Restrict _module_path_matches to known text-tower prefixes so VLM
  skip names like vision_tower.model.layers.<i>.self_attn.q_proj no
  longer falsely shadow the text alias model.layers.<i>.self_attn.q_proj.
- Pick up enable_moe_block from the config and add the per-layer dense
  MLP alongside the MoE experts in compute_total_params and
  compute_lora_params (Gemma4-style parallel dense + MoE block).
- Single-pass structured layer accounting in _compute_layer_elements,
  removing the duplicate _text_linear_dims walks.
- Drop the now-zero (activations - activations_computed) shard term in
  VramBreakdown.min_gpu_vram and the stale comment that referred to it.
- attention_implementation typed as Optional[str] to match call sites
  that pass None.
- Inline rationale comments on DOUBLE_QUANT_4BIT_FACTOR and
  NON_FLASH_ATTENTION_FACTOR pointing at VRAM_ESTIMATION.md.

* Studio: extend parallel-MoE accounting + non-prefix dense layer support

- Apply enable_moe_block / moe_has_dense_mlp symmetrically: activation
  per-layer MLP size in _layer_qkv_mlp_sizes now adds the parallel dense
  MLP for MoE layers, matching the weight and LoRA accounting added in
  the prior commit. Skip-quantizable mapping in _build_text_module_elements
  now registers both mlp.experts and per-projection mlp.{name} entries
  for MoE layers when the parallel dense block is present, so an
  llm_int8_skip_modules entry like "model.layers.N.mlp" covers both.
- Track dense layer indices as a tuple (dense_layer_indices) extracted
  from first_k_dense_replace or decoder_sparse_step + mlp_only_layers,
  and dispatch dense-vs-MoE accounting through _is_dense_mlp_layer. The
  prior count-based path silently mis-bucketed layers when mlp_only_layers
  was non-prefix (e.g. [3, 5] on an 8-layer model). num_dense_layers is
  derived from len(dense_layer_indices) for backward compatibility.
- Drop the redundant ">0" check in _is_kv_shared_layer so configs with
  num_kv_shared_layers == num_hidden_layers (every layer shared) are
  correctly recognized as shared.
- Refresh VRAM_ESTIMATION.md section 5 to note that sdpa joins
  flash_attention_2 in the linear activation path; refresh the
  VramBreakdown.activations_computed comment now that the activation
  floor is gone.

* Studio: Gemma4 PLE accounting, flex_attention, KV-share guard restore

- Add flex_attention to LINEAR_ATTENTION_IMPLS. Unsloth's
  resolve_attention_implementation returns "flex_attention" when
  HAS_FLASH_ATTENTION is False and the model class supports flex; PyTorch
  FlexAttention is a memory-efficient kernel, not a quadratic eager
  attention path. Without this, activation estimates over-charge ~36x.
- Restore the `> 0` guard in _is_kv_shared_layer. Transformers Gemma4
  (modeling_gemma4.py:1031, modular_gemma4.py:863, :926) uses
  `layer_idx >= first_kv_shared_layer_idx > 0`, so configs that mark
  every layer as KV-shared raise on construction. Reverting the
  unconditional acceptance avoids producing a detailed estimate for a
  shape the actual model code rejects.
- Extend the parallel dense MLP path (`enable_moe_block`) in
  _build_text_module_elements: when the arch is non-structured, use
  arch.intermediate_size for the dense gate/up/down dims instead of
  _text_linear_dims (which returns moe_intermediate_size via
  _get_mlp_size). Prior code under-counted skipped quantizable elements
  for the parallel dense block by up to 8x on GLM-style configs.
- Add Gemma4 per-layer-input (PLE) module accounting:
  per_layer_model_projection (one global Linear) plus per-layer
  per_layer_input_gate and per_layer_projection are added to the
  quantizable text-linear total in _compute_layer_elements;
  post_per_layer_input_norm and per_layer_projection_norm flow into
  the non-quantizable bucket. compute_lora_params adds the same three
  Linear modules to the all-linear total. References:
  transformers_versions/5.7.0/.../gemma4/modular_gemma4.py:1077-1083,
  :1247-1253.
- VRAM_ESTIMATION.md section 5 now lists flex_attention alongside sdpa
  and flash_attention_2 as linear-memory backends.

* Studio: shared-expert variants, mlp_layer_types dispatch, PLE skip, all-linear str, deepcopy resolver

Five targeted estimator corrections:

- _compute_dense_layer_indices now reads `mlp_layer_types` ahead of
  `first_k_dense_replace` / `decoder_sparse_step`. Transformers Exaone-MoE,
  Laguna, Hy_v3, GLM-MoE-DSA, GLM4-MoE-Lite, Ernie4_5_VL_MoE etc. ship the
  per-position list and may omit the prefix-style fields entirely.
- _build_text_module_elements registers per_layer_input_gate /
  per_layer_projection (per layer) and per_layer_model_projection (global)
  in the canonical element map and alias map. The PLE element count was
  added to total_quantizable in a prior commit but skip-module matching
  against names like model.layers.0.per_layer_input_gate produced 0-byte
  delta. Layer aggregate text.layers.<i> now sums all layer modules so
  prefix skip names cover the PLE pieces too.
- _targets_all_linear coerces a bare string `"all-linear"` to `["all-linear"]`
  before set comparison; the previous set comprehension iterated chars.
  PEFT LoraConfig.target_modules accepts the bare-string convention.
- ModelArchConfig gains `shared_expert_intermediate_size`. extract_arch_config
  reads `n_shared_experts` / `num_shared_experts` aliases and infers
  `n_shared_experts=1` when only `shared_expert_intermediate_size` is set.
  _compute_moe_mlp_elements and the structured + non-structured LoRA paths
  size the shared expert with its own intermediate (Qwen3.5-MoE: 512 vs
  routed moe_intermediate_size).
- _determine_attention_impl_for_gpu_estimate uses copy.deepcopy so the
  resolver does not mutate nested text_config on the cached source.
  PreTrainedConfig._attn_implementation setter walks `sub_configs` and the
  prior shallow copy still touched the inner objects.

* Studio: extend MoE/PLE/KV-share accounting to activation and skip-alias paths

Five activation-path corrections plus two LoRA / skip-alias corrections so
that shared-expert, per-layer-input, and KV-shared-layer support is symmetric
across weights, LoRA, skip-quantizable, and activation paths.

- _layer_qkv_mlp_sizes: include shared-expert FFN in mlp_size (live shared
  expert per token alongside routed experts) and keep K/V activation memory
  for KV-shared layers; only the WEIGHT path uses has_k/has_v from
  _layer_attention_dims.
- _per_layer_activation_bytes / compute_activation_bytes: account for
  per_layer_input_gate (hd-sized) and per_layer_projection (pli-sized) per
  layer plus the global per_layer_model_projection [B,S,L,PLI] tensor when
  hidden_size_per_layer_input is set.
- _build_text_module_elements: split mlp.experts into routed and
  mlp.shared_expert canonical entries; register layers.<i>.experts alias for
  Gemma4 enable_moe_block layouts and mlp.shared_experts (plural) alias for
  Exaone-MoE / Laguna / GLM4-MoE-Lite shared-expert variants.
- _compute_moe_mlp_elements: split into _compute_routed_moe_elements and
  _compute_shared_moe_elements; only count shared_expert_gate (hd->1 Linear
  per shared expert) when shared_expert_intermediate_size is set, which is
  the Qwen2-MoE / Qwen3.5-MoE discriminator. Other shared-expert families
  (Exaone-MoE, HY-V3, GLM4-MoE-Lite, Laguna) lack the gate.
- compute_lora_params: when target_modules='all-linear' bare keyword, drop
  routed and shared MoE expert LoRA contributions. PEFT's all-linear targets
  nn.Linear only; Unsloth's get_moe_target_parameters expands MoE expert
  nn.Parameter LoRA only when target_modules contains explicit
  gate_proj/up_proj/down_proj/gate_up_proj names.
- _per_layer_input_lora_params: thread target_modules through and add the
  per-PLE-module contribution when the corresponding name appears, not only
  under all-linear.

* Studio: top-k MoE activations, ERNIE list configs, suffix skips, multimodal full bytes

Six estimator corrections aligning the detailed accounting paths with real
training behavior:

- _layer_qkv_mlp_sizes scales the MoE-layer mlp_size by num_experts_per_tok
  so the active routed-expert intermediate tensors are charged for activations.
  Adds num_experts_per_tok to ModelArchConfig and extracts it from
  num_experts_per_tok / top_k_experts (Gemma4 alias) in extract_arch_config.
- compute_lora_params splits routed and shared MoE LoRA contributions so that
  bare target_modules='all-linear' zeroes routed (nn.Parameter expert tensors,
  which Unsloth's get_moe_target_parameters does NOT enable for the bare
  keyword) but keeps shared-expert LoRA (regular nn.Linear MLPs that
  Unsloth's get_peft_regex DOES match).
- extract_arch_config gains a _first_scalar helper for ERNIE-style
  moe_intermediate_size = [routed, shared] lists, plus moe_num_experts and
  moe_num_shared_experts attribute aliases. When moe_intermediate_size is a
  pair and shared_expert_intermediate_size is unset, the second element is
  treated as the shared-expert intermediate.
- estimate_required_model_memory_gb's detailed branch retains
  max(0, model_size_bytes - compute_total_params(arch) * 2) on top of the
  arch-derived breakdown.model_weights so multimodal models (vision/audio
  towers) and partially-modeled families (Gemma3n AltUp/Laurel etc.) do not
  silently drop bytes that the safetensors total includes.
- _module_path_matches accepts a tail-only match when the skip entry is
  shorter than the alias path. Transformers' BNB quantizer suffix-matches
  short skip entries like ['q_proj'] / ['lm_head'] against full module
  paths; the previous len(skip) < len(alias) early-return missed those.
- _per_layer_input_lora_params drops the all_linear branch and only counts
  PLE LoRA when the user explicitly names per_layer_input_gate /
  per_layer_projection / per_layer_model_projection. Unsloth's
  get_peft_regex requires module names to contain a component tag
  (mlp/attn/...); PLE module names lack any tag, so all-linear training
  does not attach LoRA to them.

* Studio: full-FT extra optimizer/gradient inflation, MoE top-k aliases, ERNIE position dispatch, sibling experts aggregate

When the safetensors total exceeds the text-arch fp16 estimate (multimodal
vision/audio towers, partially-modeled families), only inflate the model
weights line for adapter methods but extend optimizer + gradient bytes
under full fine-tuning, where the extra params are trainable.

DBRX exposes top-k routing as moe_top_k and Hunyuan-V1-MoE as moe_topk;
neither is aliased to num_experts_per_tok via attribute_map, so probe both
when extracting arch config.

ERNIE 4.5 MoE / VL MoE configs declare MoE layers via
moe_layer_start_index / moe_layer_end_index / moe_layer_interval (with -1
meaning the last layer); add the position-style dispatch alongside the
existing mlp_layer_types / first_k_dense_replace / decoder_sparse_step
paths.

When moe_has_dense_mlp is set (Gemma4 enable_moe_block) the routed experts
live as a sibling of self.mlp at layers.<i>.experts in the actual model
layout; keep the layer mlp aggregate to the dense path and add a separate
experts aggregate so a skip module model.layers.<i>.mlp does not collapse
the routed experts as well.

* Studio: extend MoE family extraction (Llama4 / DBRX / Hunyuan / ERNIE) and align dense vs routed MLP widths

- Llama4: pick up `config.moe_layers` (auto-populated from
  interleave_moe_layer_step) so dense layer indices reflect the actual
  is_moe_layer dispatch.
- Llama4: add a separate `dense_intermediate_size` derived from
  `intermediate_size_mlp` (used for the dense feed_forward path) and keep
  `intermediate_size` for the routed/shared expert width. Auto-attach one
  shared expert per MoE layer when the dense-vs-MoE width split is present.
- DBRX: walk the `ffn_config` sub-config when extracting MoE attrs
  (moe_num_experts / moe_top_k / ffn_hidden_size). Without this DBRX is
  misclassified as a dense arch.
- Hunyuan: normalize layer-wise `moe_topk` (and the canonical
  `num_experts_per_tok` lookup it shadows via attribute_map) through a
  worst-case scalar so the int(...) cast cannot crash on list values.
- ERNIE 4.5 MoE: switch the start/end/interval dispatch to the model's
  `(layer_idx + 1) % interval == 0` modulo gate so MoE layers match the
  decoder when interval > 1.
- ERNIE 4.5 VL MoE: drop the heuristic that read
  `moe_intermediate_size[1]` as the shared expert width; in VL configs [1]
  is the vision-routed width and shared experts are sized from [0].
- estimate_fp16_model_size_bytes: prefer the larger of config-derived and
  local-weight bytes so the multimodal extra_bytes correction can fire
  for local VLM directories.

* Add tests for VRAM estimator extensions

* Studio: trim verbose comments in VRAM estimator

Collapse multi-paragraph rationale blocks to 1-3 lines stating the single
load-bearing fact. Fix one inverted "fall through ... last" comment whose
claim disagreed with the surrounding code.

* Consolidate added tests into existing test_vram_estimation.py and test_gpu_selection.py

Move Llama4 / DBRX / ERNIE arch-extraction tests into test_vram_estimation.py
as TestLlama4ArchExtraction / TestDbrxFfnConfigExtraction /
TestErniePhaseModuloDispatch / TestErnieVlSharedExpertWidth classes. Move
estimate_fp16_model_size_bytes prefer-larger-of-config-or-local tests into
test_gpu_selection.py as TestEstimateFp16ModelSizeBytesPrefersLocalWeights.
Drop one redundant Llama4 num_dense_layers assertion already covered by the
moe_layers dispatch test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
2026-05-05 04:12:36 -07:00
Datta Nimmaturi
6b13cab746
fix KVCache estimates for gemma4 style sliding window models (#5225)
* fix KVCache estimates for gemma4 style sliding window models

Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: add per-arch SWA pattern fallback + n_kv_heads mirror for PR #5225

The pattern-aware SWA estimator added in this PR only fires when the
GGUF carries `<arch>.attention.sliding_window_pattern`. Today's
Gemma-2 / Gemma-3 / Gemma-3n / gpt-oss / Phi-3 GGUFs ship
`attention.sliding_window` but not the pattern field (llama.cpp's
converter strips it), so the new branch is bypassed and we fall back to
the legacy 1/4-global heuristic on the most popular SWA arches in our
catalogue (gemma3 alone has 6+ variants in the unsloth/* top 30 by
downloads, plus gpt-oss-20b/120b).

Two additions on top of this PR:

1. `_SWA_PATTERN_DEFAULTS_BY_ARCH` table keyed by GGUF arch name. When
   the GGUF reports a sliding window but no pattern, we synthesise the
   pattern from the architecture's canonical period (gemma2=2,
   gemma3=6, gemma3n=5, gpt_oss=2, phi3=1, cohere2=4). Periods sourced
   from a survey of the top 150 unsloth/* HF configs against
   `text_config.layer_types` and `Gemma*Config.sliding_window_pattern`.

2. Mirror `_n_kv_heads_by_layer` into the scalar `_n_kv_heads` (using
   max as a conservative upper bound) when the head_count_kv array is
   read. Without this, any non-SWA estimator path (GQA, legacy) on a
   Gemma-4-style model falls through to `n_heads`, which can be many
   times larger than the real per-layer KV head count. Also let
   `_can_estimate_kv` accept the array directly as belt-and-suspenders.

End-to-end check on `unsloth/gemma-3-270m-it-Q4_K_M.gguf` (18 layers,
sliding_window=512, no pattern field): the parser now resolves the
pattern to period=6 (3 global, 15 SWA), matching the actual
Gemma3TextConfig default. KV estimate at 32k context drops from
141 MB (legacy 1/4) to 108 MB (per-layer), a 23% reduction that
directly translates into more headroom for `_fit_context_to_vram` and
fewer cases where the slider lands on the 4096 floor.

Tests: extended `test_kv_cache_estimation.py` with
`TestArchSwaPatternDefaults` covering the six tabled arches, an
unknown-arch negative, explicit-pattern precedence, and a
no-sliding-window negative; updated `test_array_fields_parsed` to
reflect the new mirror semantics; updated
`test_end_to_end_synthetic_swa` to use the period=6 expectation. All
102 tests in the kv-cache / context-fit / max-context suites pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: tighten arch SWA table to verified non-regressing entries

Audit of every unsloth/* HF model (1334 repos, all config.json fetched
in scripts/survey_all_unsloth.py) plus end-to-end checks against five
real GGUFs (gemma-3-270m, gemma-3-1b, qwen2.5-0.5b, phi-3.5-mini,
falcon-h1-0.5b, granite-4.1-8b) confirms:

* Pure-GQA arches (llama, qwen3, mistral3, glm4, llama4, ...) and the
  qwen2 family with use_sliding_window=False all reach Path 4 GQA
  cleanly. The llama.cpp converter strips `attention.sliding_window`
  for qwen2/qwen2_vl/qwen2_5_vl when use_sliding_window=False, so the
  SWA path never fires for them. Verified on Qwen2.5-0.5B-GGUF: no
  sliding_window field in metadata.
* MLA arches (deepseek_v3/v32/v4, glm4_moe_lite, glm_moe_dsa, kimi_k25)
  emit `kv_lora_rank` -> Path 1 fires correctly.
* Hybrid Mamba/Attn arches that emit both ssm.* and
  full_attention_interval (qwen3_5, qwen3_5_moe, qwen3_next) -> Path 2
  fires correctly.

Two table changes:

1. Drop the `phi3` entry. Phi-3 GGUFs emit
   `phi3.attention.sliding_window=262144` but never emit
   `attention.key_length`/`value_length`, so the SWA path is gated
   off and the estimator falls to the legacy formula. The huge
   sliding_window also means SWA layers and global layers cache
   identical numbers of tokens at any practical context, so a fallback
   would be a no-op anyway. The previous `phi3: 1` entry was also
   semantically wrong: period=1 with the (i+1)%N!=0 rule produces
   all-global, not the all-SWA you'd want for Phi-3.
2. Document the audit findings in the table comment, including the
   two arches we deliberately skip (phi3, qwen2*) and the one
   architecture family that is not a regression vs. main but is also
   not yet optimal (mistral v0.1/v0.2 all-SWA every-layer cannot be
   expressed with the period sentinel).

Tests: added `test_non_swa_arch_uses_full_attention_path` parametrized
over llama / qwen2 / qwen3 / mistral / mistral3 / glm4 / llama4 to
pin the invariant that pure-GQA arches never receive a synthetic SWA
pattern. Removed phi3 from the parametrize list of
`test_arch_default_pattern_applied`. All 108 tests pass.

Known separately tracked (not addressed here): falcon-h1 GGUFs ship
ssm.* + key_length but no full_attention_interval, so the hybrid
path 2 cannot fire and the estimator falls to GQA path 4, which
counts every block as an attention layer. Affects 8 unsloth/* repos
(~500 downloads). Same gap exists for granitemoehybrid-class GGUFs.
Worth a follow-up that adds either a HYBRID_ATTENTION_INTERVAL_BY_ARCH
table or a tensor-name probe.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: make SWA pattern resolver dynamic so new models work without code changes

The static `_SWA_PATTERN_DEFAULTS_BY_ARCH` dict only covered
architectures we knew about at PR-merge time. New SWA models would
need a code change here every time a new arch shipped, which doesn't
scale. Replaced with a 4-tier resolver so any newly-released model
that lands on Hugging Face with a normal `config.json` is covered
automatically:

  Tier 0 (parser)  -- explicit GGUF metadata if the converter emits it
                      (BOOL array or scalar period). Already supported.
  Tier 1 (cache)   -- $UNSLOTH_STUDIO_HOME/swa_cache.json. Populated by
                      previous Tier 3 fetches. Survives restarts.
  Tier 2 (bootstrap) -- `_BOOTSTRAP_SWA_DEFAULTS` shipped with Studio.
                      Same five entries as the old static table
                      (gemma2/3/3n/gpt_oss/cohere2). Lets fully-offline
                      installs keep working for popular SWA arches
                      with zero network.
  Tier 3 (HF fetch)  -- pulls `config.json` from the GGUF's source HF
                      repo and reads `sliding_window_pattern` (int) or
                      `text_config.layer_types` (string array). Result
                      is cached to Tier 1 so subsequent loads are
                      offline-fast. Disabled by
                      `UNSLOTH_STUDIO_OFFLINE=1`. Network errors and
                      missing repos fall through silently.
  Tier 4 (caller)  -- legacy 1/4-global SWA estimate (unchanged).

The GGUF parser now also extracts a handful of `general.*` keys
(`source.huggingface.repository`, `source.url`, `source.repo_url`,
`base_model.0.repo_url`, `base_model.0.organization` + `.name`,
`organization` + `basename`) so the resolver has source-repo
candidates to try.

End-to-end smoke against a brand-new arch (`never_seen_before_arch`,
not in the bootstrap dict) pointing at `google/gemma-3-1b-it`:
resolver fetched the HF config, derived period=6, materialised the
26-layer mask with 4 global layers (indices 5/11/17/23), and wrote
`{"never_seen_before_arch": 6}` to the on-disk cache. Next load hits
Tier 1 with no network.

Tests: added `TestDynamicSwaResolver` with 10 tests covering each
tier (period derivation, aperiodic mask handling, URL parsing,
bootstrap precedence, cache precedence, HF fetch + persistence,
candidate fallback, offline env knob, network failure). All 118
tests in the kv-cache / context-fit / max-context suites pass.

The `_SWA_PATTERN_DEFAULTS_BY_ARCH` name was retired in favour of
`_BOOTSTRAP_SWA_DEFAULTS` to make the tier semantics explicit.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: add Tier 2.5 transformers introspection to SWA resolver

Slots a new tier between bootstrap and HF fetch that asks the
locally-installed `transformers` package directly. Two strategies, in
order, both offline-friendly:

  a. Default-instantiate the matching `Config` class via
     `CONFIG_MAPPING[arch]()` and read `sliding_window_pattern` /
     `text_config.layer_types`. Drills into `text_config` for
     multimodal wrappers.
  b. `inspect.getsource(cfg_class)` regex-parse for
     `sliding_window_pattern: int = N` defaults. Catches configs
     whose constructor raises (missing required args), or where the
     default is bound only in the __init__ signature. Walks
     `cfg_class.sub_configs["text_config"]` too so multimodal wrappers
     that delegate to a TextConfig still get inspected.

Resolver chain is now 5 tiers: GGUF metadata, on-disk cache,
bootstrap defaults, transformers introspection, HF Hub fetch, legacy
fallback. Tier 2.5 results are persisted to the same on-disk cache as
Tier 3 so subsequent loads skip the import overhead.

`_arch_aliases` normalises hyphen vs underscore variants (`falcon-h1`
vs `falcon_h1`) since GGUF and HF disagree for a handful of arches.

Cross-version verification (probe at `temp/swa_probe/`):

```
arch       transformers 4.57.6      transformers 5.7.0
gemma3     6                        6
gemma2     2                        2
cohere2    4                        4
gpt_oss    2                        2
gemma3n    5                        5
gemma4     ARCH-MISSING             6        <- new arch picked up automatically
falcon_h1  None                     [True]*32 <- per-layer mask used verbatim
phi3       None                     None
mistral    None                     None
qwen2      1 (all-global)           1
llama      None                     None
deepseek_v3 None                    None
```

The `gemma4` and `falcon_h1` rows are the headline: a brand-new arch
that lands in transformers (gemma4 is 5.x-only) is supported by the
resolver the moment a user upgrades the package, with zero edits to
this file. Same applies to any future arch with a `Config` class.

Tests: added `TestTransformersIntrospection` with 6 cases covering
arch-alias normalisation, real-arch resolution against the live
transformers, inspect.getsource fallback when default-init raises,
graceful behaviour when transformers is unavailable, unknown-arch
returns None, and Tier-2.5-before-Tier-3 ordering. Also adjusted the
existing Tier 3 failure test to mock Tier 2.5 out so it specifically
exercises the network-failure path. All 124 tests in the kv-cache /
context-fit / max-context suites pass.

Updated the module-level resolver comment from "4-tier" to "5-tier"
to document the new tier.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: consolidate verbose comments and docstrings in SWA resolver

Net -345 lines: -210 in llama_cpp.py, -357/+111 in
test_kv_cache_estimation.py. No behaviour change; only comments,
docstrings, and one tests-only `_SWA_FIELDS` helper to remove the
copy-pasted GGUF metadata dict from each resolver test.

Code changes only delete or shorten:
  * 5-tier resolver header collapsed from a 38-line block diagram to
    a 9-line summary; the rest is the function bodies.
  * Bootstrap dict per-arch comments collapsed to one-line `Config`
    references.
  * `_swa_cache_path`, `_save_swa_cache`, `_period_from_layer_types`,
    `_arch_aliases`, `_swa_entry_from_config_obj`,
    `_resolve_swa_pattern`, `_resolve_swa_entry_from_transformers`
    docstrings stripped to one line or removed when the body is
    self-evident.
  * Tier-by-tier inline comments inside `_resolve_swa_pattern` removed
    (function body reads top-to-bottom in tier order).
  * Path-3 SWA estimator comment shortened from a 12-line tier
    breakdown to 3 lines.
  * Parser fallback comment block (originally explained the resolver
    in-line) trimmed to two lines pointing at the resolver.
  * `_can_estimate_kv` legacy-clause comment shortened to one line.
  * GGUF `general.*` WANTED block comment shortened to one line.

Test changes:
  * Per-test docstrings dropped where the test name and body already
    explain intent.
  * Class-level docstrings reduced to one line.
  * Common GGUF field dict factored to module-level `_SWA_FIELDS`.
  * Multi-line URL/list assertions collapsed to one-liners.

All 124 tests pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: account for SWA cache double-buffering in path 3 estimate

Cross-check against llama.cpp ground truth (running llama-server with
--parallel 1 and reading the `llama_kv_cache: size = X MiB ( N cells,
M layers, ... )` log lines) showed the SWA path under-counted by
~20% on Gemma-3-shaped models:

  GGUF                                pred MiB  actual MiB  ratio
  gemma-3-270m-it-Q4_K_M               31.50    39.00       0.81
  gemma-3-1b-it-Q2_K                   43.00    54.00       0.80

Root cause: llama.cpp double-buffers the SWA cache so it can keep the
current and next windows during the shift, allocating
`2 * sliding_window` cells per SWA layer (capped at n_ctx). My formula
was using `min(n_ctx, sliding_window)` instead of
`min(n_ctx, 2 * sliding_window)`. Verified directly:

  llama_kv_cache_iswa: creating SWA KV cache, size = 1024 cells
  llama_kv_cache: size = 15.00 MiB (1024 cells, 15 layers, ...)

with `gemma3.attention.sliding_window = 512` -> 2 * 512 = 1024 cells.

Fix: introduce `swa_cells = min(n_ctx, 2 * swa)` in path 3 and use
that for both the per-layer-pattern branch and the legacy
1/4-global fallback.

Re-run after the fix:

  GGUF                                pred MiB  actual MiB  ratio
  gemma-3-270m-it-Q4_K_M               39.00    39.00       1.000
  gemma-3-1b-it-Q2_K                   54.00    54.00       1.000
  qwen2.5-0.5b-instruct-q4_k_m         96.00    96.00       1.000
  Phi-3.5-mini-instruct-Q4_K_M       3072.00  3072.00       1.000
  Falcon-H1-0.5B-Instruct-Q4_K_M      144.00   144.00       1.000
  granite-4.1-8b-Q3_K_M              1280.00  1280.00       1.000

All 5 paths now match llama.cpp's actual allocation exactly under
single-sequence inference (Studio's default).

Tests: updated `test_gemma3`, `test_gpt_oss`,
`test_gemma4_per_layer_swa_metadata`, `test_ctx_smaller_than_window`,
`test_odd_layer_count`, and `test_end_to_end_synthetic_swa` to use
the doubled SWA cell count. All 124 tests in the kv-cache /
context-fit / max-context suites pass.

* studio: tolerate truncated GGUF input so resolver fallback still runs

Wraps each iteration of the GGUF KV-pair loop in a try/except that
breaks out cleanly on `struct.error` or `UnicodeDecodeError`, instead
of letting the outer try eat the exception and skip the SWA resolver
fallback at the end.

The motivating use case is reading the GGUF metadata via an HF Hub
HTTP byte-range fetch. The first ~128 KiB of a typical GGUF contains
all the metadata we need (arch, block_count, attention.*, sliding
window, ssm, MLA fields, plus the tokenizer config) -- but for models
with large tokenizer vocabs (Gemma 3 has 262144 tokens) the tokenizer
arrays spill past the 128 KiB boundary. The truncation used to bubble
out as `unpack requires a buffer of 8 bytes`, abandoning the resolver
fallback and leaving us with no SWA pattern (so the SWA path fell
through to the legacy 1/4 estimate).

Verified end to end against `unsloth/gemma-3-1b-it-GGUF`:

  Range-fetch first 128 KiB of `gemma-3-1b-it-Q2_K.gguf` over HTTP
  (HTTP 206 Partial Content), parse:

    arch              = gemma3
    block_count       = 26
    attention.sliding_window = 512
    sliding_window_pattern   = set (4 global) <- via Tier 2 bootstrap
    KV @ ctx=8192            = 54.00 MiB     <- matches llama.cpp
                                              ground truth

This means Studio can preview KV-cache requirements (and therefore
auto-context fit) for any HF GGUF without downloading the weights.

All 124 tests pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: thread llama-server KV flags through the estimator

Adds keyword-only knobs to _estimate_kv_cache_bytes and _fit_context_to_vram
that mirror the llama-server CLI options that change KV memory:

  --swa-full         (swa_full)        SWA layers cache the full n_ctx
                                       instead of 2 * sliding_window cells
  --parallel N       (n_parallel)      number of server slots
  --kv-unified       (kv_unified)      single shared KV buffer; when off,
                                       multiplies KV by n_parallel
  --ctx-checkpoints  (ctx_checkpoints) per-slot SWA snapshots, each one
                                       sliding-window of state per SWA layer
  --kv-offload       (kv_on_gpu)       when off, KV lives in CPU RAM and is
                                       not subtracted from the VRAM budget

Defaults preserve the previous behavior (swa_full=False, n_parallel=1,
kv_unified=True, ctx_checkpoints=0, kv_on_gpu=True) so existing call sites
are unaffected. All five paths (MLA, hybrid, SWA pattern, SWA fallback,
GQA, legacy) now apply the per-slot replication factor; the SWA paths
also honor swa_full and add the checkpoint term when applicable.

Tests: TestServerFlags (17 cases) covers every flag, the no-op cases, the
swa_full + ctx_checkpoints interaction, slot multiplication on each path,
and the kv_on_gpu shortcut in _fit_context_to_vram.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: account for shared_kv_layers (Gemma 3n / Gemma 4)

Gemma 3n and Gemma 4 set <arch>.attention.shared_kv_layers in the GGUF
metadata (convert_hf_to_gguf.py lines 7578 and 7712). The trailing N
layers of the model reuse KV from earlier layers and don't allocate
their own cache, so n_layers in the per-layer formulas overcounted by
exactly that many blocks. For google/gemma-3n-E4B-it (35 layers, 15
shared), this puts ~43% of the KV estimate back on the table.

Changes:

- _read_gguf_metadata parses <arch>.attention.shared_kv_layers into
  self._shared_kv_layers; init / unload / reparse all reset it.
- _estimate_kv_cache_bytes computes n_layers_kv = max(1, n_layers -
  shared_kv_layers) and substitutes it for n_layers in:
    Path 1 (MLA), Path 3 (SWA pattern loop bound and the no-pattern
    fallback), Path 4 (GQA), Path 5 (legacy). Path 2 (hybrid) keeps
    n_layers since hybrid + shared_kv combined isn't a thing today and
    the semantics would need to specify which attention layers are
    shared.
- max(1, ...) floor protects against pathological GGUFs where shared
  >= n_layers.
- Composes naturally with --swa-full, --kv-unified / --parallel,
  --ctx-checkpoints, and the per-layer SWA pattern from the dynamic
  resolver. When the field is unset (every other arch) the math is
  byte-identical to before.

Tests: TestSharedKVLayers (13 cases) covers each path's drop, the
no-op-when-unset case, the floor at one layer, composition with the
server-flag knobs, and lifecycle reset. test_end_to_end_synthetic_shared_kv_round_trip
exercises the full GGUF parse -> estimate path on a synthetic gemma3n_text
blob. Existing TestLifecycle tests extended to cover the new field.

Full suite: 132 passing.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: only stub httpx in tests when the real lib is missing

The unit suite stubs httpx unconditionally so tests can run on a minimal
Python install. Surfaced during a fresh-venv simulation: when the stub
is installed via setdefault on a system that DOES have httpx,
huggingface_hub.errors fails to import HTTPError / Response at module
load time, which the transformers introspection tier swallows via its
bare except. Result: TestTransformersIntrospection passes in venvs
where httpx happened to be imported first (workspace) and silently
fails in venvs where it doesn't (fresh uv venv).

Switch to "only stub when real lib unavailable", and round out the stub
with HTTPError, RequestError, and Response so any test environment
without httpx still gets a complete enough surface for huggingface_hub
to import.

* studio: per-layer-type --parallel N memory accounting for SWA

Empirical verification against llama-server (see
workspace_5/temp/sim_pr5225/probe_parallel_full_matrix.py and
verify_parallel_matches_server.py) showed the prior whole-cache
slot_factor multiplication in _estimate_kv_cache_bytes was wrong for
n_parallel > 1. The actual rule, verified bit-exact across the full
(parallel x ctx) grid for both SWA and pure-GQA models:

  * non-SWA layers: total cells = n_ctx, partitioned across slots
                    (per-slot ctx = n_ctx / parallel). Total memory
                    is CONSTANT in n_parallel.
  * SWA layers:     per-slot cells = 2 * sliding_window (clamped at
                    n_ctx and at per_slot_ctx when ctx is split among
                    many slots). Total memory grows LINEARLY in
                    n_parallel.
  * --kv-unified:   no measurable difference to total memory; both
                    modes yield the same byte total in measured cases.
                    Retained as accepted-but-ignored kwarg for API
                    forward-compat.

Closed form (Path 3 with per-layer pattern):
    total_kv = sum_global_layers(n_ctx * n_kv * (k+v) * bpe)
             + parallel * sum_swa_layers(
                   min(2*sliding_window, n_ctx, n_ctx//parallel)
                   * n_kv_layer * (k_swa + v_swa) * bpe
               )
    + parallel * checkpoint_extra_per_slot   (when ctx_checkpoints > 0)

Changes to _estimate_kv_cache_bytes:
- Path 3 (SWA pattern): accumulate global_bytes and swa_bytes_per_slot
  separately; final result = global_bytes + slots * (swa_bps + cp_bps).
- Path 3 (no-pattern fallback): same split using the 1/4-global heuristic.
- Paths 1 / 2 / 4 / 5: drop the slot_factor multiplication. Non-SWA
  caches don't scale with --parallel.
- swa_full=True: SWA cells = per_slot_ctx (was n_ctx), so slots
  cancels out and total stays constant. Matches llama-server's
  --swa-full --parallel N output exactly.

Production wiring fix in start():
- Seven internal calls to _estimate_kv_cache_bytes / _fit_context_to_vram
  used the default n_parallel=1, even though load_model accepts the
  caller's n_parallel value (forwarded to llama-server via --parallel
  on the command line). Pass n_parallel through all seven so VRAM
  budgeting is correct when an operator sets parallel slots above 1.
  Studio's default ships at 1 so production today is unaffected;
  this completes the wiring for operators who tune it.

Tests:
- TestParallelSWAScaling (10 new cases): closed-form invariants per
  path, swa_full + parallel collapse, kv_unified no-op proof,
  per-slot SWA cell clamping, and the empirical Gemma-3 270m formula
  (24 + parallel * 15 MiB at ctx=8192) baked from the verifier.
- TestServerFlags: rewrote 4 assertions and renamed 2 to reflect the
  per-layer rule; non-SWA paths now correctly assert constancy.
- TestSharedKVLayers::test_composes_with_n_parallel: rewrote to assert
  only the SWA portion of the unshared layers scales.

Backward compatibility: at n_parallel=1 the output is bit-identical to
before this change (verified across 120,960 sweep combinations and the
141-test suite in both workspace and fresh-uv-venv environments).
Verifier output at --parallel in {1,2,4,8} x ctx in {4096,8192,16384}
shows ratio 1.000 against llama-server for both SWA and pure-GQA
models (24/24 cells exact match).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: accept new estimator kwargs in load-time test stubs

`test_llama_cpp_context_fit.py` and `test_llama_cpp_max_context_threshold.py`
patch `_estimate_kv_cache_bytes` with constant per-token stubs and then
call the real `_fit_context_to_vram`. After 29dcf96e threaded the
llama-server flag kwargs (`swa_full`, `n_parallel`, `kv_unified`,
`ctx_checkpoints`) through `_fit_context_to_vram`, the production method
forwards them to the stubbed estimator and the old positional-only stubs
raise `TypeError`.

These two suites exercise the load-time fit decision and the max-context
threshold property with a constant per-token KV cost; SWA / parallel-slot
accounting is intentionally out of scope, so the stubs absorb the new
kwargs and ignore them. No production change.

Restores both files to fully passing: 15/15 in `test_llama_cpp_context_fit`
and 8/8 in `test_llama_cpp_max_context_threshold`. Combined with the
existing 141/141 in `test_kv_cache_estimation`, the three KV-cache test
modules are 164/164 green.

---------

Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-05 04:06:46 -07:00
Roland Tannous
0da8af56d6
unsloth run: add --enable-tools/--disable-tools server-side tool policy (#5277)
* Add process-level tool_policy state for unsloth run

* Apply tool_policy override at chat/completions, /messages, and tool pass-through gates

* Add pure resolver for unsloth run --enable-tools/--disable-tools

* Wire --enable-tools/--disable-tools into unsloth run

* Color tool-policy notices and confirmation prompt in Claude orange

* Always show tool-status notice; print URL + API key in silent mode

* Treat any non-loopback bind as external; forward --yes after parent prompt

* Fix tool_policy double-module bug: import via state.tool_policy to share global with routes
2026-05-05 12:45:15 +04:00
Roland Tannous
dbea77e347
Studio: forward llama-server args from unsloth studio run , activate unsloth run , and allow passing model:quant to load models (#5271)
* Studio: forward unknown CLI args directly to llama-server

`unsloth studio run --model X --top-k 20 --chat-template-file foo.jinja`
now passes the unknown flags through to the llama-server subprocess.
Adds a denylist for flags Studio manages (port, -m, -c, --api-key, -ngl,
--flash-attn, --no-context-shift, --jinja, GPU-fit, model-identity, ...)
that returns HTTP 400 on collision. HTTP callers can supply the same
list via LoadRequest.llama_extra_args.

* Studio: accept `--model org/repo:variant` shorthand in `unsloth studio run`

Mirrors llama.cpp's `-hf <repo>:<quant>` and ollama's pull syntax so
`unsloth studio run --model unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL` is
equivalent to `--model unsloth/... --gguf-variant UD-Q4_K_XL`. Local
paths and Windows drive letters are preserved verbatim. If both an
embedded variant and an explicit `--gguf-variant` are given and they
disagree, the command fails with a clear error.

* Studio: register `unsloth run` as alias for `unsloth studio run`

Top-level `unsloth run --model ...` is now equivalent to
`unsloth studio run --model ...`. Same context_settings, so unknown
flags continue to pass through to llama-server.

* Studio: let users override soft-managed llama-server flags from CLI

Trims the denylist to flags Studio fundamentally cannot share with
the user (model identity, --host/--port/--path/--api-prefix,
--api-key, --ssl-*, --webui, --models-*). Soft-managed flags --
-c/--ctx-size, --parallel, --flash-attn, --no-context-shift,
--jinja, -ngl, -t/--threads, --fit* -- now pass through and override
Studio's auto-set version via llama.cpp's last-wins CLI parsing.

Lets users tune their run on the spot:
  unsloth run --model X -c 131072 --parallel 1 --threads 32

* Studio: accept `-hf` / `-hfr` / `--hf-repo` as aliases for `--model`

Matches llama-server's `-hf <repo>:<quant>` spelling so users coming
from llama.cpp can use the same flag. Typer claims the aliases before
the pass-through validator runs, so the HTTP-API denylist on those
flags is unaffected.

  unsloth run -hf unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL
2026-05-04 17:08:04 +04:00
Roland Tannous
35ab5da93c
Default Studio host to 127.0.0.1 and prompt before auto-start (#5267)
Studio bound to 0.0.0.0 by default and the installer silently auto-started
a server at end of install, exposing it on the network without consent and
contradicting the privacy-first / local-only guarantee.

- studio/backend/run.py: run_server() and argparse --host default to 127.0.0.1
- unsloth_cli/commands/studio.py: studio_default() and run() --host default to 127.0.0.1
- install.sh: drop -H 0.0.0.0 from generated launcher template; replace silent
  auto-start with a [Y/n] prompt; add cloud/network note to manual hint
- install.ps1: drop -H 0.0.0.0 from PowerShell launcher template; replace
  silent auto-start with a Read-Host [Y/n] prompt; add cloud/network note
- studio/setup.sh: drop -H 0.0.0.0 from launch hint; add cloud/network note
- README.md: simplify launch examples to `unsloth studio -p 8888`; note
  -H 0.0.0.0 is available for cloud/LAN use

Tests:
- studio/backend/tests/test_host_defaults.py
- tests/studio/test_cli_studio_defaults.py
- tests/sh/test_install_host_defaults.sh
2026-05-04 13:03:16 +04:00
Roland Tannous
456a49a350
Add Qwen3.6 support (#5257)
* qwen3.6 unsloth studio support

* Add qwen3.6 causal-conv1d detection

* Update model_mappings.py

moved qwen3.6-27B to thinking train on completion template

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-02 23:30:57 +04:00
Lee Jackson
05f46686de
Studio: Fix chat template disappearing after browser refresh (#5209)
* fix: preserve chat template on refresh

* chore: simplify chat template status lookup
2026-05-01 08:19:09 -07:00
Daniel Han
a5615426a5
Studio: fix 7 failing studio_unit_tests on main (#5216)
* Studio: fix 4 failing studio_unit_tests on main

Three of the failing tests had drifted from production:

1. test_health_response_reports_desktop_capability_fields stubbed
   `routes` with a SimpleNamespace that omitted `inference_studio_router`,
   so importing studio.backend.main raised ImportError. Add the missing
   router stub.

2. test_local_recipe_token_preserves_desktop_marker and
   test_local_recipe_token_keeps_web_marker_absent decoded the local
   provider's api_key as a JWT, but _inject_local_providers now mints
   a unified sk-unsloth-* internal API key (not a forwarded JWT), so
   jwt.decode raised "Not enough segments". Renamed and rewrote both
   tests to validate the API-key contract: starts with
   storage.API_KEY_PREFIX and authenticates via get_current_subject as
   the real admin user. The web vs desktop distinction is irrelevant
   at this layer because the unified API-key path does not carry
   session flags.

The fourth failure was a real production bug:

3. test_github_validate_skips_live_access_with_honest_note expected
   github-seed validation to return valid=True per
   _GITHUB_VALIDATE_NOTE ("GitHub access and rate limits are checked
   when the run starts"). The validate route called
   build_config_builder which lazy-imports the optional data_designer
   module; when it is missing, the bare except blocked the recipe.
   Catch ImportError specifically and treat it as a deferred check,
   matching the documented intent.

Verified all 4 tests pass and the rest of studio/backend/tests still
pass (608 total, with the only remaining failures being environment
specific: 4 GPU-aware tests on a no-GPU host and 1 Anthropic-API
smoke test, both unrelated).

* Studio: fix 3 test_gpu_selection route tests after load_model signature change

`routes/inference.load_model` gained a `fastapi_request: Request`
positional argument (used to read `app.state.llama_parallel_slots`
inside the GGUF path), but the three TestRouteErrors cases that
exercise the early validation path were not updated and failed with
`TypeError: load_model() missing 1 required positional argument:
'fastapi_request'`.

Pass a SimpleNamespace mock that satisfies the attribute path the
production code reads. The validation under test fires before the
mock is consumed, but supplying the realistic shape protects against
regressions if the validation order changes.

Affected tests:
- test_inference_route_rejects_gpu_ids_for_gguf
- test_inference_route_returns_400_for_invalid_gpu_ids
- test_inference_route_returns_400_for_uuid_parent_visibility_gpu_ids

* Studio: address review feedback on validate.py ImportError handling

Two reviewers flagged the ImportError bypass added in b0d33cf:

- chatgpt-codex-connector[bot]: catching bare ImportError marks recipes
  as valid even when build_config_builder fails for unrelated import
  problems (broken internal imports, missing transitive deps after a
  version bump), hiding real regressions until run start.
- gemini-code-assist[bot]: silent pass discourages troubleshooting;
  the deferred-validation case should be logged at debug level.

Tighten the bypass to ModuleNotFoundError where the missing module name
starts with "data_designer". Other ImportErrors propagate to the outer
handler and surface as validation failures, restoring the visibility
the reviewers asked for. Add a debug-level log entry that names the
missing module so operators can trace why validation deferred.
2026-04-28 22:43:44 -07:00
Lee Jackson
ff759ba7e4
Studio: Fix image-only chat requests failing validation (#5212)
* fix: allow image-only chat messages

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* test: deduplicate empty content validation coverage

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-28 14:49:13 -07:00
Daniel Han
b09aa82a3a
Studio: add github_repo seed reader and GitHub Support Bot recipe (#5169)
* Studio: add github_repo seed reader and GitHub Support Bot recipe

Adds a first-party Data Designer seed reader that scrapes GitHub issues,
pull requests, and commits from one or more repositories via the GraphQL
API, and a learning recipe (GitHub Support Bot) that turns those rows into
synthetic support Q&A pairs for fine-tuning.

Backend (new plugin studio/backend/plugins/data-designer-github-repo-seed):
* GitHubRepoSeedSource config: repos, token (falls back to GH_TOKEN /
  GITHUB_TOKEN env var), item_types (issues / pulls / commits),
  per-resource limit (0 means all), max_comments_per_item.
* Rate-limit-aware GraphQL client (GitHubClient + RepoScraper) shared
  across repos; flattens each item into a uniform row with columns
  item_type, repo, number, title, body, state, author, created_at,
  closed_at, url, labels, comments.
* Registered via the data_designer.plugins entry point.

Frontend:
* New seed_github block variant so the seed node card shows
  "GitHub repositories" instead of the generic "Document file"
  placeholder, with its own icon and inline summary (repo count +
  item-type list).
* Rewritten seed dialog github_repo form: repos textarea pre-filled with
  unslothai/unsloth + unslothai/unsloth-zoo, password input for the GH
  token, items-per-repo number with an "All" toggle, and the noisier
  options (item types, max comments, include comments) tucked under an
  Advanced collapsible.
* Local model auto-load on Run: if a recipe uses an is_local provider
  and the inference server is not already serving that model, the
  executions hook calls /api/inference/load first. Removes the "open
  /chat to load a model" prerequisite that users kept tripping on.
* Honor the recipe's run.rows value in the Run dialog (previously the
  store reset to 5 regardless of what the template shipped).

Recipe (studio/frontend/src/features/data-recipes/learning-recipes/
github-support-bot.json):
* Defaults to the Local Model provider + unsloth/gemma-4-E2B-it-GGUF.
* Scrapes unslothai/unsloth and unslothai/unsloth-zoo, issues and pulls,
  up to 100 items per resource.
* Two LLM blocks: normalized_question (llm-text) rewrites each thread
  into a clean support question, support_answer (llm-structured)
  produces JSON with answer / diagnosis_questions / cites / confidence.
* Run defaults to 10 rows for a quick smoke test.

Verified end-to-end on a running Studio: card renders, source-data
dialog is pre-populated, All toggle disables the limit input, the
recipe executes and produces rows against a loaded local GGUF.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: improve GitHub recipe support

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: speed up GitHub scraper and harden the support-bot recipe

Addresses a perf issue found while demoing the github_repo seed reader:

Scraper is too slow at scale. The PRs GraphQL query pulls deeply nested
fields (reviewThreads, reviews, commits, timelineItems, etc.) so the
page size was pinned at 3 to stay under GitHub's node-count ceiling. 100
PRs meant 34 serial round trips. Added lighter query variants
(PRS_PAGE_QUERY_LIGHT, ISSUES_PAGE_QUERY_LIGHT) that drop the fields the
Studio flatten layer does not use (it only reads title, body, state,
author, labels, comments). With the light query PR pages can safely go
to 25 per page and issues to 50. The plugin scraper now passes
light=True to RepoScraper so Studio always uses the fast path; the heavy
query remains available for other callers.

Recipe defaults are now demo-ready with production knobs called out:
- max_parallel_requests: 1 and max_tokens: 800 so small local models
  stay stable when running the support_answer structured column.
- support_answer prompt trimmed to 80-200 words so gemma-4-E2B GGUF can
  actually comply with the schema. The canonical 150-300 word codex
  prompt is still documented in the node3 markdown note for
  production upgrades.

* Studio: rename GitHub recipe to 'GitHub Scraper' and add Easy mode

Changes the recipe framing from a single-purpose 'Support Bot' pipeline
to a general-purpose scraper that produces {user_request,
grounded_response} training pairs. Aligns with the canonical
github_data_gatherer dataset (11 enrichment tasks mirrored in pr_requests_20
/ issue_requests_20 on the input side and explain_pr / issue_fix_plan /
issue_solution on the output side).

Recipe JSON changes:
- columns[0] renamed normalized_question -> user_request, prompt now
  inverts a GitHub thread into a realistic user ask instead of
  normalising it.
- columns[1] renamed support_answer -> coauthor_response, emits
  {response, followups, cites, task, confidence} and branches on
  issue vs PR thread type.
- Notes rewritten to document the 11-task catalog and the canonical
  production prompt to paste in for a full dataset backfill.

Frontend: Easy mode for github_repo recipes. The drag-and-drop canvas is
hidden behind an 'Advanced' tab; Easy mode is the default for any recipe
whose seed_source_type is github_repo. The Easy form reuses the existing
GithubRepoSeedForm (promoted to exported), adds a rows input bound to
previewRows, a model field bound to the model_config, and a single Run
button that calls runPreview() directly (no modal). Non-github recipes
see the same Editor / Runs tabs as before.

View mode persists per-recipe-id in localStorage under
recipe-studio:view-mode:<recipeId>.

* Studio: auto-detect server GH_TOKEN and widen Easy-mode detection

The GitHub seed form now fetches /api/data-recipe/seed/github/env-token
on mount and, when the server exposes a GH_TOKEN / GITHUB_TOKEN env var
and the token field is blank, shows a small 'Using server env var' badge
and swaps the placeholder text. The token value itself is never returned
to the UI.

Widens Easy-mode detection in recipe-studio-page.tsx so that recipes
saved before ui.seed_source_type was persisted also get the Easy tab:
falls back to recipe.seed_config.source.seed_type, which is always
present for github_repo seeds.

* fix: polish GitHub recipe UI

* Studio: default llama-server --threads to -1 (auto)

Previously we passed --threads only when the caller set an explicit
value, which meant llama-server fell back to its internal default.
That default has varied across llama.cpp builds (some versions use
hardware concurrency including hyperthreads, which hurts throughput on
CPU-heavy inference). Always passing --threads -1 pins the behaviour
to llama.cpp's auto-detect (physical cores).

Caller-supplied n_threads still wins when non-None.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: auto-switch Easy mode to Runs pane on run start

Easy mode had no progress island or canvas overlay, so after clicking Run
the only visible state was the button label flipping to "Running..." while
the screen otherwise stayed identical. This reads as stuck even though the
job is progressing.

Wire an onExecutionStart callback from recipe-studio-page.tsx through to
useRecipeExecutions so that when a run is kicked off from easy mode, the
page flips to the executions view where the Runs sidebar, progress bar,
rate/ETA panel, and live log are rendered. Advanced/editor mode keeps its
existing behavior and stays on the canvas (it already has the floating
ExecutionProgressIsland).

* fix: clean up GitHub scraper layout

* Studio: forward llm-structured output_format as llama-server response_format

Local GGUF runs of llm-structured columns used to generate the full
max_tokens budget before the prompt-level "return JSON in a ```json
fence" instruction got parsed. Small models (e.g. gemma-4-E2B-it)
routinely broke format, so each row took ~65s and frequently failed
with "No parsable JSON structure within ```json markdown fence".

For any local-provider model_config referenced by an llm-structured
column, clone the model_config and inject response_format into the
clone's inference_parameters. Uses llama.cpp server's flat shape
(tools/server/README.md):

    {"type": "json_schema", "schema": <output_format>}

Not the OpenAI-nested form; data_designer's OpenAI adapter forwards
response_format verbatim via facade._COMPLETION_REQUEST_FIELDS, and
llama-server's documented schema path expects the flat variant.

The clone is per (model_alias, column) so:
- llm-text / llm-judge columns that share the same alias keep
  free-form sampling.
- Each structured column gets its own schema, so columns with
  different output_formats don't collide.

Effect on gemma-4-E2B-it demos: every row parses cleanly, and the
model terminates immediately after the closing brace instead of
running to max_tokens. Net wall-clock is usually faster even though
grammar-constrained sampling is slightly slower per token.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: flip Easy to Runs pane before validation scrape, not after

Previously onExecutionStart fired inside runExecution, which runs AFTER
validateRecipe() -- and validation re-invokes the seed reader. For the
github_repo reader that is a full GraphQL scrape, so the user sat on a
"Running..." button with an otherwise unchanged Easy form for 10-15s
before anything moved.

Call onExecutionStart at the top of runWithValidation, right after we
have a payload to send. The view flips immediately; ensureLocalModelLoaded
+ validateRecipe now run against the Runs pane instead of a frozen Easy
form. runExecution still calls onExecutionStart downstream, but the
callback is idempotent (the page's easy -> executions guard skips the
second call), so no behaviour change for runs that pass validation.

If validation fails the toast + runErrors path still fires; the Easy
form's error banner still reads runErrors when the user switches back.

* Studio: unify data-recipe workflow auth on sk-unsloth-* keys

The previous commit (a61b4cc9) assumed storage.create_api_key(..., internal=True)
and storage.revoke_internal_api_key(key_id) existed, but those helpers were
only in the working tree, never committed. Recipe runs in local-model mode
were therefore crashing with 500 when _inject_local_providers tried to mint
a workflow key. This commit ships the missing pieces.

auth/storage.py:
- api_keys schema gains is_internal INTEGER DEFAULT 0 (with a guarded
  ALTER TABLE migration so existing auth.db files upgrade in place).
- create_api_key takes an internal=False kwarg; internal keys are flagged
  so they can be hidden from user-facing listings.
- list_api_keys takes include_internal=False so UIs never see workflow keys.
- New revoke_internal_api_key(key_id): id-only revoke for keys minted by
  non-user subjects (the JobManager does not know a username).

core/data_recipe/jobs/manager.py:
- JobManager.start accepts internal_api_key_id and stores it on Job so
  lifecycle handlers can revoke eagerly.
- _handle_event revokes on EVENT_JOB_COMPLETED / _ERROR / _CANCELLED.
- _pump_loop subprocess-died fallback also retires the key so a crashed
  worker cannot leak a live sk-unsloth-* beyond its TTL.
- Revocation is best-effort (swallow exceptions) -- the 24h TTL is the
  safety net if storage hiccups.

core/data_recipe/jobs/types.py:
- Job dataclass gains internal_api_key_id: int | None = None.

Replaces the bespoke 24h JWT path that jobs.py used to mint for local
providers. One mint/revoke/verify surface for every API key the server
issues, and revocation is now eager (seconds, not 24h) instead of TTL-only.

* Studio: plug workflow-key leak on unexpected create_job errors

Review follow-up on the sk-unsloth-* workflow-key lifecycle in
create_job. Previously the revoke handlers wrapped mgr.start(...) but
only caught RuntimeError and ValueError, and get_job_manager() sat
outside the try block entirely. Any other exception type (TypeError
from a mismatched kwarg, OSError from the queue write, etc.) would
bubble up to FastAPI and leave the minted key live until its 24h TTL.

Fix: one try block covers both get_job_manager() and mgr.start(), with
a trailing except Exception that revokes and re-raises. The
RuntimeError -> 409 and ValueError -> 400 paths are unchanged so
specific client-facing status codes still surface. Revocation is still
best-effort (_revoke_internal_api_key_safe swallows errors) because we
never want revoke failures to mask the original crash.

Severity is low -- the key can't bootstrap longer access and the 24h
TTL bounds the window -- but the reviewer's point stands: eager revoke
on every failure path is the right invariant.

* Studio: nest response_format under extra_body so pydantic accepts it

The previous commit dropped response_format at the top level of a cloned
model_config's inference_parameters, which BuilderConfig rejected with:

  ValidationError: Extra inputs are not permitted [type=extra_forbidden]
  data_designer.model_configs.1.inference_parameters.response_format

data_designer's BaseInferenceParams is a pydantic model with extra=forbid
and only a fixed set of fields (temperature, top_p, max_tokens,
max_parallel_requests, timeout, extra_body). The pass-through path for
anything the schema doesn't know about is `extra_body`, which the
OpenAI SDK spreads into the chat-completions request body at the top
level -- which is exactly where llama-server reads response_format from.

Inject under extra_body (merging with any existing extra_body contents)
so the clone validates. llama-server still receives
{"type": "json_schema", "schema": <output_format>} at the top level of
the request body, which is the flat shape llama.cpp's server expects.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: forward response_format to llama-server and fence-wrap the reply

Two-part fix for the llm-structured data-recipe path:

(1) The /v1/chat/completions proxy was dropping response_format. The
route's passthrough branch only triggered on tools / tool messages, so
requests carrying a JSON schema fell into the non-passthrough GGUF path
which calls generate_chat_completion (no response_format kwarg). The
schema never reached llama-server, so guided decoding was a no-op and
the model emitted free-form text that happened to parse a fraction of
the time. Widen the passthrough trigger and teach _build_passthrough_payload
to forward response_format so llama-server's GBNF grammar actually runs.

Guided decoding does not require supports_tools, so split the condition:
a request is now passthrough-routed if it carries tools/tool messages
(existing behavior) OR carries response_format (new). The vision guard,
streaming fork, and tools-choice defaulting are unchanged.

(2) data_designer's llm-structured parser looks for a ```json ... ```
markdown fence and discards anything else. Guided decoding emits only
the JSON object (the GBNF grammar has no fence tokens), so a
100%-valid schema-constrained run still ended up 0 ok / N failed with
"No parsable JSON structure within ```json markdown fence". In
_openai_passthrough_non_streaming, wrap each choice's content in the
expected fence when the caller asked for guided decoding. Already-fenced
content is left alone so other clients that prefer raw JSON are not
affected; the wrap is scoped to requests that carried response_format.

Net effect on the GitHub Support Bot recipe on a local GGUF: schema
actually binds during sampling, content arrives wrapped in the fence
data_designer expects, and generation terminates immediately after the
closing brace instead of running out to max_tokens.

* Studio: Easy mode runs a full run, capped at the user's row count

Easy mode used to call runPreview, which produces a test run: no
artifact persisted, reduced progress tracking, and framed in the Runs
pane as "Test run". The whole point of the form is to let a user kick
off a real dataset build with one click, so wire it to runFull instead
and bind the Rows input to fullRows (not previewRows).

runFull requires a non-empty fullRunName. The Easy form has no run-name
input, so seed a default on mount whenever Easy is active and
fullRunName is still empty. Uses `<recipe name> <iso-timestamp>` so
each Easy run gets a stable-ish default that still sorts chronologically
in the Runs pane. User can override it from the Advanced run dialog
before clicking Run.

Rename GithubScraperEasyView's rows props from previewRows/setPreviewRows
to rows/setRows so the view stays agnostic to which hook state the page
chooses to bind. Loading indicator now follows fullLoading.

* Studio: clamp GitHub scrape page size and memoize the materialization

Two wins for the "before Generating fires" gap on small previews:

(1) scrape_{issues,prs,commits} hardcoded per_page (50 / 25 / 100) and
only checked the trial limit AFTER the page was written, so a 1-row
Easy run still asked GitHub for a full 50-issue + 25-PR page, wrote
them all to JSONL, and then stopped because total_new already exceeded
the trial cap. Cap per_page at min(page_cap, trial_limit) so
github_limit=1 actually asks for first:1.

(2) GitHubRepoSeedReader.get_dataset_uri used to scrape fresh on every
invocation. data_designer calls the seed reader multiple times per
recipe job (validation, preview, per-column sampling), so a 2-repo
Easy preview ran the full GraphQL scrape three times back-to-back,
burning ~15s of dead air before any LLM generation began.

Added a module-level in-process cache keyed on
(repos, item_types, limit, include_comments, max_comments_per_item,
sha256(token)[:16]) that stores the JSONL path of the first
materialization. Subsequent calls with the same signature return the
cached path, guarded by a staleness check that drops the entry if the
file was tmp-cleaned. Raw token values never land in the key.

Net effect on a 1-row Easy run, 2 repos, limit=1: 2 GraphQL round
trips instead of ~12, and the first-to-Generating gap collapses from
~15s to roughly 2-3s.

* Studio: make Easy mode Rows input editable instead of snapping to 1

The Rows to generate input used type="number" with value bound directly
to the rows state and an onChange that coerced any non-positive parse
result back to 1. The moment the user pressed backspace to clear the
field, the parent re-rendered with value=1 and the caret jumped, making
it impossible to change the value without arrowing the browser's +/-
spinner.

Switch to a text input with inputMode="numeric" and pattern="[0-9]*"
(so mobile still shows a numeric keyboard, and the browser drops the
spinner buttons the user did not want). Add a local rowsText buffer so
the field can hold transient empty / partial digit strings while
editing without fighting the parent state; the canonical rows value
only advances when the buffer parses to a valid integer in [1, 10000],
and onBlur clamps back to 1 or 10000 if the user left it out of range.

No behavior change for valid numeric edits - the downstream runFull()
still sees a clean positive integer.

* Studio: expand dataset cells horizontally by column on click

Click a long cell to expand that whole column. Click again to collapse.
Replaces the prior row-level vertical expansion which made it hard to
compare cells across columns. State is scoped per execution and per
column; the row itself is no longer a click target.

* Studio: force expanded dataset column to grow wide enough to read

* Studio: disable thinking for local recipe inference and plumb the kwarg

Reasoning-capable models (gemma-3n, qwen3.5, etc.) emit a
<think>...</think> preamble ahead of the answer by default, which
roughly doubles the generated token count per row on a local GGUF
and pushes the actual answer past data_designer's json-fence regex
on llm-structured columns. Recipes want the terse answer, not the
scratchpad.

Two halves of the fix:

(1) routes/data_recipe/jobs.py: when _inject_local_providers walks
the recipe's model_configs to point them at the local endpoint, also
stash chat_template_kwargs={"enable_thinking": false} under each
config's inference_parameters.extra_body. OpenAI SDK spreads
extra_body into the top-level request body, so llama-server and the
Studio /v1/chat/completions route both see it.

(2) routes/inference.py: the chat-completions route previously
dropped chat_template_kwargs on the floor because the whitelist
body builder only forwarded known fields.

    - At the top of openai_chat_completions, lift
      chat_template_kwargs.enable_thinking from payload.model_extra
      onto the typed payload.enable_thinking field when the caller
      did not set the latter, so the non-passthrough GGUF path's
      generate_chat_completion(...) call honors the override.
    - Teach _build_passthrough_payload to forward a
      chat_template_kwargs dict, and have _build_openai_passthrough_body
      derive that dict from payload.enable_thinking so
      response_format requests (structured columns) also land at
      llama-server with the reasoning preamble suppressed.

Net effect on a 10-row support-bot run with gemma-4-E2B-it-GGUF:
responses arrive without <think> tags, wall-clock per call drops
roughly in half, and structured columns stop leaking reasoning
tokens through the GBNF-constrained output.

* Studio: update GitHub Support Bot learning recipe with maintainer layout

Replace the template with the hand-laid-out export from the maintainer
so note nodes ship with real x/y positions (scattered around the
graph instead of all stacked at x=480) and the edges / canvas pan look
correct on first load. Also picks up the maintainer's prompt tweaks and
output schema names (coauthor_response / user_request / followups / task /
cites / confidence).

Diff is mostly ui.nodes positions and prompt bodies; runtime shape is
unchanged (seed_config / columns still target model_1 against the Local
Model provider).

* Studio: auto-size dataset sample columns; wide text gets a wide column

Drop the per-column click-to-expand toggle and the 180-char truncation.
Every column now renders its full value. Columns with long text get a
min-w of 48rem so the text is readable without wrapping into a tall
block; narrow-content columns get a 12rem min-w. The table wrapper
already has overflow-x-auto, so wide-column totals cause a horizontal
scrollbar instead of cramming everything into the viewport.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix GitHub scrape progress

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* add resetApiBase export for test setup

* Studio: rename github-support-bot output columns to User / Assistant

Previously emitted user_request and coauthor_response, which did not
match the canonical User / Assistant chat-pair shape that downstream
SFT consumers expect. Renamed the columns in the recipe JSON (columns,
UI node ids, edges, notes, prompt Jinja refs) and the matching copy in
the learning-recipes index, data-recipes-page, and easy view.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-04-24 12:02:03 -07:00
Wasim Yousef Said
a5eb2e3d50
Add tauri (#5144)
* add unsloth studio desktop app

* Fix review findings

- studio/src-tauri/tauri.conf.json: retarget updater to staging repo
  (danielhanchen/unsloth-staging-2); switch to unslothai/unsloth on upstream merge.
- studio/src-tauri/linux/postremove.sh: drop the interactive read loop and the
  /home/* iteration. Package maintainer scripts must stay non-interactive and
  must not touch other users' data.
- studio/frontend/src/app/auth-guards.ts: honor tauriAutoAuth() boolean. Failed
  auto-auth now redirects to /login; requireGuest/requirePasswordChangeFlow
  only redirect to /chat when auth succeeds. The new early-return on failed
  auth is intentional so the login / change-password flows remain reachable
  when desktop auth is not yet established.
- studio/frontend/src/config/env.ts: keep fetched=false on health failure so
  later calls retry instead of caching the client-side platform guess.
- studio/src-tauri/src/install.rs: pick the available system package manager
  (apt-get, dnf, zypper, pacman); AppImage bundles run on non-Debian distros.
- studio/frontend/src/lib/open-link.ts + markdown-text/sources callers: return
  boolean from openLink so callers only preventDefault on handled URLs; relative
  hrefs now navigate natively.
- studio/frontend/src/features/settings/tabs/about-tab.tsx: fetch(apiUrl(...))
  so the version request targets the backend port in desktop mode. The bare
  /api/health predates the Tauri webview (blame: the earlier onboarding commit,
  which ran with same-origin frontend/backend); in desktop mode the webview
  origin is tauri://localhost so the bare path fails.
- install.ps1: gate the install_python_stack.py hotfix on a sentinel comment
  instead of a content regex; append the sentinel after applying so reruns
  are unambiguous.
- unsloth_cli/commands/studio.py _write_auth_secret: use the atomic mkstemp +
  os.replace path on Windows too; chmod calls are wrapped in try/except OSError.
- studio/src-tauri/src/preflight.rs probe_existing_backends: fan out the health
  probes concurrently; desktop-auth status still runs sequentially per candidate.
  reqwest::Client is internally Arc-wrapped so the in-loop .clone() is a
  refcount bump, not a deep clone; annotated inline.
- studio/src-tauri/src/preflight.rs run_cli_probe: wait() after kill() to reap
  the child, matching probe_cli_capability.
- studio/src-tauri/src/process.rs + main.rs: add stop_backend_detached and use
  it from the tray quit handler so the 5s graceful-wait does not block the
  Tauri main loop. RunEvent::Exit keeps the synchronous safety-net call.
- studio/backend/main.py: drop the permissive localhost CORS regex in
  api-only mode; the explicit allow_origins list is sufficient.
- .github/workflows/release-desktop.yml: drop max-parallel: 1 so platform
  builds run in parallel, and lift releaseBody to an env var so the three
  tauri-action invocations share one source of truth.

* Fix review findings (loop 2)

- studio/backend/auth/storage.py update_password: clear_desktop_secret()
  alongside clear_bootstrap_password() so rotating the admin password
  also revokes any previously provisioned .desktop_secret. Without this,
  an old local desktop credential keeps minting fresh admin tokens via
  /api/auth/desktop-login after a password rotation.
- studio/src-tauri/src/desktop_auth.rs provision_desktop_auth: wrap
  cmd.output().await in tokio::time::timeout(30s). DESKTOP_AUTH_LOCK is
  held across the whole desktop_auth flow, and previously a hanging
  `unsloth studio provision-desktop-auth` subprocess would pin the lock
  indefinitely and freeze every subsequent desktop_auth call.

* Add review tests

* Consolidate review tests

Merge review-added tests into the existing studio/backend/tests/test_desktop_auth.py
(the PR's authoritative desktop-auth test file). Drops three scaffolding files under
tests/python/ in favor of five focused tests next to the tests they extend:
- test_update_password_clears_desktop_secret (runtime)
- test_update_password_on_unknown_user_leaves_desktop_secret_intact (runtime)
- test_cli_provisioning_delegates_to_storage_create_desktop_secret (source-level)
- test_cli_connect_auth_db_reads_storage_db_path (source-level)
- test_desktop_auth_provision_has_bounded_timeout (Rust source-level)

* Revert auth-guards.ts Tauri branches to unconditional form

The review loop on PR 5144 introduced a regression: the isTauri branch of
requireAuth redirected to /login when tauriAutoAuth() returned false, and
requireGuest / requirePasswordChangeFlow silently fell through on the same
condition. The Tauri desktop app authenticates via a local auto-generated
secret; it must never surface /login or /change-password to the user. A
failed auto-auth should let the startup layer retry, not expose a password
form.

Restore the three Tauri branches to the author's original unconditional
form (requireAuth: return; requireGuest / requirePasswordChangeFlow: throw
redirect({to: '/chat'})). Keep the rest of the review fixes -- the
apiUrl() fetch wrapping, authRedirect helper, and fetchAuthStatus refactor
are all legitimate improvements and are preserved.

* Revert release-desktop.yml to author's version

The review loop's workflow-file tweaks (drop max-parallel: 1, lift releaseBody
to an env var) are cosmetic. OAuth tokens cannot push workflow-file changes,
and fine-grained PATs cannot honor maintainerCanModify on a third-party fork.
Reverting the workflow file to wasimysaid's version lets the push go through
without needing a classic PAT with both repo and workflow scopes.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Daniel Han <unslothai@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-23 04:50:10 -07:00
Daniel Han
41a6cc8692
Studio: split vision-cache exception test to match transient vs permanent (#5145)
`TestVisionCacheOnException::test_exception_result_cached` currently
patches `load_model_config` with `side_effect=OSError("network down")`
and asserts `assert_called_once()`. That assertion is impossible by
design: `_is_vision_model_uncached` in
`studio/backend/utils/models/model_config.py` intentionally returns
`None` for `OSError` so `is_vision_model` does not cache the fallback
and retries on the next call. The module docstring on
`_vision_detection_cache` itself spells this out:

    Only definitive results (True/False from successful detection) are
    cached; transient failures (network errors, timeouts) are NOT
    cached so they can be retried.

The test has been failing identically on every downstream review run
against `unslothai/unsloth` main (e.g. `unsloth#5115`, `unsloth#5080`),
but the failure is not introduced by any of those PRs and does not
gate correctness.

Fix the collision by splitting the class into the two contracts the
code actually implements:

1. `test_permanent_exception_result_cached` keeps the original
   intent ("exception falls back to False and that False is cached")
   but uses `ValueError`, which is one of the exception types
   `_is_vision_model_uncached` treats as permanent and caches. No
   `huggingface_hub` import needed.

2. `test_transient_exception_not_cached` pins the opposite contract
   with the original `OSError("network down")`: the call returns
   False but the second invocation re-runs detection
   (`call_count == 2`). This guards against a future regression
   where somebody caches transient failures and then users with a
   flaky network permanently see wrong detection for a model.

Both tests use `assert ... is False` on the public API and mock-count
assertions on `load_model_config`; no private helpers are touched.
2026-04-23 00:22:40 -07:00
Roland Tannous
92cee0ff3e
Studio: support images on /v1/messages (Anthropic-compat) (#5128)
* Studio: support images on /v1/messages (Anthropic-compat)

Translate Anthropic `image` content blocks (base64 and url sources) to
OpenAI `image_url` multimodal parts so the Anthropic endpoint reaches
llama-server's native vision path. Mirrors the `/v1/chat/completions`
vision behavior: 400 when the active GGUF isn't a vision model, and
embedded images are re-encoded to PNG (stb_image format coverage).
Server-side agentic loop is disabled when images are present, matching
the existing `not image_b64` gate on /v1/chat/completions.

Adds translator + normalizer unit tests.

* Studio: address gemini-code-assist review on /v1/messages image support

- Preserve interleaving of Anthropic text + image content blocks in the
  translator (previously flattened all text first, then all images).
- Let _normalize_anthropic_openai_images return has_image so the route
  skips the second scan it was doing.
- Use module-level base64/io in the helper instead of re-importing.
2026-04-22 03:25:07 +04:00
Roland Tannous
21e9a91a57
Studio: forward standard OpenAI tools / tool_choice on /v1/responses (Codex compat) (#5122)
* Studio: forward standard OpenAI tools / tool_choice on /v1/responses

Mirrors the /v1/chat/completions client-side tool pass-through from #5099
so clients (OpenAI Codex CLI, OpenAI Python SDK, ...) that target the
Responses API receive structured function_call output items instead of
plain text with tool-call tokens leaking into content.

- ResponsesRequest: type tools/tool_choice properly, add parallel_tool_calls;
  accept function_call and function_call_output input items for multi-turn
- Translate flat Responses tool / tool_choice shape to the nested Chat
  Completions shape before forwarding to llama-server
- _normalise_responses_input: map function_call_output -> role="tool",
  function_call -> assistant tool_calls (preserving call_id)
- Non-streaming: map returned tool_calls -> top-level function_call
  output items keyed by call_id
- Streaming: emit response.output_item.added (function_call),
  response.function_call_arguments.delta/.done, and response.output_item.done
  per tool call while keeping the text message at output_index 0
- Pytest coverage: tools/tool_choice translation, multi-turn input mapping,
  non-streaming tool_calls mapping, response round-trip

* Studio: merge system messages and close inner stream on /v1/responses

Fixes two issues surfacing when OpenAI Codex CLI drives /v1/responses
against a GGUF with a strict chat template (gpt-oss harmony, Qwen3, ...).

1. "System message must be at the beginning" upstream errors
   Codex sends `instructions` AND a `role:"developer"` message in `input`,
   producing two separate system-role messages. Strict templates raise
   when a second system message exists or when one appears after a user
   turn. _normalise_responses_input now hoists all instructions / system /
   developer content into a single merged system message at the top of
   the Chat Completions message list.

2. "async generator ignored GeneratorExit" / "Attempted to exit cancel
   scope in a different task"
   _responses_stream consumed the inner chat-completions body_iterator
   without an explicit aclose() in a finally block. On client disconnect
   (Codex frequently cancels mid-stream), Python 3.13 finalized the inner
   async generator on a different task, tripping anyio's cancel-scope
   check. Mirrored the same try/finally + aclose pattern used by the
   /v1/messages, /v1/chat/completions, and /v1/completions passthroughs.

Tests: hoisting of instructions + developer, developer mid-conversation,
multiple system messages in input, no-system passthrough.

* Studio: accept Codex multi-turn shapes and fix cross-task stream close on /v1/responses

Two issues observed driving /v1/responses from OpenAI Codex CLI against a
GGUF backend.

1. 422 on every turn after the first
   Codex replays prior assistant turns with
   `content:[{"type":"output_text","text":...,"annotations":[],"logprobs":[]}]`
   and carries forward `reasoning` items (o-series / gpt-5) between turns.
   Our `ResponsesContentPart` union only accepted input_text / input_image,
   and `ResponsesInputItem` only message / function_call / function_call_output,
   so Pydantic failed the whole list and FastAPI returned
   `"Input should be a valid string"` against the `str` branch of the
   outer union.

   - Add `ResponsesOutputTextPart` for assistant-replay content.
   - Add `ResponsesUnknownContentPart` and `ResponsesUnknownInputItem`
     as permissive catch-alls (drop during normalisation).
   - Wire an explicit `Discriminator` so dispatch is deterministic and
     the fallthrough reaches the catch-all instead of misreporting via
     the outer `Union[str, list[...]]`.
   - `_normalise_responses_input` now accepts output_text parts, flattens
     single-part assistant text to a plain string (keeps legacy chat
     templates happy), and silently drops reasoning / unknown items.

2. "async generator ignored GeneratorExit" / cross-task cancel scope
   `_responses_stream` awaited `openai_chat_completions` in the parent
   route-handler task, which opens the httpx client for the inner
   passthrough on *that* task. The outer `StreamingResponse` then iterates
   in a child task, so the asyncgen GC finalises the inner httpcore byte
   stream on the child task, tripping anyio's "Attempted to exit cancel
   scope in a different task". Move the `await` inside `event_generator`
   so the httpx lifecycle stays within the single streaming child task,
   and surface any HTTPException as a `response.failed` SSE frame.

Tests: assistant output_text replay, reasoning-item tolerance, unknown
content-part tolerance, end-to-end Codex-shape payload (developer + user +
reasoning + function_call + function_call_output + assistant output_text +
user), and single-part assistant flattening to plain string.

* Studio: call llama-server directly from streaming /v1/responses

The previous fix (running the inner await inside event_generator) was not
enough. Wrapping the existing `openai_chat_completions` pass-through still
stacks two async generators: when the outer generator is closed, the
innermost `HTTP11ConnectionByteStream.__aiter__` in httpcore doesn't
receive GeneratorExit before Python's asyncgen GC finalises it in a
sibling task, tripping "Attempted to exit cancel scope in a different
task" and "async generator ignored GeneratorExit" — the same Python 3.13
+ httpcore 1.0.x interaction already seen in PRs #4956, #4981, #5099.

Cure both pass-throughs had: a single same-task httpx lifecycle with
explicit `aiter_lines().aclose()` BEFORE `resp.aclose()` / `client.aclose()`
in the generator's finally block.

Apply it at the Responses layer by dropping the wrapper entirely for GGUF:
open httpx, consume `resp.aiter_lines()`, parse `chat.completion.chunk`,
emit Responses SSE events, close everything in finally — all in the
single StreamingResponse child task. Non-GGUF streaming is rejected with
a 400 (wrapping the transformers backend would re-introduce the
double-layer pattern and isn't a Codex-compatible path today anyway).

Also surfaces upstream httpx.RequestError / non-200 as a
`response.failed` SSE frame rather than a dropped stream now that the
request is dispatched after SSE headers have gone out.

* Studio: silence benign httpcore asyncgen GC warnings on Python 3.13

The streaming pass-throughs (/v1/chat/completions, /v1/messages,
/v1/responses, /v1/completions) all use the proven #4981 / #5099 pattern
— single-task httpx lifecycle with explicit aiter_lines().aclose() ahead
of resp.aclose() / client.aclose() in the generator's finally block.
That handles our own iterators correctly.

The residual noise ("async generator ignored GeneratorExit" /
"Attempted to exit cancel scope in a different task") comes from an
innermost HTTP11ConnectionByteStream.__aiter__ that httpcore creates
internally inside its pool. We hold no reference to it, so we cannot
aclose it ourselves. Python 3.13's asyncgen GC hook finalises it on the
finaliser task, its aclose path enters an anyio CancelScope shield, and
Python flags the cross-task exit. The response has already been
delivered with a 200 by then — it is purely log noise, not a functional
failure. Same interaction seen in modelcontextprotocol/python-sdk #831,
agno #3556, chainlit #2361, langchain-mcp-adapters #254.

Install a targeted sys.unraisablehook that swallows this specific tuple
— RuntimeError mentioning "cancel scope" or "GeneratorExit" plus an
object repr referencing HTTP11ConnectionByteStream — and defers to the
default hook for every other unraisable. Idempotent; guarded by a
sentinel attribute so repeated imports don't stack filters.
2026-04-21 13:17:20 +04:00
Roland Tannous
ac2daf8b7a
Studio: forward standard OpenAI tools / tool_choice to llama-server (#5099)
* fix(studio): forward OpenAI tools/tool_choice to llama-server (#4999)

Studio's /v1/chat/completions silently stripped standard OpenAI `tools`
and `tool_choice` fields, so clients using standard function calling
(opencode, Claude Code, Cursor, Continue, ...) never got structured
tool_calls back. Adds a client-side pass-through path mirroring the
existing Anthropic /v1/messages flow: when `tools` is present without
Studio's `enable_tools` shorthand, the request is forwarded to
llama-server verbatim so the client sees native id, finish_reason
("tool_calls"), delta.tool_calls, and accurate usage tokens.

Also wires Anthropic tool_choice forwarding: /v1/messages previously
accepted tool_choice on the request model but silently dropped it with
a warning. Translate the four Anthropic shapes to OpenAI format and
forward them so agentic clients can actually enforce tool use.

- ChatCompletionRequest: add tools, tool_choice, stop; extra="allow"
- ChatMessage: accept role="tool", optional tool_call_id / tool_calls /
  name; content is now optional (assistant with only tool_calls)
- routes/inference.py: _openai_passthrough_stream /
  _openai_passthrough_non_streaming helpers, routing branch in
  openai_chat_completions, vision+tools via content-parts injection
- _build_passthrough_payload: tool_choice parameter (default "auto")
- anthropic_compat: anthropic_tool_choice_to_openai() translator
- tests/test_openai_tool_passthrough.py: Pydantic + translator unit tests
- tests/test_studio_api.py: 5 new E2E tests (non-stream, stream,
  multi-turn, OpenAI SDK, Anthropic tool_choice=any regression)

* fix(studio): surface httpx transport errors from OpenAI passthrough

When the managed llama-server subprocess crashes mid-request, the
async pass-through helpers in routes/inference.py used to return a
bare 500 (non-streaming) or an "An internal error occurred" SSE chunk
(streaming) because _friendly_error only recognized the sync path's
"Lost connection to llama-server" substring -- httpx transport
failures (ConnectError / ReadError / RemoteProtocolError /
ReadTimeout) stringify differently and fell through to the generic
case.

- _friendly_error: map any httpx.RequestError subclass to the same
  "Lost connection to the model server" message the sync chat path
  emits. Placed before the substring heuristics so the streaming path
  automatically picks it up via its existing except Exception catch.
- _openai_passthrough_non_streaming: wrap the httpx.AsyncClient.post
  in a try/except httpx.RequestError and re-raise as HTTPException
  502 with the friendly detail.
- tests/test_openai_tool_passthrough.py: new TestFriendlyErrorHttpx
  class pinning the mapping for ConnectError, ReadError,
  RemoteProtocolError, ReadTimeout, and confirming non-httpx paths
  (context-size heuristic, generic fallback) are unchanged.

* fix(studio): close aiter_bytes/aiter_lines explicitly in passthroughs

The httpcore asyncgen cleanup fix in 5cedd9a5 is incomplete on Python
3.13 + httpcore 1.0.x: it switched to manual client/response lifecycle
but still used anonymous `async for raw_line in resp.aiter_lines():`
patterns in all three streaming paths. Python's async for does NOT
auto-close the iterator on break/return, so the aiter_lines /
aiter_bytes async generator remains alive, reachable only from the
surrounding coroutine frame. Once `_stream()` returns the frame is
GC'd and the orphaned asyncgen is finalized on a LATER GC pass in a
DIFFERENT asyncio task, where httpcore's
HTTP11ConnectionByteStream.aclose() enters anyio.CancelScope.__exit__
with a mismatched task and prints "Exception ignored in: <async
generator>" / "async generator ignored GeneratorExit" / "Attempted
to exit cancel scope in a different task" to the server log.

User observed this on /v1/messages after successful (status 200)
requests, with the traceback pointing at HTTP11ConnectionByteStream
.__aiter__ / .aclose inside httpcore.

Fix: save resp.aiter_lines() / resp.aiter_bytes() as a variable and
explicitly `await iter.aclose()` in the finally block BEFORE
resp.aclose() / client.aclose(). This closes the asyncgen inside the
current task's event loop, so the internal httpcore byte stream is
cleaned up before Python's asyncgen GC hook has anything orphaned to
finalize. Each aclose is wrapped in try/except Exception so nested
anyio cleanup noise can't bubble out.

Applied to all three streaming passthrough paths:
- _anthropic_passthrough_stream (/v1/messages client-side tool path)
- _openai_passthrough_stream (/v1/chat/completions client-side tool
  path, new in this PR)
- openai_completions (/v1/completions bytes proxy from PR #4956)

* fix(studio): default ChatCompletionRequest.stream to false per OpenAI spec

OpenAI's /v1/chat/completions spec defaults `stream` to false, so
clients that omit the field (naive curl, minimal integrations) expect
a single JSON response back. Studio was defaulting to true, silently
switching those clients into SSE and breaking any parser that didn't
also handle streaming. ResponsesRequest and AnthropicMessagesRequest
already default to false correctly; only ChatCompletionRequest was
wrong.

Studio's own frontend always sets `stream` explicitly on every
chat-adapter / chat-api / runtime-provider call site, so the flip has
no UI impact. SDK users (OpenAI Python/JS SDK, opencode, Claude Code,
Cursor, Continue) also always pass `stream` explicitly, so they're
unaffected. The only clients feeling the change are raw-curl users
who were relying on the wrong default -- those get the correct OpenAI
behavior now.

Added a regression test pinning the default so it can't silently
flip back.

* fix(studio): reject images in OpenAI tool passthrough for text-only GGUFs

The new tool passthrough branch runs before _extract_content_parts,
skipping the existing not is_vision guard. Requests combining tools
with an image on a text-only tool-capable GGUF were forwarded to
llama-server, producing opaque upstream errors instead of the
pre-existing clear 400. Restore the guard inline at the dispatch
point, checking both legacy image_base64 and inline image_url parts.

* fix(studio): require tool_call_id on role=tool chat messages

Enforce the OpenAI spec rule that role="tool" messages must carry a
tool_call_id. Without it, upstream backends cannot associate a tool
result with the assistant's prior tool_calls entry and the request
fails in non-obvious ways through the passthrough path. Reject at the
request boundary with a 422 instead.

* fix(studio): harden OpenAI tool passthrough validation and error surfacing

Three related fixes called out by the PR review:

1. Preserve upstream status codes in the streaming passthrough. The
   httpx request is now dispatched before the StreamingResponse is
   constructed. Non-200 upstream responses and httpx RequestError
   transport failures raise HTTPException with the real status
   instead of being buried inside a 200 SSE error frame, so OpenAI
   SDK clients see APIError/BadRequestError/... as expected.

2. Require non-empty content on user/system/tool messages. Per the
   OpenAI spec, content may only be omitted on assistant messages
   that carry tool_calls; enforce that at the request boundary so
   malformed messages never reach the passthrough path.

3. Role-constrain tool-call metadata. tool_calls is only valid on
   role=assistant, tool_call_id and name only on role=tool. Without
   this, a user/system message with tool_calls would flip the
   passthrough branch on and be forwarded to llama-server, surfacing
   as an opaque upstream error.

* fix(studio): normalize image mode and passthrough JSON verbatim

Two Gemini-code-assist review findings on PR #5099:

1. Unconditionally convert decoded images to RGB before PNG encoding.
   The prior code only handled RGBA, letting CMYK/I/F images crash
   at img.save(format="PNG") and surface as opaque 400s. Applied to
   both the passthrough helper and the non-passthrough GGUF path
   that originally carried this pattern, keeping the two sites in
   sync.

2. Return the upstream JSON body as raw bytes via Response rather
   than parse-then-re-serialize with JSONResponse. Matches the
   passthrough helper's "verbatim" contract and drops a redundant
   round-trip.

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-04-18 12:53:23 +04:00
Daniel Han
f0d03655e8
Studio: add folder browser modal for Custom Folders (#5035)
* Studio: add folder browser modal for Custom Folders

The Custom Folders row in the model picker currently only accepts a
typed path. On a remote-served Studio (Colab, shared workstation) that
means the user has to guess or paste the exact server-side absolute
path. A native browser folder picker can't solve this: HTML
`<input type="file" webkitdirectory>` hides the absolute path for
security, and the File System Access API (Chrome/Edge only) returns
handles rather than strings, neither of which the server can act on.

This PR adds a small in-app directory browser that lists paths on the
server and hands the chosen string back to the existing
`POST /api/models/scan-folders` flow.

## Backend

* New endpoint `GET /api/models/browse-folders`:
  * `path` query param (expands `~`, accepts relative or absolute; empty
    defaults to the user's home directory).
  * `show_hidden` boolean to include dotfiles/dotdirs.
  * Returns `{current, parent, entries[], suggestions[]}`. `parent` is
    null at the filesystem root.
  * Immediate subdirectories only (no recursion); files are never
    returned.
  * `entries[].has_models` is a cheap hint: the directory looks like it
    holds models if it is named `models--*` (HF hub cache layout) or
    one of the first 64 children is a .gguf/.safetensors/config.json/
    adapter_config.json or another `models--*` subfolder.
  * Sort order: model-bearing dirs, then plain, then hidden; case-
    insensitive alphabetical within each bucket.
  * Suggestions auto-populate from HOME, the HF cache root, and any
    already-registered scan folders, deduplicated.
  * Error surface: 404 for missing path, 400 for non-directory, 403 on
    permission errors. Auth-required like the other models routes.

* New Pydantic schemas `BrowseEntry` and `BrowseFoldersResponse` in
  `studio/backend/models/models.py`.

## Frontend

* New `FolderBrowser` component
  (`studio/frontend/src/components/assistant-ui/model-selector/folder-browser.tsx`)
  using the existing `Dialog` primitive. Features:
  * Clickable breadcrumb with a `..` row for parent navigation.
  * Quick-pick chips for the server-provided suggestions.
  * `Show hidden` checkbox.
  * In-flight fetch cancellation via AbortController so rapid
    navigation doesn't flash stale results.
  * Badges model-bearing directories inline.

* `chat-api.ts` gains `browseFolders(path?, showHidden?)` and matching
  types.

* `pickers.tsx` adds a folder-magnifier icon next to the existing `Add`
  button. Opening the browser seeds it with whatever the user has
  already typed; confirming fills the text input, leaving the existing
  validation and save flow unchanged.

## What it does NOT change

* The existing text-input flow still works; the browser is additive.
* No new permissions or escalation; the endpoint reads only directories
  the server process is already allowed to read.
* No model scanning or filesystem mutation happens from the browser
  itself -- it just returns basenames for render.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: cap folder-browser entries and expose truncated flag

Pointing the folder browser at a huge directory (``/usr/lib``,
``/proc``, or a synthetic tree with thousands of subfolders) previously
walked the whole listing and stat-probed every child via
``_looks_like_model_dir``. That is both a DoS shape for the server
process and a large-payload surprise for the client.

Introduce a hard cap of 2000 subdirectory entries and a
``truncated: bool`` field on the response. The frontend renders a small
hint below the list when it fires, prompting the user to narrow the
path. Below-cap directories are unchanged.

Verified end-to-end against the live backend with a synthetic tree of
2050 directories: response lands at 2000 entries, ``truncated=true``,
listing finishes in sub-second time (versus tens of seconds if we were
stat-storming).

* Studio: suggest LM Studio / Ollama dirs + 2-level model probe

Three improvements to the folder-browser, driven by actually dropping
an LM Studio-style install (publisher/model/weights.gguf) into the
sandbox and walking the UX:

## 1. Quick-pick chips for other local-LLM tools

`well_known_model_dirs()` (new) returns paths commonly used by
adjacent tools. Only paths that exist are returned so the UI never
shows dead chips.

* LM Studio current + legacy roots + user-configured
  `downloadsFolder` from its `settings.json` (reuses the existing
  `lmstudio_model_dirs()` helper).
* Ollama: `$OLLAMA_MODELS` env override, then `~/.ollama/models`,
  `/usr/share/ollama/.ollama/models`, and `/var/lib/ollama/.ollama/models`
  (the systemd-service install path surfaced in the upstream "where is
  everything?" issue).
* Generic user-choice locations: `~/models`, `~/Models`.

Dedup is stable across all sources.

## 2. Two-level model-bearing probe

LM Studio and Ollama both use `root/publisher/model/weights.gguf`.
The previous `has_models` heuristic only probed one level, so the
publisher dir (whose immediate children are model dirs, not weight
files) was always marked as non-model-bearing. Pulled the direct-
signal logic into `_has_direct_model_signal` and added a grandchild
probe so the classic layout is now recognised.

Still O(PROBE^2) worst-case, still returns immediately for
`models--*` names (HF cache layout) and for any direct weight file.

## 3. model_files_here hint on response body

A leaf model dir (just GGUFs, no subdirs) previously rendered as
`(empty directory)` in the modal, confusing users into thinking the
folder wasn't scannable. Added a `model_files_here` count on the
response (capped at 200) and a small hint row in the modal: `N model
files in this folder. Click "Use this folder" to scan it.`

## Verification

Simulated an LM Studio install by downloading the real 84 MB
`unsloth/SmolLM2-135M-Instruct-Q2_K.gguf` into
`~/.lmstudio/models/unsloth/SmolLM2-135M-Instruct-GGUF/`. Confirmed
end-to-end:

* Home listing suggests `~/.lmstudio/models` as a chip.
* Browsing `~/.lmstudio/models` flags `unsloth` (publisher) as
  `has_models=true` via the 2-level probe.
* Browsing the publisher flags `SmolLM2-135M-Instruct-GGUF` (model
  dir) as `has_models=true`.
* Browsing the model dir returns empty entries but
  `model_files_here=1`, and the frontend renders a hint telling the
  user it is a valid target.

* Studio: one-click scan-folder add + prominent remove + plain search icon

Three small Custom Folders UX fixes after real-use walkthrough:

* **One-click add from the folder browser**. Confirming `Use this
  folder` now submits the path directly to
  `POST /api/models/scan-folders` instead of just populating the text
  input. `handleAddFolder` takes an optional explicit path so the
  submit lands in the same tick as `setFolderInput`, avoiding a
  state-flush race. The typed-path + `Add` button flow is unchanged.

* **Prominent remove X on scan folders**. The per-folder delete
  button was `text-muted-foreground/40` and hidden entirely on
  desktop until hovered (`md:opacity-0 md:group-hover:opacity-100`).
  Dropped the hover-only cloak, bumped color to `text-foreground/70`,
  added a red hover/focus background, and sized the icon up from
  `size-2.5` to `size-3`. Always visible on every viewport.

* **Plain search icon for the Browse button**. `FolderSearchIcon`
  replaced with `Search01Icon` so it reads as a simple "find a
  folder" action alongside the existing `Add01Icon`.

* Studio: align Custom Folders + and X buttons on the same right edge

The Custom Folders header used `px-2.5` with a `p-0.5` icon button,
while each folder row used `px-3` with a `p-1` button. That put the
X icon 4px further from the right edge than the +. Normalised both
rows to `px-2.5` with `p-1` so the two icons share a column.

* Studio: empty-state button opens the folder browser directly

The first-run empty state for Custom Folders was a text link reading
"+ Add a folder to scan for local models" whose click toggled the
text input. That's the wrong default: a user hitting the empty state
usually doesn't know what absolute path to type, which is exactly
what the folder browser is for.

* Reword to "Browse for a models folder" with a search-icon
  affordance so the label matches what the click does.
* Click opens the folder browser modal directly. The typed-path +
  Add button flow is still available via the + icon in the
  section header, so users who know their path keep that option.
* Slightly bump the muted foreground opacity (70 -> hover:foreground)
  so the button reads as a primary empty-state action rather than a
  throwaway hint.

* Studio: Custom Folders header gets a dedicated search + add button pair

The Custom Folders section header had a single toggle button that
flipped between + and X. That put the folder-browser entry point
behind the separate empty-state link. Cleaner layout: two buttons in
the header, search first, then add.

* Search icon (left) opens the folder browser modal directly.
* Plus icon (right) toggles the text-path input (unchanged).
* The first-run empty-state link is removed -- the two header icons
  cover both flows on every state.

Both buttons share the same padding / icon size so they line up with
each other and with the per-folder remove X.

* Studio: sandbox folder browser + bound caps + UX recoveries

PR review fixes for the Custom Folders folder browser. Closes the
high-severity CodeQL path-traversal alert and addresses the codex /
gemini P2 findings.

Backend (studio/backend/routes/models.py):

* New _build_browse_allowlist + _is_path_inside_allowlist sandbox.
  browse_folders now refuses any target that doesn't resolve under
  HOME, HF cache, Studio dirs, registered scan folders, or the
  well-known third-party model dirs. realpath() is used so symlink
  traversal cannot escape the sandbox. Also gates the parent crumb
  so the up-row hides instead of 403'ing.
* _BROWSE_ENTRY_CAP now bounds *visited* iterdir entries, not
  *appended* entries. Dirs full of files (or hidden subdirs when
  show_hidden is False) used to defeat the cap.
* _count_model_files gets the same visited-count fix.
* PermissionError no longer swallowed silently inside the
  enumeration / counter loops -- now logged at debug.

Frontend (folder-browser.tsx, pickers.tsx, chat-api.ts):

* splitBreadcrumb stops mangling literal backslashes inside POSIX
  filenames; only Windows-style absolute paths trigger separator
  normalization. The Windows drive crumb value is now C:/ (drive
  root) instead of C: (drive-relative CWD-on-C).
* browseFolders accepts and forwards an AbortSignal so cancelled
  navigations actually cancel the in-flight backend enumeration.
* On initial-path fetch error, FolderBrowser now falls back to HOME
  instead of leaving the modal as an empty dead end.
* When the auto-add path (one-click "Use this folder") fails, the
  failure now surfaces via toast in addition to the inline
  paragraph (which is hidden when the typed-input panel is closed).

* Studio: rebuild browse target from trusted root for CodeQL clean dataflow

CodeQL's py/path-injection rule kept flagging the post-validation
filesystem operations because the sandbox check lived inside a
helper function (_is_path_inside_allowlist) and CodeQL only does
intra-procedural taint tracking by default. The user-derived
``target`` was still flowing into ``target.exists`` /
``target.is_dir`` / ``target.iterdir``.

The fix: after resolving the user-supplied ``candidate_path``,
locate the matching trusted root from the allowlist and rebuild
``target`` by appending each individually-validated segment to
that trusted root. Each segment is rejected if it isn't a single
safe path component (no separators, no ``..``, no empty/dot).
The downstream filesystem ops now operate on a Path constructed
entirely from ``allowed_roots`` (trusted) plus those validated
segments, so CodeQL's dataflow no longer sees a tainted source.

Behavior is unchanged for all valid inputs -- only the
construction of ``target`` is restructured. Live + unit tests
all pass (58 selected, 7 deselected for Playwright env).

* Studio: walk browse paths from trusted roots for CodeQL

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Ubuntu <ubuntu@h100-8-cheapest.us-east5-a.c.unsloth.internal>
2026-04-15 08:04:33 -07:00
Lee Jackson
f9ef639dde
Studio: support GGUF variant selection for non-suffixed repos (#5023)
* fix: support GGUF variant selection for non-suffixed repos

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: harden GGUF detection across cached models and picker flows

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* chore: use shared GGUF picker helper for search rows

* fix: avoid mixed cache duplication and preserve GGUF fallback detection

* fix: unify GGUF cache matching and merge picker hints

* fix: normalize local GGUF matching across picker and model config

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: robust cached-gguf classification + hint-aware click routing

- _repo_gguf_size_bytes: treat size_on_disk=None as 0 and dedupe fallback
  by commit_hash so partial/interrupted downloads don't TypeError out of
  sum() and wipe the entire cached list.
- list_cached_gguf / list_cached_models: narrow per-repo try/except so
  one malformed repo no longer poisons the whole response.
- handleModelClick: route through isKnownGgufRepo instead of the
  suffix-only isGgufRepo, so non-suffixed GGUF repos still open the
  variant expander from every call site.
- Replace the modelIsGgufById/resultIsGgufById Maps with Sets of known
  GGUF ids to stop conflating "no hint" with "known not-GGUF".
- Make HfModelResult.isGguf required (it is always set in makeMapModel).
- Add regression tests for the None size case, mixed-repo inclusion in
  cached-gguf, and per-repo error isolation.

* fix: exclude mmproj from GGUF classification and case-normalize hint lookups

- _repo_gguf_size_bytes now filters mmproj vision-adapter files so
  safetensors+mmproj.gguf repos stay on the cached-models path and
  non-GGUF rows no longer show zero pickable variants. A vision-capable
  GGUF repo (main weight + mmproj adapter) still classifies as GGUF and
  reports the main weight size.
- modelGgufIds / resultGgufIds now key on lowercased ids and
  isKnownGgufRepo lowercases its lookup, so store and HF-search ids
  that differ only by casing still match the same GGUF hint.
- New regression tests: mmproj-only repo excluded from cached-gguf,
  same repo included in cached-models, vision-capable repo still
  classified as GGUF with correct size.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
2026-04-15 15:32:01 +04:00
Roland Tannous
13928b5f0e
Add configurable PyTorch mirror via UNSLOTH_PYTORCH_MIRROR env var (#5024)
* Add configurable PyTorch mirror via UNSLOTH_PYTORCH_MIRROR env var

When set, UNSLOTH_PYTORCH_MIRROR overrides the default
https://download.pytorch.org/whl base URL in all four install scripts
(install.sh, install.ps1, studio/setup.ps1, studio/install_python_stack.py).
When unset or empty, the official URL is used. This lets users behind
corporate proxies or in regions with poor connectivity to pytorch.org
point at a local mirror without patching scripts.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add pytest for UNSLOTH_PYTORCH_MIRROR in install_python_stack.py

Tests that _PYTORCH_WHL_BASE picks up the env var when set, falls back
to the official URL when unset or empty, and preserves the value as-is
(including trailing slashes).

* Remove stale test assertions for missing install.sh messages

* Fix GPU mocking in test_get_torch_index_url.sh

Extract _has_usable_nvidia_gpu and _has_amd_rocm_gpu alongside
get_torch_index_url so the GPU-presence checks work in tests.
Add -L flag handling to mock nvidia-smi so it passes the GPU listing
check. All 26 tests now pass on CPU-only machines.

* Strip trailing slash from UNSLOTH_PYTORCH_MIRROR to avoid double-slash URLs

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-15 11:39:11 +04:00
Daniel Han
5aa8c15246
Studio: hard-stop at n_ctx with a 'Context limit reached' toast (#5021)
* Studio: hard-stop at n_ctx with a dedicated 'Context limit reached' toast

llama-server's default behavior when the KV cache fills is to silently
drop the oldest non-``n_keep`` tokens and keep generating. The UI has
no way to tell the user that earlier turns were evicted -- they just
see degraded continuity and a confusing ``5,361 / 4,096`` on the
context usage bar.

Launch llama-server with ``--no-context-shift`` so it returns a clean
error once the request would exceed ``n_ctx``. In the chat adapter,
catch the error, identify it as a context-limit error via
``isContextLimitError()``, and surface a dedicated toast that names
the exact control to adjust: the ``Context Length`` field in the chat
Settings panel.

Also add a lightweight tooltip hint on ``ContextUsageBar`` when usage
crosses 85%, so users see the "raise Context Length in Settings"
suggestion before they hit the hard stop.

Tests:

  * ``test_llama_cpp_no_context_shift.py`` pins the ``--no-context-shift``
    flag in the static launch-command template, and pins it inside the
    unconditional ``cmd = [ ... ]`` block so a future refactor can't
    hide it behind a branch.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Shorten --no-context-shift comment to 1 line

* Match backend _friendly_error rewrite in isContextLimitError

Codex review on PR caught that ``backend/routes/inference.py::_friendly_error``
rewrites the raw llama-server text
  "request (X tokens) exceeds the available context size (Y tokens)"
into
  "Message too long: X tokens exceeds the Y-token context window. ..."
on the main streaming GGUF path. The heuristic only looked for
"context size" / "exceeds the available context" / "context shift",
none of which survive the rewrite, so the new "Context limit reached"
toast would never fire for the most common case. Add matches for
"message too long" and "context window" so both wordings hit.

Also addresses Gemini feedback on the launch-flag test:
  * Use ``inspect.getsource(LlamaCppBackend.load_model)`` instead of
    reading ``__file__`` directly; scopes the assertions to the
    function that actually launches llama-server.
  * Replace the hardcoded ``"            ]"`` indent search with a
    line-at-a-time scan for a line that is just ``]``, so the test
    survives reformatting.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-14 10:58:20 -07:00
Daniel Han
5861a7ce15
Studio: split model-load progress label across two rows (#5020)
* Studio: split model-load progress label across two rows

The chat flow and training overlay both compose a progress label like
"112.6 of 122.3 GB • 331.0 MB/s • 30s left" and render it next to the
percent badge in a single flex row. Once the rate + ETA part shows up,
the label outgrows the row width and wraps mid-phrase, orphaning the
percent ("19 left %") onto a second ragged line.

Fix in model-load-status.tsx: split the label on the first " • " into
a primary (size) chunk that stays on row 1 with the percent, and a
secondary (rate/ETA) chunk that renders on its own muted row below.
Labels without a bullet (e.g. "22.8 GB downloaded") collapse cleanly
to one row. The inline-status variant keeps only the primary and
surfaces the full label via the tooltip.

Also extracts the rate/ETA math out of useTransferStats into a pure
``transfer-stats.ts`` module (appendSample + computeTransferStats) so
it can be reasoned about and tested without React. The hook is now a
thin wrapper that feeds sample history through the pure functions.

Backend: adds two companion test files for load_progress():

  * test_llama_cpp_load_progress_matrix.py (21 tests) -- platform
    matrix (Linux /proc, macOS/Windows absence), VmRSS parsing
    variants (tab/space/missing/malformed), filesystem edges (HF-cache
    symlinks, broken symlinks, nonexistent paths, relative paths),
    shard aggregation (partial multi-shard, two series in same dir,
    mmproj-* exclusion, single-file), lifecycle races, concurrent
    sampling (10 threads x 50 iters against real /proc), fraction
    bounds.
  * test_llama_cpp_load_progress_live.py (5 tests) -- no-mock live
    integration: real subprocess allocating 100 MB to match VmRSS,
    real ready phase, real dead-pid degradation, real shard
    aggregation, repeated polling. Skipped on non-Linux.

Both complement the existing test_llama_cpp_load_progress.py.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Hoist splitProgressLabel out of JSX IIFE (review feedback)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-14 10:58:16 -07:00
Daniel Han
bb14ab144a
Studio: live model-load progress + rate/ETA on download and load (#5017)
* Studio: live model-load progress + rate/ETA on download and load

Two UX fixes for the opaque multi-minute wait between clicking Load
and being able to chat, visible most clearly on large MoE GGUFs like
MiniMax-M2.7 (131 GB of weights on a 97 GB GPU):

1. **Model-load phase is now observable.** The existing chat flow
   transitions the toast to "Starting model..." as soon as the
   download hits 100%, then shows a spinner with no other feedback
   until llama-server reports healthy. For a 130 GB model that spinner
   freezes for five-plus minutes while the kernel pages shards into
   the page cache. A new `GET /api/inference/load-progress` endpoint
   samples `/proc/<pid>/status VmRSS` on the llama-server subprocess
   against the sum of shard file sizes on disk, so the UI can render
   a real bar plus rate / ETA during that window.

2. **Rate and ETA on downloads and loads.** Both the chat toast and
   the training-start overlay used to show a static pair of numbers
   (for example "15.4 of 140.8 GB"). A rolling 15-second window over
   the existing byte-series now surfaces "85.3 MB/s, 24m 23s left"
   beside that pair. The estimator is shared between the download
   and load phases so the numbers don't reset when the phase flips.

Also fixes a pre-existing assignment bug uncovered while wiring this
up: `load_model` was storing the caller's `gguf_path` kwarg into
`self._gguf_path`, which is `None` on the HF-download code path. The
resolved on-disk path (`model_path`) is what llama-server actually
mmaps; downstream consumers need that. No existing reader used
`_gguf_path`, so this is a correctness fix for the new endpoint.

- Backend: `LlamaCppBackend.load_progress()`, `GET /api/inference/load-progress`, `LoadProgressResponse` Pydantic model.
- Frontend: `useTransferStats` hook, `formatRate` / `formatEta` helpers, `getLoadProgress` client, rewired chat toast and `DownloadRow` in the training overlay.
- Tests: `studio/backend/tests/test_llama_cpp_load_progress.py` covers empty states, mmap phase, ready phase, sharded total aggregation, missing gguf_path, and unreadable /proc (7 cases). `tsc -b` and `vite build` on the frontend both clean.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-14 09:46:22 -07:00
Daniel Han
7252410ccc
studio: stream export worker output into the export dialog (#4897)
* studio: stream export worker output into the export dialog

The Export Model dialog only showed a spinner on the "Exporting..."
button while the worker subprocess was doing the actual heavy lifting.
For Merged to 16bit and GGUF / Llama.cpp exports this meant several
minutes (or more, for large models) of opaque silence, with no way to
tell whether save_pretrained_merged, convert_hf_to_gguf.py, or
llama-quantize was making progress.

This adds a live terminal-style output panel inside the export dialog,
rendered just above the Cancel / Start Export buttons and scrollable
with auto-follow-tail. It shows stdout and stderr from both the worker
process itself and any child process it spawns (GGUF converter,
llama-quantize), coloured by stream.

Backend

- core/export/worker.py: new _setup_log_capture(resp_queue) installed
  before LogConfig.setup_logging. It saves the original stdout/stderr
  fds, creates pipes, os.dup2's the write ends onto fds 1 and 2 (so
  every child process inherits the redirected fds), and spins up two
  daemon reader threads. Each thread reads bytes from a pipe, echoes
  them back to the original fd (so the server console keeps working),
  splits on \n and \r, and forwards each line to the resp queue as
  {"type":"log","stream":"stdout|stderr","line":...,"ts":...}.
  PYTHONUNBUFFERED=1 is set so nested Python converters flush
  immediately.

- core/export/orchestrator.py:
  - Thread-safe ring buffer (collections.deque, maxlen 4000) with a
    monotonically increasing seq counter. clear_logs(),
    get_logs_since(cursor), get_current_log_seq(), is_export_active().
  - _wait_response handles rtype == "log" by appending to the buffer
    and continuing the wait loop. Status messages are also surfaced as
    a "status" stream so users see high level progress alongside raw
    subprocess output.
  - load_checkpoint, _run_export, and cleanup_memory now wrap their
    bodies with the existing self._lock (previously unused), clear the
    log buffer at the start of each op, and flip _export_active in a
    try/finally so the SSE endpoint can detect idle.

- routes/export.py:
  - Wrapped every sync orchestrator call (load_checkpoint,
    cleanup_memory, export_merged_model, export_base_model,
    export_gguf, export_lora_adapter) in asyncio.to_thread so the
    FastAPI event loop stays free during long exports. Without this
    the new SSE endpoint could not be served concurrently with the
    blocking export POST.
  - New GET /api/export/logs/stream SSE endpoint. Honors
    Last-Event-ID and a since query param for reconnect, emits log /
    heartbeat / complete / error events, uses the id field to carry
    the log seq so clients can resume cleanly. On first connect
    without an explicit cursor it starts from the current seq so old
    lines from a previous run are not replayed.

Frontend

- features/export/api/export-api.ts: streamExportLogs() helper that
  authFetches the SSE endpoint and parses id / event / data fields
  manually (same pattern as streamTrainingProgress in train-api.ts).

- features/export/components/export-dialog.tsx:
  - Local useExportLogs(exporting) hook that opens the SSE stream on
    exporting transitions to true, accumulates up to 4000 lines in
    component state, and aborts on cleanup.
  - New scrollable output panel rendered above DialogFooter, only
    shown for Merged to 16bit and GGUF / Llama.cpp (LoRA adapter is
    a fast disk write with nothing to show). Dark terminal styling
    (bg-black/85, emerald text, rose for stderr, sky for status),
    max-height 14rem, auto-scrolls to the bottom on new output but
    stops following if the user scrolls up. A small streaming / idle
    indicator is shown next to the panel title.
  - DialogContent widens from sm:max-w-lg to sm:max-w-2xl when the
    output panel is visible so the logs have room to breathe.

Verified

- Python smoke test (tests/smoke_export_log_capture.py): spawns a
  real mp.get_context("spawn") process, installs _setup_log_capture,
  confirms that parent stdout prints, parent stderr prints, AND a
  child subprocess invoked via subprocess.run (both its stdout and
  stderr) are all captured in the resp queue. Passes.
- Orchestrator log helpers tested in isolation: _append_log,
  get_logs_since (with and without a cursor), clear_logs not
  resetting seq so reconnecting clients still progress. Passes.
- routes.export imports cleanly in the studio venv and /logs/stream
  shows up in router.routes.
- bun run build: tsc -b plus vite build, no TypeScript errors.

No existing export behavior is changed. If the subprocess, the SSE
endpoint, or the frontend hook fails, the export itself still runs to
completion the same way it did before, with or without logs visible.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* export dialog: trim bootstrap noise, scope logs per screen, show realpath

Several follow-ups to the live export log work:

1. Worker bootstrap noise (transformers venv activation, Unsloth banner,
   "Top GGUF/hub models" lists, vision detection, 2k-step weight load
   bar) is dropped from the export-dialog stream. A threading.Event
   gate in worker.py defaults closed and only opens once _handle_export
   actually starts; until then the reader thread still echoes lines to
   the saved console fd for debugging but does not push them onto the
   resp_queue. The orchestrator already spawns a fresh subprocess for
   every checkpoint load, so the gate is naturally reset between runs.

2. tqdm in non-tty mode defaults to a 10s mininterval, which makes
   multi-step bars look frozen in the panel. Set TQDM_MININTERVAL=0.5
   in the worker env so any tqdm-driven progress emits more often.

3. The dialog's useExportLogs hook now also clears its line buffer
   when exportMethod or open changes, so re-opening the dialog into a
   different action's screen no longer shows the previous action's
   saved output. A useElapsedSeconds tick + "Working Xs" badge in the
   log header gives users a visible sign that long single-step phases
   (cache copies, GGUF conversion) are still running when no new lines
   are arriving.

4. ExportBackend.export_{merged,base,gguf,lora} now return
   (success, message, output_path); the worker forwards output_path on
   each export_*_done response, the orchestrator's _run_export passes
   it to routes/export.py, which surfaces it via
   ExportOperationResponse.details.output_path. The dialog's Export
   Complete screen renders the resolved on-disk realpath under "Saved
   to" so users can find their exported model directly.

* fix(cli): unpack 3-tuple return from export backend

ExportOrchestrator.export_{merged,base,gguf,lora} now return
(success, message, output_path) so the studio dialog can show
the on-disk realpath. The CLI still unpacked 2 values, so every
`unsloth export --format ...` crashed with ValueError before
reporting completion. Update the four call sites and surface
output_path via a "Saved to:" echo.

* fix(studio): anchor export log SSE cursor at run start

The export dialog SSE defaulted its cursor to get_current_log_seq()
at connect time, so any line emitted between the POST that kicks
off the export and the client opening the stream was buffered with
seqs 1..k and then skipped (seq <= cursor). Long-running exports
looked silent during their first seconds.

Snapshot _log_seq into _run_start_seq inside clear_logs() and
expose it via get_run_start_seq(). The SSE default cursor now uses
that snapshot, so every line emitted since the current run began
is reachable regardless of when the client connects. Old runs
still can't leak in because their seqs are <= the snapshot.

* fix(studio): reconnect export log SSE on stream drop

useExportLogs launched streamExportLogs once per exporting
transition and recorded any drop in .catch(). Long GGUF exports
behind a proxy with an idle kill-timeout would silently lose the
stream for the rest of the run even though the backend already
supports Last-Event-ID resume. The "retry: 3000" directive emitted
by the backend is only meaningful to native EventSource; this
hook uses a manual fetch + ReadableStream parse so it had no
effect.

Wrap streamExportLogs in a retry loop that tracks lastSeq from
ExportLogEvent.id and passes it as since on reconnect. Backoff is
exponential with jitter, capped at 5s, reset on successful open.
The loop stops on explicit backend `complete` event or on effect
cleanup.

* fix(studio): register a second command so Typer keeps `export` as a subcommand

The CLI export unpacking tests wrap `unsloth_cli.commands.export.export`
in a fresh Typer app with a single registered command. Typer flattens a
single-command app into that command, so the test's
`runner.invoke(cli_app, ["export", ckpt, out, ...])` treats the leading
`"export"` token as an unexpected extra positional argument -- every
parametrized case failed with:

    Got unexpected extra argument (.../out)

Register a harmless `noop` second command so Typer preserves subcommand
routing and the tests actually exercise the 3-tuple unpack path they
were written to guard.

Before: 4 failed
After:  4 passed

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: studio-install <studio@local.install>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
2026-04-14 08:55:43 -07:00
Daniel Han
44082cf88e
Studio: anchor ctx-slider warning threshold at 4096 when weights exceed VRAM (#5014)
* Studio: anchor ctx-slider warning threshold at 4096 when weights exceed VRAM

The chat settings sheet's ctx slider reads `max_context_length` from
`/api/inference/status` and renders

    Exceeds estimated VRAM capacity (N tokens). The model may use
    system RAM.

when the user drags the slider above that value. For models whose
weights fit on some GPU subset, `_max_context_length` was already set
to the binary-search cap and the warning fired correctly.

For models whose weights exceed 90% of every GPU subset's free memory
(e.g. MiniMax-M2.7-GGUF at 131 GB on a 97 GB GPU), the ceiling-probe
loop never matched a subset, so `max_available_ctx` stayed at the
native context (e.g. 196608). The slider ran all the way to native
with no indication that any value above the 4096 spec default would
trigger `--fit on` and degrade performance.

Anchor `max_available_ctx` at `min(4096, native_context_length)` when
no subset fits, so the warning fires at the right threshold and the
user sees the correct safe-zone / warning-zone split:

    Before (MiniMax-M2.7 on 97 GB GPU):
      slider 0 .. 196608, warning threshold = 196608  (never fires)

    After:
      slider 0 .. 196608, warning threshold = 4096    (fires correctly)

No frontend changes required: `chat-settings-sheet.tsx` already
consumes `ggufMaxContextLength` (= status.max_context_length) as the
warning threshold and `ggufNativeContextLength` as the slider max.

Adds tests/test_llama_cpp_max_context_threshold.py covering
weights-exceed-VRAM (single / multi-GPU), a native-ctx below the 4096
fallback case (don't lie about supported ctx), fittable-model
regressions (small / multi-GPU / tiny on huge GPU), and the
`max_context_length` property's fallback semantics.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-14 08:53:49 -07:00
Daniel Han
b2f80f210e
Studio: make GGUF disk-space preflight cache-aware (#5012)
* Studio: make GGUF disk-space preflight cache-aware

The pre-download disk check in LlamaCppBackend.load_model compared the
repo's total GGUF size against free disk without crediting bytes
already present in the Hugging Face cache. Re-loading a large cached
model (e.g. MiniMax-M2.7-GGUF at 131 GB) then failed cold with
"Not enough disk space to download any variant" whenever free disk
was below the full weight footprint, even though nothing actually
needed to be downloaded.

Subtract bytes already on disk via try_to_load_from_cache before
comparing against free space. A partial blob (interrupted download) is
not credited, so a second attempt still allocates room to finish the
download. The log line now also surfaces how much is already cached.

Adds tests/test_llama_cpp_cache_aware_disk_check.py covering the
fully-cached, partial-cache-insufficient-disk, partial-cache-enough-disk,
cold-cache, incomplete-blob, and zero-size-path-info cases. Sparse
tempfiles keep the GB-scale scenarios cheap to simulate.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-14 08:53:37 -07:00
Daniel Han
767fa8cade
Studio: honor explicit GGUF ctx and default to 4096 when weights exceed VRAM (#5011)
* Studio: honor explicit GGUF ctx and default to 4096 when weights exceed VRAM

The load-time auto-fit in LlamaCppBackend.load_model had two issues for
models whose weights do not fit on any GPU subset (the common case for
large MoE GGUFs such as MiniMax-M2.7, Qwen3.5-397B-A17B, etc.):

1. Auto mode (max_seq_length=0) left effective_ctx at the model's native
   context when no subset passed the 90% fit check. The UI slider then
   landed on e.g. 196608 for MiniMax-M2.7, far above anything usable.
   Default the auto-pick to 4096 so the UI starts at a sane value; the
   slider ceiling stays at the native context so the user can still
   opt in to longer contexts and receive the "might be slower" warning.

2. Explicit ctx was silently shrunk when weights fit but the requested
   KV overflowed the 90% budget. The shrink loop emitted -c <capped>
   -ngl -1 without informing the caller, so a user who had opted into
   a longer context via the UI never actually got it. Drop the shrink
   loop on the explicit path and emit -c <user_ctx> --fit on instead,
   letting llama-server flex -ngl (CPU layer offload).

Adds tests/test_llama_cpp_context_fit.py covering both paths, the
file-size-only fallback when KV metadata is missing, non-regression on
fittable auto-pick, and platform-agnostic input shape.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-14 08:53:25 -07:00
Datta Nimmaturi
da78c6be71
[Studio] Install flash attn at setup time for linux (#4979)
* [Studio] Install flash attn at setup time for linux

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup changes

Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Test cases

* wheel_utils: narrow url_exists exceptions and log at debug level

---------

Signed-off-by: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
2026-04-14 16:40:17 +04:00
Datta Nimmaturi
dccc0ebada
[Studio] Show non exported models in chat UI (#4892)
* Show non exported models in chat UI

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Distinguish b/w LoRa and full fine tune saves. Cleanup

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
2026-04-14 15:03:58 +04:00
Daniel Han
9a261aec5f
Studio: Expose openai and anthropic compatible external API end points (#4956)
* Studio: add API key authentication for programmatic access

External users want to hit the Studio API (chat completions with tool
calling, training, export, etc.) without going through the browser
login flow. This adds sk-unsloth- prefixed API keys that work as a
drop-in replacement for JWTs in the Authorization: Bearer header.

Backend:
- New api_keys table in SQLite (storage.py)
- create/list/revoke/validate functions with SHA-256 hashed storage
- API key detection in _get_current_subject before the JWT path
- POST/GET/DELETE /api/auth/api-keys endpoints on the auth router

Frontend:
- /api-keys page with create form, one-time key reveal, keys table
- API Keys link in desktop and mobile navbar
- Route registered with requireAuth guard

Zero changes to any existing route handler -- every endpoint that uses
Depends(get_current_subject) automatically works with API keys.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use actual origin in API key usage examples

The examples on /api-keys were hardcoded to localhost:8888 which is
wrong for remote users. Use window.location.origin so the examples
show the correct URL regardless of where the user is connecting from.

* Add `unsloth studio run` CLI command for one-liner model serving

Adds a `run` subcommand that starts Studio, loads a model, creates an
API key, and prints a ready-to-use curl command -- similar to
`ollama run` or `vllm serve`.

Usage: unsloth studio run -m unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add end-to-end tests for `unsloth studio run` and API key usage

Tests the 4 usage examples from the API Keys page:
1. curl basic (non-streaming) chat completions
2. curl streaming (SSE) chat completions
3. OpenAI Python SDK streaming completions
4. curl with tools (web_search + python)

Also tests --help output, invalid key rejection, and no-key rejection.
All 7 tests pass against Qwen3-1.7B-GGUF.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add /v1/completions, /v1/embeddings, /v1/responses endpoints and --parallel support

- llama_cpp.py: accept n_parallel param, pass to llama-server --parallel
- run.py: plumb llama_parallel_slots through to app.state
- inference.py: add /completions and /embeddings as transparent proxies to
  llama-server, add /responses as application-level endpoint that converts
  to ChatCompletionRequest; thread n_parallel through load_model
- studio.py: set llama_parallel_slots=4 for `unsloth studio run` path

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make /v1/responses endpoint match OpenAI Responses API format

The existing /v1/responses shim returned Chat Completions format, which
broke OpenAI SDK clients using openai.responses.create(). This commit
replaces the endpoint with a proper implementation that:

- Returns `output` array with `output_text` content parts instead of
  `choices` with `message`
- Uses `input_tokens`/`output_tokens` instead of `prompt_tokens`/
  `completion_tokens` in usage
- Sets `object: "response"` and `id: "resp_..."`
- Emits named SSE events for streaming (response.created,
  response.output_text.delta, response.completed, etc.)
- Accepts all OpenAI Responses API fields (tools, store, metadata,
  previous_response_id) without erroring -- silently ignored
- Maps `developer` role to `system` and `input_text`/`input_image`
  content parts to the internal Chat format

Adds Pydantic schemas for request/response models and 23 unit tests
covering schema validation, input normalisation, and response format.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: add Anthropic-compatible /v1/messages endpoint (#4981)

* Add Anthropic-compatible /v1/messages endpoint with tool support

Translate Anthropic Messages API format to/from internal OpenAI format
and reuse the existing server-side agentic tool loop. Supports streaming
SSE (message_start, content_block_delta, etc.) and non-streaming JSON.
Includes offline unit tests and e2e tests in test_studio_run.py.

* Add enable_tools, enabled_tools, session_id to /v1/messages endpoint

Support the same shorthand as /v1/chat/completions: enable_tools=true
with an optional enabled_tools list uses built-in server tools without
requiring full Anthropic tool definitions. session_id is passed through
for sandbox isolation. max_tokens is now optional.

* Strip leaked tool-call XML from Anthropic endpoint content

Apply _TOOL_XML_RE to content events in both streaming and
non-streaming tool paths, matching the OpenAI endpoint behavior.

* Emit custom tool_result SSE event in Anthropic stream

Adds a non-standard tool_result event between the tool_use block close
and the next text block, so clients can see server-side tool execution
results. Anthropic SDKs ignore unknown event types.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Split /v1/messages into server-side and client-side tool paths

enable_tools=true runs the existing server-side agentic loop with
built-in tools (web_search/python/terminal). A bare tools=[...] field
now triggers a client-side pass-through: client-provided tools are
forwarded to llama-server and any tool_use output is returned to the
caller with stop_reason=tool_use for client execution.

This fixes Claude Code (and any Anthropic SDK client) which sends
tools=[...] expecting client-side execution but was previously routed
through execute_tool() and failing with 'Unknown tool'.

Adds AnthropicPassthroughEmitter to convert llama-server OpenAI SSE
chunks into Anthropic SSE events, plus unit tests covering text
blocks, tool_use blocks, mixed, stop reasons, and usage.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix httpcore GeneratorExit in /v1/messages passthrough stream

Explicitly aclose aiter_lines() before the surrounding async with
blocks unwind, mirroring the prior fix in external_provider.py
(a41160d3) and cc757b78's RuntimeError suppression.

* Wire stop_sequences through /v1/messages; warn on tool_choice

Plumb payload.stop_sequences to all three code paths (server-side
tool loop, no-tool plain, client-side passthrough) so Anthropic SDK
clients setting stop_sequences get the behavior they expect. The
llama_cpp backend already accepted `stop` on both generate_chat_
completion and generate_chat_completion_with_tools; the Anthropic
handler simply wasn't passing it.

tool_choice remains declared on the request model for Anthropic SDK
compatibility (the SDK often sets it by default) but is not yet
honored. Log a structured warning on each request carrying a non-
null tool_choice so the silent drop is visible to operators.

* Wire min_p / repetition_penalty / presence_penalty through /v1/messages

Align the Anthropic endpoint's sampling surface with /v1/chat/completions.
Adds the three fields as x-unsloth extensions on AnthropicMessagesRequest
and threads them through all three code paths: server-side tool loop,
no-tool plain, and client-side passthrough.

The passthrough builder emits "repeat_penalty" (not "repetition_penalty")
because that is llama-server's field name; the backend methods already
apply the same rename internally.

* Fix block ordering and prev_text reset in non-streaming tool path

_anthropic_tool_non_streaming was building the response by appending
all tool_use blocks first, then a single concatenated text block at
the end — losing generation order and merging pre-tool and post-tool
text into one block. It also never reset prev_text between synthesis
turns, so the first N characters of each post-tool turn were dropped
(where N = length of the prior turn's final cumulative text).

Rewrite to build content_blocks incrementally in generation order,
matching the streaming emitter's behavior: deltas within a turn are
merged into the trailing text block, tool_use blocks interrupt the
text sequence, and prev_text is reset on tool_end so turn N+1 diffs
against an empty baseline.

Caught by gemini-code-assist[bot] review on #4981.

* Make test_studio_run.py e2e tests pytest-compatible

Add a hybrid session-scoped studio_server fixture in conftest.py that
feeds base_url / api_key into the existing e2e test functions. Three
invocation modes are now supported:

1. Script mode (unchanged) — python tests/test_studio_run.py
2. Pytest + external server — point at a running instance via
   UNSLOTH_E2E_BASE_URL / UNSLOTH_E2E_API_KEY env vars, no per-run
   GGUF load cost
3. Pytest + fixture-managed server — pytest drives _start_server /
   _kill_server itself via --unsloth-model / --unsloth-gguf-variant,
   CI-friendly

The existing _start_server / _kill_server helpers and main() stay
untouched so the script entry point keeps working exactly as before.
Test function signatures are unchanged — the (base_url, api_key)
parameters now resolve via the new fixtures when running under
pytest.

* Rename test_studio_run.py -> test_studio_api.py

The file is entirely about HTTP API endpoint testing (OpenAI-compatible
/v1/chat/completions, Anthropic-compatible /v1/messages, API key auth,
plus a CLI --help sanity check on the command that runs the API). None
of its tests cover training, export, chat-UI, or internal-Python-API
concerns.

The old name misleadingly suggested "tests for the unsloth studio run
CLI subcommand" — the new name reflects the actual scope.

Updates:
- git mv the file (rename tracked, history preserved)
- Rewrite opening docstring to state the API surface focus and call
  out what is explicitly out of scope
- Update all 4 Usage-block path references to the new filename
- LOG_FILE renamed to test_studio_api.log
- conftest.py fixture import rewritten from test_studio_run to
  test_studio_api, plus 7 docstring/comment references updated

No functional changes to test logic, signatures, or main().

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Fix httpcore asyncgen cleanup in /v1/messages and /v1/completions

The earlier fix in 985e92a9 was incomplete: it closed aiter_lines()
explicitly but still used `async with httpx.AsyncClient()` /
`async with client.stream()` inside the generator. When the generator
is orphaned (e.g. client disconnects mid-stream and Starlette drops
the StreamingResponse iterator without explicitly calling aclose()),
Python's asyncgen finalizer runs the cleanup in a DIFFERENT task than
the one that originally entered the httpx context managers. The
`async with` exits then trigger httpcore's HTTP11ConnectionByteStream
.aclose(), which enters anyio.CancelScope.__exit__ with a mismatched
task and raises RuntimeError("Attempted to exit cancel scope in a
different task"). That error escapes any user-owned try/except
because it happens during GC finalization.

Replace `async with` with manual client/response lifecycle in both
/v1/messages passthrough and /v1/completions proxy. Close the
response and client in a finally block wrapped in
`try: ... except Exception: pass`. This suppresses RuntimeError (and
other Exception subclasses) from the anyio cleanup noise while
letting GeneratorExit (a BaseException, not Exception) propagate
cleanly so the generator terminates as Python expects.

Traceback observed in user report:
  File ".../httpcore/_async/connection_pool.py", line 404, in __aiter__
      yield part
  RuntimeError: async generator ignored GeneratorExit
...
  File ".../anyio/_backends/_asyncio.py", line 455, in __exit__
      raise RuntimeError(
  RuntimeError: Attempted to exit cancel scope in a different task

* Expand unsloth studio run banner with SDK base URL and more curl examples

Add an explicit "OpenAI / Anthropic SDK base URL" line inside the info
box so SDK users don't accidentally copy the bare server URL (without
/v1) into their OpenAI/Anthropic SDK constructors and hit 404s.

Replace the single /v1/chat/completions curl example with three
labeled blocks: chat/completions, Anthropic /messages, and OpenAI
Responses. The Anthropic example includes max_tokens (Anthropic SDKs
require it even though Studio accepts None).

All examples derived from a computed sdk_base_url so the /v1 prefix
stays in sync if the public path ever changes.

* Hash API keys with HMAC-SHA256 + persistent server secret

Stores the HMAC secret in a new app_secrets singleton table. Fixes
CodeQL py/weak-sensitive-data-hashing alert on storage.py:74-76,
394-395. Refresh tokens stay on plain SHA-256 (unchanged _hash_token)
so existing user sessions survive upgrade — API keys are new on this
branch so there is no migration.

* Use PBKDF2 for API key hashing per CodeQL recommendation

HMAC-SHA256 was still flagged by py/weak-sensitive-data-hashing.
Switch to hashlib.pbkdf2_hmac, which is in CodeQL's recommended
allowlist (Argon2/scrypt/bcrypt/PBKDF2). Persistent server-side
salt stays in app_secrets for defense-in-depth. 100k iterations to
match auth/hashing.py's password hasher.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
2026-04-13 21:08:11 +04:00
Daniel Han
cad8c6ad05
Add AMD ROCm/HIP support across installer and hardware detection (#4720)
* Add ROCm detection to install.sh and expand shell tests

Add AMD ROCm GPU detection to get_torch_index_url() in install.sh.
When nvidia-smi is not found, probe for ROCm via amd-smi, /opt/rocm
version file, hipconfig, dpkg-query, and rpm.

Includes validation guard for malformed _rocm_tag, Debian epoch prefix
stripping, ROCm 7.2+ cap to rocm7.1 index, bitsandbytes AMD install,
and status messaging. Shell tests expanded to 23 cases.

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add ROCm torch reinstall support to install_python_stack.py

Add _detect_rocm_version() and _ensure_rocm_torch() to detect when a
Linux host has ROCm but the venv received CPU-only torch, and reinstall
with the correct ROCm wheels. Covers ROCm 6.0 through 7.1 with a
30-second timeout on the torch GPU probe subprocess.

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add ROCm support to llama.cpp prebuilt installer

Add has_rocm field to HostInfo, extend detect_host() to probe for ROCm
via hipcc/amd-smi/rocm-smi/ROCM_PATH, and route ROCm hosts to upstream
prebuilts (Linux ROCm 7.2 prebuilt with source fallback, Windows HIP
prebuilt with CPU fallback). Add linux-rocm and windows-hip install
kinds to runtime_patterns_for_choice().

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add IS_ROCM hardware flag and fix AMD error message

Add IS_ROCM flag to hardware.py detect_hardware() (set when
torch.version.hip is present, DeviceType stays CUDA). Export IS_ROCM
from __init__.py. Add "rocm" key to get_package_versions().

Replace "We do not support AMD" error in tokenizer_utils.py with a
helpful message pointing to ROCm installation docs.

Co-authored-by: Daniel Han <danielhanchen@gmail.com>

* Add comprehensive ROCm support test suite (68 tests)

Add tests/studio/install/test_rocm_support.py covering all ROCm code
paths across install_llama_prebuilt.py, install_python_stack.py,
hardware.py, tokenizer_utils.py, and install.sh. All tests use mocks
and run without AMD hardware.

Covers: asset selection (11), runtime patterns (5), HostInfo (4),
ROCm version detection (9), torch reinstall (9), index mapping (8),
hardware flag (8), tokenizer message (2), install.sh structure (10),
and live regression (1).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden ROCm support: probe error handling, version cap, validation

Address review findings from 8 independent reviewers:

- Wrap _ensure_rocm_torch() torch probe in try/except for
  TimeoutExpired and OSError so a hung or broken torch import does not
  crash the installer (8/8 reviewers flagged this)
- Add torch>=2.4,<2.11.0 version cap to the ROCm reinstall path to
  prevent installing unsupported torch 2.11.0 from the rocm7.1 index
- Use with-statement for file reads in _detect_rocm_version() to avoid
  resource leaks
- Handle ROCM_PATH="" correctly (use `or "/opt/rocm"` instead of
  default parameter to avoid relative path resolution)
- Strengthen shell validation guard from rocm[0-9] to rocm[1-9] to
  reject rocm0.x tags that would produce nonexistent PyTorch index URLs
- Switch shell version cap from blocklist to allowlist (rocm6.*|rocm7.0*
  |rocm7.1* pass through, everything else caps to rocm7.1) so future
  ROCm 10+ does not fall through to a nonexistent index
- Add sorted() to _ROCM_TORCH_INDEX lookup for defensive ordering
- Fix test_probe_timeout_handled: replace zero-assertion test with
  proper assertions verifying reinstall proceeds after timeout

* Clean up rocm_paths list construction in detect_host()

Filter None from the ROCM_PATH env var lookup at list construction time
instead of relying on the inline `if p` guard in the any() call.

* Require actual AMD GPU presence before selecting ROCm paths

All 8 reviewers across 2 cycles independently flagged that ROCm
detection used toolkit/filesystem hints (hipcc, /opt/rocm, rocm-core)
as a proxy for GPU presence, which would misroute CPU-only or NVIDIA
hosts that happen to have ROCm tools installed.

Now all 3 detection points (install.sh, install_python_stack.py,
install_llama_prebuilt.py) probe for an actual AMD GPU before
entering the ROCm path:

- install.sh: check rocminfo for gfx* GPU names, or amd-smi list
  for device rows, before version detection
- install_python_stack.py: new _has_rocm_gpu() function probes
  rocminfo and amd-smi list before _ensure_rocm_torch() proceeds
- install_llama_prebuilt.py: detect_host() probes rocminfo/amd-smi
  list instead of just checking tool existence or directory paths

Also:
- Shell test mock amd-smi now handles "list" subcommand
- Python tests updated to mock _has_rocm_gpu where needed
- Added test_no_gpu_with_rocm_tools_skips to verify the new guard
- Test index lookups now use sorted() to match production code

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden hipconfig version parsing and torch probe compatibility

- Add parts[1].isdigit() check in hipconfig version parsing to handle
  versions like "6.3-HIP" where the minor component has non-numeric
  suffix (strip "-" prefix before int() conversion)
- Use getattr() in torch probe subprocess to safely handle old or
  custom torch builds that may lack torch.version.hip/cuda attributes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Strengthen AMD GPU detection and add NVIDIA precedence guard

- Change amd-smi list detection from any-non-empty-output to requiring
  "gpu" marker in output, matching the shell-side NR>1 check. Prevents
  false positives from header-only amd-smi list output.
- Add nvidia-smi check at the top of _ensure_rocm_torch() so mixed
  AMD+NVIDIA hosts preserve NVIDIA precedence (matching install.sh and
  install_llama_prebuilt.py behavior).
- Apply the same amd-smi marker fix to install_llama_prebuilt.py
  detect_host() for consistency.

* Add Windows-specific ROCm/HIP detection in detect_host()

The previous detect_host() ROCm check used rocminfo and amd-smi list
which are Linux-only tools. On Windows, has_rocm would always be False,
making the Windows HIP prebuilt path at line 1794 unreachable.

Now detect_host() uses platform-specific detection:
- Linux: rocminfo (check for gfx GPU names) or amd-smi list
- Windows: hipinfo.exe, amd-smi, or amdhip64.dll on PATH

This allows Windows AMD users to get the HIP prebuilt binary instead
of silently falling through to the CPU prebuilt.

* Add AMD ROCm gaps: Mamba/SSM source builds, GPU monitoring, Windows messaging, RDNA expansion

- worker.py: Add HIP detection to causal-conv1d/mamba-ssm probe, check
  for hipcc before ROCm source builds, improve status messages and error
  reporting, add timeout and uv support for the source build fallback
- amd.py: New AMD GPU monitoring module via amd-smi metric --json,
  mirroring nvidia.py structure (utilization, temperature, power, VRAM)
- hardware.py: Branch to amd.py when IS_ROCM is True for GPU utilization,
  visible GPU queries, and physical GPU count
- install_python_stack.py: Detect AMD GPUs on Windows and warn that
  ROCm-enabled PyTorch must be installed manually
- kernels/utils.py: Expand is_rdna() to cover RDNA2 (gfx1030-1032),
  RDNA3 (gfx1102-1103), RDNA3.5 (gfx1150-1152) alongside existing entries
- tests: Add 32 new tests covering all changes (95/95 pass)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden ROCm detection, fix VRAM heuristic, and expand RDNA2 coverage

- Windows ROCm detection: validate actual GPU presence via hipinfo/amd-smi
  output markers instead of just checking tool existence on PATH
- _ensure_rocm_torch: validate nvidia-smi actually reports a GPU before
  giving NVIDIA precedence (fixes AMD-only hosts with stale NVIDIA tools)
- amd.py _parse_numeric: handle dict-shaped metric objects from newer
  amd-smi versions ({"value": 10, "unit": "W"}) and strip MiB/GiB units
- amd.py VRAM heuristic: raise threshold from 100k to 10M to correctly
  handle MI300X (192 GB = 196608 MB) and other high-VRAM GPUs
- amd.py visible GPU: use AMD-reported GPU IDs instead of enumerate index
  so non-dense sets like CUDA_VISIBLE_DEVICES=1,3 report correctly
- install.sh: add ROCm <6.0 minimum version guard (no PyTorch wheels
  exist for older versions); fix rocm7.1* glob to not match rocm7.10+
- is_rdna: add gfx1033-1036 for RDNA2 mobile GPUs (RX 6600M etc.)
- worker.py: increase ROCm source build timeout from 600s to 1800s;
  fix success log message for ROCm source builds
- Tests: update mocks for _has_usable_nvidia_gpu, add RDNA2 target asserts

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add HIP_VISIBLE_DEVICES support, unit-aware VRAM parsing, Windows GPU validation

- hardware.py: check HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES on ROCm
  before falling back to CUDA_VISIBLE_DEVICES, so multi-GPU AMD setups with
  HIP-specific env vars report the correct visible device set
- amd.py: add _parse_memory_mb() that reads "unit" from dict-shaped amd-smi
  JSON (e.g. {"value": 192, "unit": "GiB"}) and converts to MB correctly;
  fixes MI300X VRAM misreported as 0.19 GB instead of 192 GB
- install_python_stack.py: Windows AMD warning now validates actual GPU
  presence via hipinfo/amd-smi output markers before printing
- install_llama_prebuilt.py: restore amdhip64.dll fallback for Windows HIP
  detection after tool-based checks, so Windows HIP installs without CLI
  tools on PATH are still detected
- hardware.py: fix IS_ROCM comment to accurately describe its role

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix HIP_VISIBLE_DEVICES empty-string handling in GPU visibility spec

Use explicit None checks instead of Python `or` operator when reading
HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES, so that an empty string
("") is correctly honored as "no visible GPUs" rather than silently
falling through to CUDA_VISIBLE_DEVICES on mixed ROCm+CUDA systems.

* Fix IS_ROCM test assertion for multi-line formatting

* Cap torchvision/torchaudio versions, remove amdhip64.dll fallback, fix visible GPU count

- Cap torchvision<0.26.0 and torchaudio<2.11.0 alongside torch<2.11.0 in
  both install.sh and install_python_stack.py to prevent resolver from
  selecting incompatible companion packages from ROCm wheel index
- Remove amdhip64.dll fallback in Windows ROCm detection (DLL presence
  without hipinfo/amd-smi is not proof of GPU existence)
- Fix get_visible_gpu_count() to use _get_parent_visible_gpu_spec() which
  respects HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES on ROCm hosts

* Attribute is_rdna() RDNA2/3/3.5/4 expansion to PR #4428

The is_rdna() expansion to cover RDNA2 (gfx1030-1036), RDNA3
(gfx1100-1103), RDNA3.5 (gfx1150-1152), and RDNA4 (gfx1200-1201)
architectures is based on the original work from PR #4428.

Co-authored-by: GoldenGrapeGentleman <yueyuan@amd.com>
Co-authored-by: billishyahao <bill.he@amd.com>

* Support AMD Radeon for studio (#4770)

Co-authored-by: Iswarya Alex <iswarya.alex@amd.com>

* Remove ROCm test files from main PR

Move test_rocm_support.py and shell test additions to a separate PR
to keep the main ROCm support PR focused on implementation changes.

* Fix installer and hardware detection issues for PR #4720

- Fix empty _tri_arg passed to uv pip install in Radeon path (causes
  "Empty field is not allowed for PEP508" error)
- Fix Radeon fallback: use ROCm index instead of CPU-only when
  repo.radeon.com is unreachable (TORCH_INDEX_URL already has ROCm)
- Use $TORCH_CONSTRAINT in fallback paths instead of hardcoded strings
- Fix _pick_radeon_wheel: relax suffix to match manylinux_2_28_x86_64
  wheels (AMD Radeon repo does not use bare linux_x86_64 platform tag)
- Fix IS_ROCM export: use __getattr__ so callers always see the live
  value after detect_hardware() runs
- Fix apply_gpu_ids: set HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES
  on ROCm so _get_parent_visible_gpu_spec picks up narrowed GPU set
- Fix _parse_memory_mb: distinguish GB (1000 MB) from GiB (1024 MiB)
- Add amd-smi version as a fallback in _detect_rocm_version
- Fix trailing whitespace and missing newline at EOF in install.sh

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix GPU detection false positives and add missing health groups

- Fix _has_rocm_gpu() false positive: require "GPU: <number>" data rows
  from amd-smi list, not just header containing "gpu"
- Apply same fix in detect_host() in install_llama_prebuilt.py
- Add runtime_payload_health_groups for linux-rocm and windows-hip so
  partial/corrupt ROCm/HIP prebuilt installs are properly detected
- Add bitsandbytes install to Radeon fallback paths (was only in the
  success path, skipped when repo.radeon.com was unreachable)
- Keep DEVICE/CHAT_ONLY as direct imports in __init__.py (matching main)
  and only use __getattr__ for IS_ROCM

* Fix _ensure_rocm_torch and Windows AMD warning false positives

- _ensure_rocm_torch: only skip when HIP is already present, not for
  CUDA builds (which are unusable on AMD-only hosts). Fixes the case
  where a venv has a stale CUDA wheel and the repair step is skipped.
- Windows AMD warning: use GPU data row check (same as Linux fix) to
  avoid false positives from amd-smi list header-only output.

* Fix amd-smi GPU detection for GPU[N] output format

Older amd-smi versions output "GPU[0] : Card series: ..." instead of
"GPU: 0". The regex now matches both "GPU: <digit>" and "GPU[<digit>"
formats to detect actual GPU data rows.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden AMD GPU detection against false positives

- install.sh: replace weak amd-smi list check (awk 'NR>1 && NF') with
  strict pattern matching GPU data rows (/^GPU[[:space:]]*[:\[]/)
- All files: reject rocminfo gfx000 (CPU HSA agent) by requiring
  gfx[1-9] instead of gfx[0-9] in the rocminfo GPU probe
- Fixes false positives on hosts with ROCm tools but no AMD GPU

* Remove duplicate comment from pre-commit merge

* Refactor: deduplicate AMD detection, consolidate bitsandbytes, clean up imports

- Extract _has_amd_rocm_gpu() shell function to avoid duplicating the
  rocminfo/amd-smi GPU detection logic in get_torch_index_url and
  the Radeon auto-detect block
- Consolidate bitsandbytes install into a single case block after torch
  install (was duplicated 4 times across Radeon success/fallback paths)
- Move math and re imports to top of amd.py (were inline in functions)
- Add _smi_query() helper in hardware.py to centralize IS_ROCM backend
  selection for get_gpu_utilization and get_visible_gpu_utilization

Addresses Gemini code review suggestions.

* Fix VRAM parsing for string values and GB/GiB consistency

- Extract unit from string-valued VRAM fields (e.g. "192 GiB") so
  _parse_memory_mb correctly applies the unit multiplier instead of
  treating the value as bare MB
- Treat GB and GiB identically (both as binary x1024) since GPU tools
  including amd-smi use binary units even when labeling them "GB"
- Fixes incorrect VRAM reporting on MI300-class cards (was showing
  ~0.19 GB instead of 192 GB for string-valued outputs)

* Add --no-cache to uv for ROCm HIP source builds

Avoid stale cache artifacts from partial HIP source builds when
uv is used for causal-conv1d/mamba-ssm compilation on ROCm.
The pip path already uses --no-cache-dir; this adds the uv equivalent
(--no-cache) only when is_hip is True.

* Fix critical: initialize _amd_gpu_radeon before case block

_amd_gpu_radeon was only set inside the */rocm*) case arm, so on
NVIDIA/CPU/macOS paths where TORCH_INDEX_URL does not contain "rocm",
the variable was unbound. With set -u (nounset) enabled, this crashes
the installer for every non-AMD user.

Move initialization to before the case block so it is always defined.

* Fix Windows AMD: route has_rocm hosts to HIP prebuilt path

resolve_release_asset_choice was selecting windows-cpu for all Windows
x86_64 hosts including those with has_rocm=True. Windows AMD users
should fall through to resolve_upstream_asset_choice which tries the
HIP prebuilt first. Add "not host.has_rocm" guard to the published
windows-cpu selection.

* Harden ROCm detection, Radeon wheel fallback, and HIP visibility

Addresses review findings from parallel reviewers on PR #4720:

- install.sh: add _has_usable_nvidia_gpu() helper requiring nvidia-smi -L
  to actually list a GPU before treating the host as NVIDIA. Fixes the
  stale-nvidia-smi-on-PATH regression where AMD-only hosts fell into the
  CUDA branch.
- install.sh: fix hipconfig awk blocks to propagate a non-zero exit code
  when the output is not a recognisable version string, so the ||-chain
  continues to dpkg-query / rpm instead of terminating early.
- install.sh: fail-closed on Radeon wheel fallback. When torch,
  torchvision or torchaudio is missing from the Radeon repo for the
  active Python tag, fall back to the standard ROCm index instead of
  silently mixing Radeon wheels with PyPI defaults. Quote all wheel
  arguments individually so wheel filenames cannot be word-split or
  glob-expanded.
- install_llama_prebuilt.py: detect_host() now requires nvidia-smi -L to
  list a GPU before setting has_physical_nvidia. Routes AMD ROCm hosts
  with a broken leftover nvidia-smi to the ROCm path instead of
  misclassifying them as NVIDIA.
- install_llama_prebuilt.py: scan upstream assets for any rocm-<version>
  prebuilt instead of hard-coding rocm-7.2, so ROCm 6.x / 7.0 / 7.1 / 7.3+
  users pick up a matching upstream prebuilt when one exists.
- install_llama_prebuilt.py: validate_server() adds --n-gpu-layers 1 for
  linux-rocm and windows-hip hosts, so new HIP prebuilts are preflighted
  on the GPU path instead of passing validation on CPU only.
- install_llama_prebuilt.py: restore the published windows-cpu fallback
  for AMD Windows hosts without a HIP prebuilt so hash-approved bundles
  are still preferred over the raw upstream CPU asset.
- install_python_stack.py: drop the /opt/rocm / hipcc gate in
  _ensure_rocm_torch() and rely on _has_rocm_gpu(). Runtime-only ROCm
  installs (package-managed minimal installs, Radeon software) that ship
  amd-smi / rocminfo without hipcc can now repair a CPU-only venv via
  "unsloth studio update". Adds an explicit IS_WINDOWS / IS_MACOS guard.
- studio/backend/utils/hardware/amd.py: honour HIP_VISIBLE_DEVICES /
  ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES in
  get_primary_gpu_utilization(). A process restricted to GPU 2 now
  reports metrics for GPU 2 instead of physical GPU 0. Tighten the plain
  bytes unit detection to an explicit allowlist.
- studio/backend/utils/hardware/hardware.py: route
  get_backend_visible_gpu_info()'s backend_cuda_visible_devices field
  through a helper that reads HIP_VISIBLE_DEVICES on ROCm. Drop the
  unconditional "(rocm=False)" suffix in apply_gpu_ids() logs.

* Fix round 2 regressions: ROCm validate_server and Windows HIP routing

Follow-up to 810b833b addressing review findings on the first round of
hardening commits:

- install_llama_prebuilt.py validate_server: gate --n-gpu-layers on the
  resolved install_kind instead of host.has_rocm. AMD Windows hosts
  without a HIP prebuilt fall back to windows-cpu and must not be
  validated with GPU layers; thread install_kind through from the
  caller.
- install_llama_prebuilt.py resolve_release_asset_choice: reinstate the
  "not has_rocm" guard on the published windows-cpu bundle so AMD
  Windows hosts reach resolve_upstream_asset_choice() where the new
  HIP prebuilt path lives. Prefer a published windows-hip bundle first
  when one exists, fall through to upstream HIP + upstream CPU
  otherwise.
- install_llama_prebuilt.py detect_host: also set has_physical_nvidia
  when the secondary --query-gpu block confirms a working NVIDIA GPU,
  so older nvidia-smi versions without -L support do not silently skip
  the Linux diagnostics that key off has_physical_nvidia.
- install_llama_prebuilt.py: drop redundant "import re as _re" /
  "import re as _re_rocm" local aliases in favour of the existing
  top-level "import re".
- install_python_stack.py _ensure_rocm_torch: run the AMD
  bitsandbytes install unconditionally after the HIP-torch probe so
  "unsloth studio update" on venvs that already have ROCm torch still
  gains the AMD bitsandbytes build.
- install.sh: add a non-x86_64 early-exit to get_torch_index_url() so
  aarch64 / arm64 Linux hosts do not hit the ROCm wheel index
  (PyTorch only publishes ROCm wheels for linux_x86_64).
- install.sh: add bitsandbytes install to the migrated-environment
  branch so upgrades pick it up for ROCm hosts instead of only the
  fresh-install path.
- install.sh: in the Radeon wheel path, pass version constraints +
  --no-index --find-links to uv instead of explicit wheel URLs so a
  version-compatible torch / torchvision / torchaudio triple is
  resolved, rather than picking the highest-version wheel for each
  package independently.
- studio/backend/utils/hardware/amd.py _first_visible_amd_gpu_id: fall
  through to lower-priority visibility env vars when the first entry
  is malformed (leading comma, all-whitespace first token) instead of
  silently returning GPU 0.

* Fix round 3 findings: x86_64 guard, ROCm version clip, Radeon deps

Address issues surfaced by the round 3 reviewers on top of 8636fa63:

- install_python_stack.py _ensure_rocm_torch: add the same `x86_64`
  guard that install.sh already has. Linux aarch64 / arm64 ROCm hosts
  must skip the repair path entirely; PyTorch only publishes ROCm
  wheels for linux_x86_64, and without this guard
  `unsloth studio update` aborts with a missing-wheel error on non
  x86_64 hosts.
- install_llama_prebuilt.py resolve_upstream_asset_choice: add a
  best-effort _detect_host_rocm_version() helper (reading
  /opt/rocm/.info/version, amd-smi version, hipconfig --version) and
  filter rocm_candidates to entries whose major.minor is <= host
  version. Falls back to the newest candidate only when no compatible
  one exists, so a ROCm 6.4 host downloads rocm-6.4 instead of being
  handed the numerically newest rocm-7.2 bundle (which fails preflight
  and forces a source build).
- install.sh: remove the round 2 --no-index switch from the Radeon
  wheel branch. --no-index forced uv to ignore PyPI entirely, which
  broke transitive dependency resolution (filelock, sympy, networkx,
  jinja2, fsspec, setuptools, typing-extensions, ...) on a fresh venv.
  Restore the round 1 explicit wheel URL invocation but add a
  torch / torchvision / torchaudio version-pair sanity check so a
  mismatched trio (e.g. torch 2.9.1 + torchvision 0.23.0 + torchaudio
  2.9.0) falls back to the standard ROCm index instead of installing a
  broken combination.
- install_python_stack.py _ensure_rocm_torch: restructure the
  "tag is None" path so it no longer short-circuits the bitsandbytes
  install. On a ROCm runtime older than anything in
  _ROCM_TORCH_INDEX, print the "no wheel" warning but still run the
  AMD bitsandbytes install.
- studio/backend/core/training/worker.py: restore the pre-PR
  "no timeout" behaviour for non-HIP causal-conv1d / mamba-ssm source
  builds. The round 2 "timeout = 1800 if is_hip else 300" cap aborts
  slow non-HIP builds (Linux aarch64, unsupported torch/CUDA combos)
  after 5 minutes; omit timeout for the non-HIP branch so the cap
  only applies to ROCm source builds.

* Fix round 4 findings: apply_gpu_ids env inheritance, Radeon X.Y, bitsandbytes gate

Address remaining issues surfaced by the round 4 reviewers:

- studio/backend/utils/hardware/hardware.py apply_gpu_ids: mirror the
  selection into HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES whenever
  the caller already had a ROCm visibility env var set, not only when
  IS_ROCM has already been set by detect_hardware(). Training and
  inference workers call apply_gpu_ids() before detect_hardware()
  runs, so the old guard would leave a forked ROCm worker with a
  stale HIP_VISIBLE_DEVICES mask that no longer matched the
  narrowed CUDA_VISIBLE_DEVICES selection.
- install.sh get_radeon_wheel_url: accept X.Y ROCm versions in
  addition to X.Y.Z. The `/opt/rocm/.info/version` file and some
  hipconfig versions report only two components, and the Radeon
  repository publishes both rocm-rel-X.Y.Z/ and rocm-rel-X.Y/
  directories, so treating X.Y as invalid caused Radeon hosts to fall
  back to the generic ROCm index even when a matching AMD wheel set
  existed.
- install_python_stack.py _ensure_rocm_torch: only install the AMD
  bitsandbytes build when the venv actually has a ROCm-compatible
  torch (either already present or just installed by this function).
  Previously the bitsandbytes install ran unconditionally, which
  could leave an AMD bitsandbytes layered on top of a CPU/CUDA torch
  on hosts where the ROCm runtime is older than any entry in
  _ROCM_TORCH_INDEX. Also add --force-reinstall so an existing
  CPU/CUDA bitsandbytes is replaced by the AMD build during upgrades.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix gemini findings: amd-smi metric envelope validation and dict-wrapped GPU id

Two medium-severity defensive fixes from the gemini-code-assist review on
the AMD monitoring backend:

1. _extract_gpu_metrics may return a dict where every value is None when
   amd-smi succeeds (zero exit) but the JSON envelope contains no usable
   fields (error response, unsupported card). The new _has_real_metrics
   helper lets get_primary_gpu_utilization surface available:False and
   lets get_visible_gpu_utilization skip ghost device rows so the UI
   does not render placeholder cards with empty numbers.

2. Newer amd-smi versions wrap scalar fields as {"value": 0, "unit":
   "none"}, including the per-GPU id. The previous int(raw_id) call
   silently fell back to the enumeration index in that case, losing the
   real GPU id. Routing raw_id through the existing _parse_numeric
   helper handles bare ints, floats, strings, and the dict shape
   uniformly, with a debug log on parse failure.

* Fix gemini round 2 findings: explicit length guard on ROCm version file parser

Both _detect_rocm_version (install_python_stack.py) and
_detect_host_rocm_version (install_llama_prebuilt.py) read /opt/rocm/.info/version
or $ROCM_PATH/lib/rocm_version, split on "." and unconditionally accessed
parts[1]. The surrounding broad `except Exception: pass` already swallowed
the resulting IndexError, so a one-component file like "6\n" did fall
through to the next detection source -- but the control flow relied on
exception handling instead of an explicit check.

Add `if len(parts) >= 2:` guards in both helpers so the loop falls through
on its own without raising. Behaviour is unchanged for the common multi-
component case; the previously-silent IndexError path becomes an explicit
no-op.

* Fix gemini round 3: include has_rocm in validate_server fallback path

When validate_server is called without an explicit install_kind (older
call sites that have not been updated), the fallback was only enabling
--n-gpu-layers for NVIDIA and macOS arm64 hosts. AMD ROCm Linux hosts
fell through to the CPU validation path even though the prebuilt being
exercised was a HIP binary.

Add host.has_rocm to the fallback expression so the GPU offload flag is
applied consistently with the install_kind=='linux-rocm' / 'windows-hip'
branches above.

* Fix gemini round 4: remove risky bytes-vs-MB heuristic in _parse_memory_mb

The previous heuristic divided any bare number above 10_000_000 by
1024*1024 on the assumption that large unit-less values were bytes.
This misclassified small VRAM allocations: 5 MB of used VRAM reported
as 5_242_880 bytes without a unit would be taken at face value and
render as 5_242_880 MB (~5 TB) in the monitoring UI.

Modern amd-smi always provides explicit units (MiB/GiB dict form),
and legacy amd-smi returns bare numbers in MB -- the heuristic never
had a real workload to handle. Drop it and default to MB for bare
numeric input, keeping the existing unit-aware branches for dict /
string inputs unchanged.

The unrelated gemini suggestion to "default minor to 0" in the
amd-smi version awk parser was intentionally NOT applied: rocm7.0
and rocm7.1 ship different wheel sets, so silently substituting 0
for a missing minor could install the wrong wheels. The existing
reject-and-fall-through behaviour is safer.

* Fix gemini round 5: POSIX compliance and leading-comma visibility parsing

Three medium findings from gemini-code-assist addressed in this commit:

1. _pick_radeon_wheel used grep -o and sort -V, both GNU extensions
   that are not in POSIX and break on BSD/BusyBox coreutils. install.sh
   has a #!/bin/sh shebang so the whole pipeline was rewritten as a
   single awk script that extracts all href="..." hits on each line,
   filters to wheels matching the package prefix and python tag, and
   picks the newest version via zero-padded lexical comparison. No
   external sort or grep is needed.

2. _first_visible_amd_gpu_id in the AMD monitoring backend treated a
   leading comma (e.g. HIP_VISIBLE_DEVICES=",1") as "fall through to
   the next env var", which is surprising given the clear intent to
   narrow to device 1. Filter empty tokens after the split and return
   the first real one. An all-commas value ("," / ",,,") still falls
   through because no real tokens exist; the empty-string and "-1"
   explicit-zero cases are unchanged.

The unrelated amd-smi version awk parser suggestion was not applied
(see round 4 commit message for rationale: defaulting a missing minor
to 0 could silently install the wrong ROCm wheel set).

* Fix 20-reviewer.py findings: base drift, Radeon %2B, dpkg/rpm fallback, bnb, backend label

Consolidated fix batch from a 20-parallel reviewer.py run on the current
head. Each fix is drawn from a high-consensus finding and addresses a
real bug or feature gap, not a stylistic preference.

1. install.sh: bump `unsloth>=2026.4.2` -> `unsloth>=2026.4.4` at five
   call sites so this branch no longer regresses main's version floor
   (main bumped to 2026.4.4 in #4876). Without this, merging 4720 would
   silently downgrade the minimum version pin for fresh installs.

2. install.sh: URL-decode Radeon wheel names before extracting the
   torch / torchvision / torchaudio version strings. Real wheel URLs
   from repo.radeon.com are percent-encoded ("torch-2.10.0%2Brocm7.2.0...")
   so the previous `[+-]` terminator in the sed regex never matched,
   `_torch_ver` stayed empty, `_radeon_versions_match` stayed false,
   and every Radeon consumer install silently fell back to the generic
   ROCm index. Now decode %2B -> + first, then extract, then validate.

3. install.sh: the two AMD bitsandbytes install lines were running
   `uv pip install "bitsandbytes>=0.49.1"` without `--force-reinstall`,
   so upgrades where the venv already has a CPU/CUDA bitsandbytes
   satisfying the constraint would keep the stale non-AMD wheel. Add
   `--force-reinstall --no-cache-dir` to both call sites, matching the
   pattern already used in install_python_stack.py::_ensure_rocm_torch.

4. install_python_stack.py and install_llama_prebuilt.py: add
   `dpkg-query -W rocm-core` and `rpm -q rocm-core` fallbacks to the
   Python-side ROCm version detectors so they match the chain in
   install.sh::get_torch_index_url. Package-managed ROCm installs
   (Debian/Ubuntu/RHEL/Fedora distro packages) can expose GPUs via
   rocminfo/amd-smi but still lack /opt/rocm/.info/version, hipconfig,
   or amd-smi `version` output -- without these fallbacks, `unsloth
   studio update` on such hosts returned None and skipped the ROCm
   torch repair. Also strip the dpkg epoch prefix ("1:6.3.0-1") before
   parsing so epoch-annotated packages parse correctly.

5. hardware.py: add a `_backend_label(device)` helper that returns
   "rocm" when IS_ROCM is set and the device is DeviceType.CUDA, and
   use it for every `"backend": ...` emission in JSON responses served
   to the Studio frontend. Internally we still represent ROCm hosts as
   DeviceType.CUDA (ROCm torch reuses the whole torch.cuda.* API
   surface), but the user-facing API now correctly reports "rocm" on
   AMD boxes instead of labeling them as "cuda".

All 250 simulation scenarios pass (was 233 before this batch: added 17
new regression tests covering the version pin, %2B decoding, bnb
force-reinstall flags, dpkg/rpm fallback presence, and the
_backend_label helper's four-way truth table).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix gemini round 6 + URL audit: amd.py defensive checks, rocm6.5+ clip to 6.4

Two rounds of fixes in one commit, plus a full URL audit of every PyPI /
download.pytorch.org / repo.radeon.com reference the PR introduces.

amd.py (4 medium gemini findings on commit b3627bc2):

1. _extract_gpu_metrics used `and vram_total_mb` as part of the vram_util
   gate. The follow-up `vram_total_mb > 0` already handles the division
   guard, but the truthiness check was redundant and slightly surprising
   for a 0.0 valid value. Replace with explicit `is not None and > 0`
   for both vram_util and power_util.

2. get_physical_gpu_count called `data.get("gpu", ...)` without guarding
   for non-dict envelopes. A scalar / string JSON response from amd-smi
   would raise AttributeError. Add an isinstance(data, dict) check and
   return None for unexpected shapes.

3. get_visible_gpu_utilization had the same .get() exposure on the outer
   envelope. Rewrite the gpu_list extraction as an explicit
   list/dict/else cascade so a malformed scalar envelope produces
   gpu_list=[data] and continues without raising.

4. The same function's per-entry loop also called gpu_data.get() on
   whatever was inside gpu_list. If a scalar ever leaks into the list
   (directly or via the previous fix's fallback), _extract_gpu_metrics
   would raise on the first .get() inside the helper. Skip non-dict
   entries in the loop before extracting metrics.

install.sh (URL audit finding, previously flagged by 20-reviewer as #13):

5. get_torch_index_url used `rocm6.*` in the rocm tag case statement,
   which matched rocm6.5 and rocm6.6 and emitted
   download.pytorch.org/whl/rocm6.5 -- which returns HTTP 403 because
   PyTorch only publishes rocm 5.7, 6.0-6.4, 7.0-7.2. Enumerate the
   supported 6.x minors explicitly and add a rocm6.* fallback branch
   that clips to rocm6.4 (the last supported 6.x wheel set).

URL audit results (all URLs PR 4720 references):
- 14/14 download.pytorch.org/whl/{cpu,cu118,cu124,cu126,cu128,cu130,
  rocm6.0..6.4,rocm7.0..7.2} return HTTP 200.
- 9/9 repo.radeon.com/rocm/manylinux/rocm-rel-{5.7,6.0,6.1,6.2,6.3,
  6.4,7.0,7.1,7.2}/ return HTTP 200.
- X.Y.Z patch directories exist for 7.0.2, 7.1.1, 7.2.1 but NOT for
  6.3.0, 6.4.0, 6.2.1 -- install.sh already handles this via the X.Y.Z
  -> X.Y fallback sed in the Radeon wheel install block.
- Docs links (rocm.docs.amd.com, docs.unsloth.ai AMD guide) and the
  llama.cpp GitHub releases API endpoint all return 200.

Test suite: 255 -> 258. New regression coverage:
- U17: get_physical_gpu_count tolerates scalar amd-smi envelope
- U18: get_visible_gpu_utilization tolerates scalar envelope
- U19a-c: vram_util / power_util return None on zero total, but
  vram_total_gb still echoes 0.0 (not None)
- A_rocm{6.5,6.6,6.9}_clips_to_rocm64: install.sh clips unsupported
  6.x minors to rocm6.4 instead of producing a 403 index URL

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix reviewer.py round 2: tokenizer AMD multi-GPU, --no-torch bnb, main.py backend label

Three high-confidence findings from a second 20-parallel reviewer.py run
on commit 7effb3ae. Triaged 15 total findings and applied the three that
were confirmed as real bugs; the rest were either false positives (e.g.
"migrated AMD venv not repaired" -- _ensure_rocm_torch runs downstream
via setup.sh regardless), design decisions (e.g. visibility mask env
vars not consulted in installer detection), or edge cases the existing
fallback logic already handles.

1. unsloth/tokenizer_utils.py [6/20]: the multi-GPU guard's shell probe
   runs `nvidia-smi --query-gpu=memory.used`, catches the failure, then
   only raises if `torch.cuda.is_available()` is False. On ROCm torch,
   torch.cuda.is_available() returns True (ROCm reuses the torch.cuda.*
   API), so the guard becomes dead code on AMD hosts and multi-GPU AMD
   setups slip through even though unsloth does not support them yet.
   Add a torch.cuda.device_count() > 1 fallback inside the except so
   AMD multi-visible-device setups are flagged consistently with the
   original CUDA memory check.

2. install.sh [1/20]: the fresh-install bitsandbytes block for AMD ROCm
   ran unconditionally when TORCH_INDEX_URL matched `*/rocm*`, even when
   SKIP_TORCH=true (from --no-torch or Intel Mac auto-detect). A user
   running `install.sh --no-torch` on an AMD host would still pull in
   bitsandbytes despite explicitly asking for GGUF-only mode. Wrap the
   case block in an outer `[ "$SKIP_TORCH" = false ]` guard.

3. studio/backend/main.py [3/20]: the /api/system endpoint returned
   `"device_backend": get_device().value`, which is "cuda" on ROCm
   hosts (because ROCm torch piggybacks on torch.cuda). Other endpoints
   (hardware.py) already use the _backend_label helper which swaps
   "cuda" -> "rocm" when IS_ROCM. Route /api/system through the same
   helper so the Studio UI reports the backend consistently across all
   endpoints.

4. studio/backend/tests/test_utils.py: update test_backend_matches_device
   to call _backend_label(get_device()) instead of raw get_device().value
   so the test matches the new contract and still passes on CUDA hosts.

Tests: 258 -> 261. New regression coverage:
- X08 main.py /api/system uses _backend_label
- X09 tokenizer multi-GPU guard has device_count() fallback
- X10 fresh-install bnb case block gated on SKIP_TORCH=false

* fix: prevent bitsandbytes from overwriting ROCm torch with CUDA wheels

During install, bitsandbytes was installed without --no-deps, causing
uv to resolve torch from PyPI (CUDA build) and silently overwrite the
ROCm wheels that were just installed in the previous step.

This happened in three places:
- install.sh: bitsandbytes install in both migrated and fresh paths
- install_python_stack.py: bitsandbytes install inside _ensure_rocm_torch()

Additionally, multiple install steps in install_python_stack.py (extras,
overrides, studio deps) can pull in CUDA torch via transitive
dependencies. A final _ensure_rocm_torch() call at the end of the
install sequence ensures ROCm torch is always in place at runtime.

All changes are gated behind ROCm-specific conditions and do not affect
NVIDIA, CPU-only, macOS, or Windows install paths.

Tested on AMD Instinct MI300X VF with ROCm 7.2.0 -- confirms
torch==2.10.0+rocm7.1 with HIP 7.1.25424 after install.

* fix: ROCm inference fallback -- skip Unsloth patching and bnb 4-bit on HIP

On AMD ROCm (HIP), two issues prevent the normal Unsloth inference path:

1. Unsloth's global monkey-patching of transformers model classes
   (LlamaRotaryEmbedding, attention modules) triggers
   _assert_async_cuda_kernel crashes on HIP during generation.
   Training uses different code paths and works fine.

2. bitsandbytes 4-bit matmul kernels also trigger HIP assertion
   failures on MI300X (CDNA3 / gfx942), even without Unsloth patching.

This commit adds a ROCm-specific inference fallback that:
- Skips importing Unsloth at module level (prevents global patching)
- Loads models in 16-bit with plain transformers + PEFT instead
- Resolves pre-quantized model names (e.g. "xxx-bnb-4bit" -> "xxx")
  since pre-quantized HF repos still trigger bnb codepaths
- Guards get_chat_template calls (unavailable without Unsloth import)
- Fixes max_seq_length=0 being passed to from_pretrained (GGUF
  semantics don't apply to transformers path)

The NVIDIA path is completely unchanged -- Unsloth import and
for_inference() optimization remain active. GGUF inference (via
llama-server/HIP) is unaffected since it never imports Python model
classes. AMD GPUs typically have large VRAM (e.g. 192GB on MI300X)
so 16-bit loading is practical for inference.

Tested on AMD Instinct MI300X VF (ROCm 7.2, HIP 7.1.25424):
- Simple generation: PASS
- Compare mode (base vs finetuned): PASS
- GGUF inference + tool calling: PASS (unaffected by this change)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: guard audio/vision inference on ROCm, remove unused import

- Add clear RuntimeError for audio/vision model inference on ROCm
  (these paths use Unsloth's FastModel/FastVisionModel which would
  crash on HIP; GGUF inference is the supported path on AMD)
- Remove unused `import os as _os` from the ROCm changes

* fix: amd-smi parsing for newer output format (gpu_data wrapper, mem_usage, temperature)

amd-smi on recent ROCm versions (7.x) wraps metric output in a
{"gpu_data": [...]} envelope instead of returning a raw list. This
caused get_primary_gpu_utilization() and get_visible_gpu_utilization()
to fail silently (returning available=False) because the GPU data
dict was never unwrapped.

Additionally:
- VRAM data moved from "vram" to "mem_usage" with "total_vram" /
  "used_vram" keys. Added fallback key lookup.
- Temperature "edge" sensor returns "N/A" on MI300X VF; the previous
  dict.get() chain returned the "N/A" string instead of falling
  through to "hotspot". Changed to a loop that checks each key until
  a parseable value is found.

Tested on AMD Instinct MI300X VF (ROCm 7.2, amd-smi 24.x):
- GPU utilization: 0% (idle), up to 100% during training
- Temperature: 40-44C (from hotspot sensor)
- VRAM: 0.28/191.69 GB (idle)
- Power: 158-211W draw

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Bug fix detecting radeon (#4940)

* Bug fix detecting radeon

* Expanding GPU target for gfx1100*

* Generalize gfx family-prefix filter to cover gfx10/gfx12 as well

rocminfo on ROCm 6.1+ emits LLVM generic-family ISA lines alongside the
specific GPU (e.g. gfx11-generic next to gfx1100). The outer grep captures
the bare family prefix from the generic line, and passing that to
-DGPU_TARGETS breaks the HIP build because clang only accepts specific
gfxNNN ids.

The previous filter only special-cased gfx11. Generalize it so any bare
2-digit family prefix (gfx10, gfx11, gfx12, ...) is dropped whenever a
specific sibling target is present in the same list. No real AMD GPU has
a 2-digit gfx id, so the filter can only ever drop family prefixes and
never a real target.

Covers the existing gfx11 cases unchanged, and extends the same fix to
gfx10-1-generic / gfx10-3-generic (RDNA1/2) and gfx12-generic (RDNA4),
which would otherwise hit the same build failure on newer rocminfo.

---------

Co-authored-by: Iswarya Alex <iswarya.alex@amd.com>
Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>

---------

Co-authored-by: Eda Z <eda.zhou@amd.com>
Co-authored-by: GoldenGrapeGentleman <yueyuan@amd.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: billishyahao <bill.he@amd.com>
Co-authored-by: Iswarya Alex <47045679+iswaryaalex@users.noreply.github.com>
Co-authored-by: Iswarya Alex <iswarya.alex@amd.com>
Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
2026-04-10 01:56:12 -07:00
Roland Tannous
f801e59c29
split venv_t5 into tiered 5.3.0/5.5.0 and fix trust_remote_code (#4878)
* split venv_t5 into venv_t5_530 and venv_t5_550 for tiered transformers 5.x support

* fix bfloat16 crash on T4 for FORCE_FLOAT32 models and disable trust_remote_code auto-enable for native t5 models

* revert FORCE_FLOAT32 dtype change

* restrict trust_remote_code auto-enable to Nemotron models only

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use config.json model_type for tier detection, add unsloth/nvidia namespace guard

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Revert "[pre-commit.ci] auto fixes from pre-commit.com hooks"

This reverts commit fb43d468e2.

* Revert "use config.json model_type for tier detection, add unsloth/nvidia namespace guard"

This reverts commit fc49ae2453.

* add unsloth/nvidia namespace guard to Nemotron trust_remote_code auto-enable

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* reorder tier checks: all substring matches before config.json fetches

* extract shared activate_transformers_for_subprocess into transformers_version.py

* narrow Nemotron trust_remote_code to nemotron_h/nemotron-3-nano, add to export worker

* clean venv_t5 dirs before re-install in setup.sh, clarify version alias comment

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* run venv_t5 migration outside deps fast-path gate in both setup scripts

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-07 20:05:01 +04:00
Daniel Han
ab65b47c73
Add tests for is_vision_model() caching behaviour (#4855)
* Add tests for is_vision_model() caching behaviour

* Fix review feedback: remove dead helper, fix exception test

- Remove unused _make_config() helper function (dead code)
- Fix test_exception_result_cached to actually exercise the exception path
  by mocking load_model_config to raise OSError instead of using
  side_effect=[False] which only tested normal False returns

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use strict mock specs so tests exercise intended detection paths

Use MagicMock(spec=[]) for all config mocks so hasattr() only returns
True for explicitly set attributes. Without this, MagicMock defaults
make all hasattr checks truthy, allowing tests to pass via unintended
detection paths (e.g. img_processor instead of vision_config).

---------

Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-06 06:41:40 -07:00
Daniel Han
4020a70a93
Add tests for cache case resolution (from PR #4822) (#4823)
Tests for resolve_cached_repo_id_case and get_model_config case
resolution, separated from the runtime changes in PR #4822.
2026-04-03 13:58:26 -07:00
Daniel Han
39fe23ded8
Tests for architecture-aware KV cache estimation (#4760)
* test: add 66 tests for architecture-aware KV cache estimation

Covers all 5 estimation paths (MLA, Hybrid Mamba, Sliding Window,
Standard GQA, Legacy), GGUF parser for 8 new metadata fields,
_can_estimate_kv gate conditions, quantization scaling, edge cases,
path priority ordering, and lifecycle (init/unload/reparse).

Zero external dependencies beyond pytest. No GPU or network required.
Cross-platform (Linux, macOS, Windows, WSL).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-01 06:13:37 -07:00
Daniel Han
653eb3819a
fix(studio): allow context length slider to reach model's native limit (#4746)
* fix(studio): allow context length slider to reach model's native limit

The context length slider was hard-capped to the VRAM-estimated maximum,
preventing users from requesting higher context even though the backend
already handles it safely (multi-GPU selection, --fit fallback). Expose
the model's native context length from GGUF metadata as a separate API
field and use it as the slider ceiling instead. Add an amber warning
when the selected context exceeds the estimated VRAM capacity.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Raise VRAM budget to 90% and add native_context_length tests

Increase the GPU memory utilization threshold from 70% to 90% across
_select_gpus and _fit_context_to_vram, allowing longer context lengths
before VRAM capping kicks in.

Add 33 tests for the native_context_length feature covering the backend
property, context value separation invariants, Pydantic models, route
completeness, edge cases, and cross-platform binary I/O.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-01 06:12:52 -07:00
Daniel Han
9a8b622306
Studio: simplify tool-call dedup and replace html2text with builtin converter (#4722)
* Simplify tool-call dedup: drop hashlib, inline helpers

The duplicate tool-call detector only compares calls within a single
request from the same JSON parser, so dict key order is guaranteed
identical for identical calls (Python 3.7+ insertion-ordered dicts).

- Replace hashlib.md5(json.dumps(...)) with name + str(args)
- Inline _tool_call_key, _is_duplicate_call, _record_tool_call
  since each was a one-liner used once
- Remove unused hashlib import

* Remove tool_calling_benchmark_results.md from repo

* Replace html2text with builtin HTML-to-Markdown converter

Drop the external html2text (GPL-3.0) dependency and its regex
fallback. Add _html_to_md.py (~190 lines, stdlib only) using
html.parser.HTMLParser that handles headings, links, bold/italic,
lists, tables, blockquotes, code blocks, and entity decoding.
Strips script/style/head tags entirely.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use json.dumps(sort_keys=True) for tool-call dedup key

str(dict) is sensitive to insertion order, so semantically identical
calls with different key ordering would bypass duplicate detection.
Switch to json.dumps with sort_keys=True for a canonical representation.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Revert dedup key to str(arguments)

json.dumps(sort_keys=True) is unnecessary here -- the arguments dict
always comes from the same JSON parser within a single request, so
key insertion order is deterministic (Python 3.7+).  str() is faster
and sufficient for consecutive-call dedup.

* Address review comments on _html_to_md.py

- Remove "hr" from _BLOCK_TAGS so the dedicated hr handler is reachable
- Prefix all newlines with ">" inside blockquotes (multi-line support)
- Emit full ![alt](url) for images instead of alt text only
- Replace newlines with spaces inside table cells
- Track header cells per-row (_row_has_th) instead of last-cell-only
- Strip trailing tabs in addition to spaces in cleanup regex

* Fix blockquote rendering, truncated-HTML buffer flush, and dedup key canonicalization

_html_to_md.py:
- Rewrite blockquote handling with stack-based buffer approach so nested
  blockquotes, pre blocks inside blockquotes, and multi-paragraph quotes
  all render correctly with proper "> " prefix on every line.
- Add flush_pending() to recover content from truncated HTML where closing
  tags are missing (common when _fetch_page_text caps the download size).
  Flushes open <a>, <td>, <pre>, and blockquote buffers.
- Skip <img> tags to match prior html2text ignore_images=True behavior
  and avoid data-URI amplification consuming the output budget.
- Collapse all whitespace (including newlines) in non-pre content per
  standard HTML whitespace rules: \s+ -> single space.
- Escape pipe characters in table cell content to prevent column breakage.
- Emit separator row after the first row for tables without <th> headers.
- Guard against IndexError on _ol_counter for orphan <li> elements.
- Normalize CRLF line endings before parsing.

llama_cpp.py:
- Restore canonical dedup key with json.dumps(sort_keys=True) so that
  semantically identical tool calls with different JSON key order are
  correctly detected as duplicates.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix table optional end tags, inline code whitespace, and link text normalization

_html_to_md.py:
- Extract _finish_cell() and _finish_row() helpers to handle HTML tables
  that omit optional </td>, </th>, or </tr> end tags. This is valid HTML
  and common on real web pages -- previously the parser would silently
  drop earlier cells and entire rows.
- Call _finish_cell()/_finish_row() from handle_starttag for <tr>/<td>/<th>,
  handle_endtag for </tr>/<td>/<th>/<table>, and flush_pending() so all
  three paths (normal close, implicit close, truncated HTML) use the same
  row-finalization logic including header separator emission.
- Add _in_inline_code flag so handle_data() preserves literal whitespace
  inside <code> spans instead of collapsing it. Source like
  <code>pip  install   unsloth</code> now correctly renders as
  `pip  install   unsloth` rather than `pip install unsloth`.
- Extract _finish_link() helper that normalizes accumulated link text with
  \s+ -> single space before building the Markdown link. Prevents block-
  level content inside <a> tags (e.g. <a><div>one</div><div>two</div></a>)
  from producing multiline [one\n\ntwo](href) link labels.
- Empty blockquotes now produce no output instead of a stray ">".
- Remove unused _bq_depth field (all routing uses _bq_stack).
- Flush open cells and rows in handle_endtag("table") for robustness.

* Support <ol start=N>, <dl>/<dt>/<dd>, and preserve code block whitespace

_html_to_md.py:
- Honor <ol start="N"> attribute so ordered lists preserve their original
  numbering instead of always restarting from 1. Important for docs/tutorials
  that continue numbering across sections.
- Add dl, dt, dd to _BLOCK_TAGS so definition lists (common on MDN, Python
  docs, Django docs) produce separated text instead of concatenated blobs.
- Rewrite _cleanup() to be fence-aware: content inside fenced code blocks
  is now preserved verbatim (intentional blank lines in <pre> content are
  no longer collapsed). Outside code blocks, blank runs are limited to one
  and trailing whitespace is stripped.
- Fix _prefix_blockquote() to strip trailing whitespace before collapsing
  blank lines, preventing the "\n\n \n\n" pattern from sneaking through.

* Suppress whitespace-only text nodes between table structural elements

Indented HTML tables (nearly all real-world pages) produce whitespace
text nodes between <table>, <tr>, </tr> etc. that land in the output
as leading spaces before table rows, breaking Markdown table alignment.

Skip whitespace-only text nodes when inside a table but not inside a
cell, so indentation from source HTML does not leak into the output.

* Revert dedup key to str(arguments) with explanatory comment

json.dumps(sort_keys=True) is unnecessary overhead here: arguments
always comes from json.loads on model output within a single request,
so dict insertion order is deterministic in Python 3.7+. A repeated
call from the model produces the same JSON, which parses to the same
dict repr. str() avoids re-serialization on every tool call.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-31 06:15:18 -07:00
Daniel Han
e159b93b97
studio: improve GGUF tool calling accuracy and reliability (#4700)
* studio: improve GGUF tool calling accuracy and reliability

- Add URL fetching to web_search tool so models can read full page
  content instead of only getting search snippets. Uses html2text for
  clean markdown conversion with regex fallback.
- Inject current date and behavioral guidance (URL fetch workflow,
  no repeated queries, use code for data processing) into the
  tool-use system prompt.
- Append error recovery nudge to tool results that indicate failure,
  helping small models avoid looping on the same broken call.
- Strip leaked <tool_call> XML from assistant messages in conversation
  history and from the outgoing SSE stream.
- Raise default max tool iterations from 10 to 25 across backend,
  model schema, and frontend defaults.
- Increase _MAX_PAGE_CHARS from 4k to 16k so fetched pages contain
  enough content for the model to extract useful information.
- Add "IMPORTANT: These are only short snippets" hint to search
  results so models know to fetch full pages when needed.

Tested with Qwen3.5-4B-GGUF (UD-Q4_K_XL), 10 runs before/after:
- XML leaks in responses: 10/10 -> 0/10
- URL fetch usage: 0 -> 4/10 runs
- Runs producing actual correct answers: 0/10 -> 2/10
- Average tool calls per query: 5.5 -> 3.8 (more efficient)
- Average response time: 12.3s -> 9.8s

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add tool calling benchmark results across model sizes and quants

Tested 16 configurations (4 models x 2 quants x 2 KV cache types)
with 10 runs each on NVIDIA B200.

Best config: 27B UD-Q4_K_XL + bf16 KV -- 6/10 runs found all 4
correct songs, 0 XML leaks, 131s average response time.

* Add duplicate tool-call detection and final-answer synthesis

When the model repeats the exact same tool call (same name + arguments)
twice in a row, skip execution and return a redirect message telling it
to try a different approach. This prevents the 8x-repeated-query loops
observed on 27B and 35B models.

When the tool iteration cap (25) is reached, inject a "provide your
final answer now" message before the final streaming pass. This lets
the model synthesize a useful answer from everything it gathered
instead of being silently cut off.

Tested on Qwen3.5-27B UD-Q4_K_XL (10 runs):
- Repeated query runs: 4/10 -> 2/10
- Cap hits: 1/10 -> 0/10
- All 4/4 accuracy: 5/10 -> 7/10

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix CodeQL alert: handle whitespace in script/style closing tags

The regex fallback for HTML stripping did not match closing tags
with whitespace before the angle bracket (e.g. </script >).
Use \s* before > in both script and style patterns.

* Address reviewer findings: SSRF, timeout crash, XML regex, dedup

- SSRF: resolve hostname via getaddrinfo and reject private, loopback,
  link-local, multicast, and reserved addresses before fetching
- Timeout: handle timeout=None (unlimited mode) in URL fetch path
  by defaulting to 60s instead of crashing on min(None, 60)
- Download cap: read at most max_chars*4+1 bytes instead of the
  full response body before truncating
- XML regex: match both <tool_call> and <function=...> markup in
  the history/stream cleanup (inference.py)
- CodeQL: use [^>]* in closing script/style tags to handle any
  whitespace or attributes before >
- Dedup: track whether each tool call failed so retries after
  transient errors are allowed; only block consecutive identical
  calls that both succeeded
- Final-answer synthesis: guard on max_tool_iterations > 0 so
  callers who disable tools do not get a false "used all calls" turn

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix redirect SSRF, SSE streaming regression, dedup off-by-one

- SSRF redirect bypass: disable auto-redirect in urllib, manually
  follow up to 5 hops with host validation at each step. Prevents
  public URLs from redirecting to loopback/private targets.
- SSE streaming: track prev_text on the raw cumulative and strip
  XML from the delta only, so completed tool_call tags do not cause
  the cumulative to shrink and drop trailing real text.
- Dedup off-by-one: check the immediately previous call (window=1)
  instead of requiring 2 matching history entries, so the second
  identical successful call is blocked rather than the third.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix redirect HTTPError handling and tighten error prefixes

- Redirect fix: urllib raises HTTPError (not a normal response) when
  the redirect handler returns None. Catch HTTPError for 3xx codes
  and extract the Location header from the exception object.
- Error prefixes: remove overly broad "No " prefix that matched
  "No results found." (a valid empty-search outcome, not an error).
  Replace with specific prefixes like "Blocked:", "No query provided",
  "Failed to resolve". This ensures empty search results are correctly
  classified as non-errors for duplicate-call tracking.

* Fix SSE cross-chunk XML leaks, cleanup review findings

- SSE streaming: sanitize the full cumulative text before diffing
  against the previous sanitized snapshot, so XML tags that span
  chunk boundaries are stripped correctly. The previous delta-based
  approach leaked split tags.
- DRAINING fallback: use _strip_tool_markup() helper instead of a
  manual regex that only handled <tool_call> but not <function=...>.
- Move hashlib import, _TOOL_XML_RE compile, and datetime import to
  module level per style guide.
- Remove unused _hit_tool_cap variable.

* Fix DNS rebinding, charset detection, HTTPError handling, dedup double-record

- DNS rebinding: resolve hostname once via getaddrinfo, pin the
  returned IP, rewrite the URL to connect to the pinned IP with
  a Host header. Each redirect hop re-resolves and re-validates.
  Closes the TOCTOU window between validation and connection.
- Charset: use resp.headers.get_content_charset() instead of
  hardcoding utf-8, so pages with other encodings decode correctly.
- HTTPError: return descriptive "HTTP {code} {reason}" instead of
  re-raising into a generic "Search failed" message.
- Dedup: remove redundant _record_tool_call in the duplicate branch;
  the single call at the end of the loop handles all cases.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-03-31 03:06:44 -07:00
Datta Nimmaturi
3b5a49776b
[studio] multi gpu: revert to balanced for inference. (#4698)
* Revert to balanced for inference

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused for_inference parameter from get_device_map

Since inference and training both use "balanced" now, the for_inference
flag is dead code. Remove it from the function signature, the call site
in inference.py, and simplify the tests accordingly.

* Remove redundant TestDeviceMapForInference test class

TestGpuAutoSelection already covers the same multi-gpu and single-gpu
device_map assertions. The TestDeviceMapForInference class was left
over from when for_inference had distinct behavior.

* Remove redundant test_get_device_map_multi_gpu_uses_balanced

Its assertions ([0,1] -> balanced, [0] -> sequential) are already
covered by test_get_device_map_uses_explicit_gpu_selection.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-31 01:24:41 -07:00
Datta Nimmaturi
9311df2b29
[Studio] multi gpu finetuning/inference via "balanced_low0/sequential" device_map (#4602)
* [WIP] balanced device map for studio

* gpus as a request parameter

* API for multi GPU stuff

* return multi gpu util in new API

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use balanced_low0 instead of balanced

* Use balanced_low0 instead of balanced

* Fix device_map typo, UUID parsing crash, set() filter bug, and broken tests

- balanced_low0 -> balanced_low_0 (transformers/accelerate rejects the old string)
- get_parent_visible_gpu_ids() now handles UUID/MIG CUDA_VISIBLE_DEVICES
  gracefully instead of crashing on int() parse
- _get_backend_visible_gpu_info() set() or None bug: empty set is falsy so
  CUDA_VISIBLE_DEVICES=-1 would disable filtering and report all GPUs
- test_gpu_selection.py: add missing get_visible_gpu_utilization import and
  add required job_id arg to start_training() calls

* Smart GPU determinism using estimates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* disallow gpu selection for gguf for now

* cleanup

* Slightly larger baseline

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Treat empty list as auto

* Verbose logging/debug

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup and revert unnecessary deletions

* Cleanup excessive logs and guard against disk/cpu offload

* auth for visibility API. cleanup redundant imports. Adjust QLoRA estimate

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* support for non cuda gpus

* Fix multi-GPU auto-selection memory accounting

The multi_gpu_factor was applied uniformly to all GPUs including the
first one, which unfairly penalizes single-GPU capacity when
transitioning to multi-GPU. This created a discontinuity where a model
that barely fits 1 GPU would suddenly require 2 GPUs because the first
GPU's free memory was discounted by 20%.

Now the first GPU keeps its full free memory, and only additional GPUs
have an overhead factor (0.85) applied to account for inter-GPU
communication and sharding overhead. This gives more accurate
auto-selection and avoids unnecessary multi-GPU for models that
comfortably fit on one device.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add sandbox tests for multi-GPU selection logic

24 tests covering model size estimation, memory requirements, automatic
GPU selection, device map generation, GPU ID validation, and multi-GPU
overhead accounting. All tests use mocks so they run without GPUs on
Linux, macOS, and Windows.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix reviewer findings: 4bit inference estimate, fallback, GGUF gpu_ids, retry

1. 4-bit inference now uses reduced memory estimate (model_size/3 + buffer)
   instead of the FP16 1.3x multiplier. This prevents over-sharding
   quantized models across unnecessary GPUs.

2. When model size estimation fails, auto_select_gpu_ids now falls back to
   all visible GPUs instead of returning None (which could default to
   single-GPU loading for an unknown-size model).

3. GGUF inference route now treats gpu_ids=[] as auto-selection (same as
   None) instead of rejecting it as an unsupported explicit request.

4. Training retry path for "could not get source code" now preserves the
   gpu_ids parameter so the retry lands on the same GPUs.

5. Updated sandbox tests to cover the new 4-bit inference estimate branch.

* Remove accidentally added unsloth-zoo submodule

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix UUID/MIG visibility and update test expectations

1. nvidia.py: When CUDA_VISIBLE_DEVICES uses UUID/MIG tokens, the
   visibility APIs now return "unresolved" with empty device lists instead
   of exposing all physical GPUs. This prevents the UI from showing GPUs
   that the backend process cannot actually use.

2. test_gpu_selection.py: Updated test expectations to match the new
   multi-GPU overhead accounting (first GPU at full capacity, 0.85x for
   additional GPUs) and 4-bit inference memory estimation formula.
   All 60 tests now pass.

* Add CPU/disk offload guard to audio inference path

The audio model loading branch returned before the common
get_offloaded_device_map_entries() check, so audio models loaded with a
multi-GPU device_map that spilled layers to CPU/disk would be accepted
instead of rejected. Now audio loads also verify no modules are offloaded.

* Improve VRAM requirement estimates

* Replace balanced_low_0 with balanced

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refine calculations for slightly easier nums

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* adjust estimates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use nums instead of obj to avoid seralisation error

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Harden nvidia-smi parsing and fix fallback GPU list

1. nvidia.py: Wrap int() casts for GPU index and memory in try/except
   so MIG slices, N/A values, or unexpected nvidia-smi output skip the
   unparseable row instead of aborting the entire GPU list.

2. nvidia.py: Handle GPU names containing commas by using the last
   field as memory instead of a fixed positional index.

3. hardware.py: fallback_all now uses gpu_candidates (GPUs with verified
   VRAM data) instead of raw devices list, which could include GPUs
   with null VRAM that were excluded from the ranking.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* consolidate raise_if_offload

* Improve MoE support. Guard against nvidia-smi failures

* Improve MoE support. Guard against nvidia-smi failures

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix shared-expert LoRA undercount, torch VRAM fallback, and apply_gpu_ids edge case

1. vram_estimation.py: compute_lora_params now includes shared experts
   (n_shared_experts) alongside routed experts when computing MoE LoRA
   adapter parameters. Previously only n_experts were counted, causing
   the estimator to undercount adapter, optimizer, and gradient memory
   for DeepSeek/GLM-style models with shared experts.

2. hardware.py: _torch_get_per_device_info now uses mem_get_info (which
   reports system-wide VRAM usage) instead of memory_allocated (which
   only reports this process's PyTorch allocations). This prevents
   auto-selection from treating a GPU as mostly free when another
   process is consuming VRAM. Falls back to memory_allocated when
   mem_get_info is unavailable.

3. hardware.py: apply_gpu_ids([]) now returns early instead of setting
   CUDA_VISIBLE_DEVICES="" which would disable CUDA entirely. Empty
   list inherits the parent visibility, same as None.

4. hardware.py: Upgraded fallback_all GPU selection log from debug to
   warning so operators are notified when the model likely will not fit
   in available VRAM.

* Guard nvidia-smi subprocess calls against OSError and TimeoutExpired

get_visible_gpu_utilization and get_backend_visible_gpu_info now catch
OSError (nvidia-smi not found) and TimeoutExpired internally instead
of relying on callers to wrap every invocation. Returns the standard
available=False sentinel on failure so the torch-based fallback in
hardware.py can take over.

* Guard get_primary_gpu_utilization and reset GPU caches between tests

1. nvidia.py: get_primary_gpu_utilization now catches OSError and
   TimeoutExpired internally, matching the pattern already used in
   get_visible_gpu_utilization and get_backend_visible_gpu_info. All
   three nvidia-smi callers are now self-contained.

2. test_gpu_selection.py: Added _GpuCacheResetMixin that resets the
   module-level _physical_gpu_count and _visible_gpu_count caches in
   tearDown. Applied to all test classes that exercise GPU selection,
   device map, or visibility functions. This prevents stale cache
   values from leaking between tests and causing flaky results on
   machines with real GPUs.

* Fix nvidia-smi fallback regression and physical GPU count validation

1. hardware.py: get_gpu_utilization, get_visible_gpu_utilization, and
   get_backend_visible_gpu_info now check result.get("available") before
   returning the nvidia-smi result. When nvidia-smi is unavailable or
   returns no data (e.g., containers without nvidia-smi, UUID/MIG masks),
   the functions fall through to the torch-based fallback instead of
   returning an empty result. This fixes a regression where the internal
   exception handling in nvidia.py prevented the caller's except block
   from triggering the fallback.

2. hardware.py: resolve_requested_gpu_ids now separates negative-ID
   validation from physical upper-bound validation. The physical count
   check is only enforced when it is plausibly a true physical count
   (i.e., higher than the largest parent-visible ID), since
   torch.cuda.device_count() under CUDA_VISIBLE_DEVICES returns the
   visible count, not the physical total. The parent-visible-set check
   remains authoritative in all cases. This prevents valid physical IDs
   like [2, 3] from being rejected as "out of range" when nvidia-smi is
   unavailable and CUDA_VISIBLE_DEVICES="2,3" makes torch report only
   2 devices.

* Fix UUID/MIG torch fallback to enumerate devices by ordinal

When CUDA_VISIBLE_DEVICES uses UUID or MIG identifiers,
get_parent_visible_gpu_ids() returns [] because the tokens are
non-numeric. The torch fallback in get_visible_gpu_utilization() and
get_backend_visible_gpu_info() previously passed that empty list to
_torch_get_per_device_info(), getting nothing back.

Now both functions detect the empty-list case and fall back to
enumerating torch-visible ordinals (0..device_count-1) with
index_kind="relative". This means the UI and auto-selection still
see real device data in Kubernetes, MIG, and Slurm-style UUID
environments where nvidia-smi output cannot be mapped to physical
indices.

Updated test_uuid_parent_visibility to verify the new torch fallback
path returns available=True with relative ordinals.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add type hint for gpu_ids parameter in InferenceOrchestrator.load_model

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-30 02:33:15 -07:00