Commit graph

1,853 commits

Author SHA1 Message Date
pre-commit-ci[bot]
91780a250e [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-07-01 11:26:13 +00:00
Daniel Han
172a4eb642 Merge branch 'diffusion-phase4-native' into diffusion-phase6-features
# Conflicts:
#	studio/backend/core/inference/sd_cpp_engine.py
#	studio/backend/tests/test_sd_cpp_engine.py
2026-07-01 11:24:14 +00:00
Daniel Han
48628252bd Merge remote-tracking branch 'origin/image-generation' into diffusion-phase4-native
# Conflicts:
#	scripts/diffusion_bench.py
#	scripts/diffusion_quality.py
#	studio/backend/core/inference/diffusion.py
#	studio/backend/core/inference/diffusion_device.py
#	studio/backend/core/inference/diffusion_families.py
#	studio/backend/core/inference/diffusion_memory.py
#	studio/backend/core/inference/diffusion_precision.py
#	studio/backend/core/inference/diffusion_speed.py
#	studio/backend/models/inference.py
#	studio/backend/routes/inference.py
#	studio/backend/tests/test_diffusion_backend.py
#	studio/backend/tests/test_diffusion_device.py
#	studio/backend/tests/test_diffusion_memory.py
#	studio/backend/tests/test_diffusion_precision.py
#	studio/backend/tests/test_diffusion_speed.py
2026-07-01 11:19:15 +00:00
Daniel Han
2cae1c1e7c Merge remote-tracking branch 'origin/main' into image-generation
# Conflicts:
#	studio/backend/routes/models.py
#	studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
2026-07-01 10:47:21 +00:00
Daniel Han
8cc05ac89c
Reduce comments across recent fixes (#6776)
Condense the verbose comments and docstrings added by the recent
chat template, GPT-OSS detection, PEFT tensor-parallel, and Studio
inference proxy fixes. Comments and whitespace only; no code changes.
2026-06-30 23:13:36 -07:00
oobabooga
8b69f22809 Fix fp8 text-encoder quant crashing generation on T5 and tied-embedding encoders 2026-06-30 23:22:13 -03:00
Anish Umale
d0f8d40c36
studio: allow updating HF models through UI (#5388)
* add models for /update endpoint

* add logic for identifying out of date hf models

* add endpoint for updating hf models

* add relevant field to GgufVariantDetail

* make exception handling better

* add update_available flag for cached_models, and moved /update endpoint from inference -> models

* hook up /update endpoint on the frontend

* implement update scenarios for the model picker

* fix bug where downloaded flag for an older revision was being wrongly set to false

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

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

* fix import and make hf calls async

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

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

* remove has_vision from UpdateRequest

* fix ci

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

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

* clear cancel event before updating gguf variant

* set _cancel_event back if it was set initially

* add hf_token to get_paths_info

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

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

* studio: harden model update endpoint and update checks

- update_hf_model: pass snapshot_download local_dir (local_path is not a
  valid kwarg and 500s when updating bicodec audio models)
- get_gguf_variants: wrap the remote update check so a network, rate-limit,
  gated, or offline failure degrades to "no update info" instead of failing
  the whole variant listing, matching list_cached_models
- add regression tests for both paths

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

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

* Studio: HF model update detection and Update action for cached models

Surface an "Update available" cue and a managed Update action for cached
on-device models. /api/hub/update-status compares each cached main GGUF
file's local blobs against the remote main revision using set membership
across all cached revisions, so a repo that was already updated (and still
holds the old snapshot alongside the new one) is not falsely flagged.

The Update action re-downloads through the download manager so it shows in
the Downloads panel with progress and cancel. The frontend wires the Update
button into the GGUF, on-device, and model-selector cards and keeps the
quant label fully visible when the action buttons crowd the row.

Adds regression tests for the multi-revision update check.

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

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

* Studio: accept force_download kwarg in hf_xet_fallback test double

The download seam now passes force_download to the attempt callable; the _FakeAttempt mock did not accept it, failing 6 tests with TypeError. Add the keyword (default False) so the scripted-results double matches the seam.

* Fix Studio model update regressions

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

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

* Address Studio update review feedback

* Address Studio update edge cases

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

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

* Share GGUF update status helper

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

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

* Fix GGUF update detection and cache cleanup

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

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

* Fix cached GGUF update badges

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: shimmyshimmer <107991372+shimmyshimmer@users.noreply.github.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-07-01 01:54:57 +03:00
pre-commit-ci[bot]
53077b5ae3 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-30 22:34:29 +00:00
Daniel Han
6ae6fb1c45
Studio diffusion (Phase 2): memory planner, streamed offload, fp8 TE, speed layer, quality harness (#6675)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-30 19:30:57 -03:00
Daniel Han
ecf028780d
Studio diffusion (Phase 1): cross-platform device policy, fp16 guard, lock split, validate-before-evict (#6670)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-30 16:33:47 -03:00
pre-commit-ci[bot]
43f5b619be [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-30 19:04:34 +00:00
YuzukoUnderson
d915a139eb
feat(i18n): add Japanese locale support for the Studio UI catalog (#6705)
* i18n: register Japanese language support in messages

* i18n: add Japanese locale support

* Update studio/frontend/src/i18n/locales/ja.ts

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update studio/frontend/src/i18n/locales/ja.ts

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* test(i18n): add ja locale to parity check

* i18n: fill remaining missing keys for Japanese locale

* i18n: fix terminal string localization in ja locale

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-30 19:37:04 +01:00
Tai An
7337729e57
fix(studio/llama_cpp): disable trust_env on the loopback health probe (#6750) (#6752)
* fix(studio/llama_cpp): disable trust_env on the loopback health probe

_wait_for_health() polls http://127.0.0.1:<port>/health with the default
httpx trust_env=True, so an ambient HTTP(S)_PROXY in the environment is
applied to the loopback request. A proxy that returns 503 for 127.0.0.1
makes every probe fail, so the loop runs until timeout and Studio load
hangs (trust_env=False returns 200 immediately).

Pass trust_env=False so the local readiness probe never goes through a
proxy. This mirrors the existing trust_env=False handling in the sibling
llama_http / external_provider HTTP clients.

* test(offline_gguf_cache): accept trust_env kwarg in fake_get mock

_wait_for_health now calls httpx.get(..., trust_env=False); update the retry test's fake_get to accept the kwarg so it doesn't raise TypeError.

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

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

* fix(studio/llama_cpp): bypass proxies for loopback clients

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

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

* fix(studio/routes): bypass proxies for llama streams

* [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: wasimysaid <wasimysdev@gmail.com>
2026-06-30 19:09:26 +02:00
Daniel Han
b72a8c4263
studio: explicit Cloudflare tunnel notice and public-exposure warning at startup (#6515)
* studio: announce Cloudflare tunnel state and warn about public exposure on startup

The startup banner only printed a line when a tunnel URL was up, so a plain
`unsloth studio -H 0.0.0.0` launch silently created a public trycloudflare.com
URL with no indication that Studio had become reachable from the internet. The
only hint at the tunnel was the CLI help, shown when an invalid command was typed.

Make the banner always state the tunnel state for wildcard binds:
- ON: the public URL plus a warning that anyone with it can reach Studio from
  outside the network, and that --no-cloudflare keeps it local-only.
- FAILED: requested but did not start (local network only).
- OFF: --no-cloudflare was passed (local network only).
Secure mode keeps its existing wording (the authenticated tunnel is intended and
--no-cloudflare is not valid there). Clarify the --cloudflare help text in both
the argparse and typer definitions. Default behavior is unchanged.

Also surface the state on the `unsloth studio run` banner, which runs the server
with silent=True and prints its own banner: it now calls _print_cloudflare_line
too, so the ON/OFF/FAILED notice and public-exposure warning are no longer
skipped on that path (previously it only echoed the URL when a tunnel was up).

For the OFF and FAILED notices, do not claim "local network only" when the
reachability probe just confirmed the raw port is reachable from the public
internet: --no-cloudflare and a failed tunnel disable only the Cloudflare link,
not the wildcard bind, so the message is reworded to flag the public raw port.

* Fix/adjust Cloudflare banner warnings for PR #6515

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

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

* Fix/adjust Cloudflare banner comments for PR #6515

* Fix/adjust IPv6 Cloudflare tunnel gate for PR #6515

* Fix/adjust Cloudflare review comments for PR #6515

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

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

* Fix silent run Cloudflare notice

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

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

---------

Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-30 17:47:48 +02:00
Michael Han
0a3e5a3172
Studio: quick eject from the model selector (#6654)
* Studio: quick eject from the model selector

Add a one-click eject shortcut to the loaded-model pill so users do not
have to open the picker to unload a model.

- The loaded-status indicator shows a green checkmark at rest and swaps to
  a red eject icon on pill hover, with an "Eject model" tooltip. Clicking
  it ejects without opening the picker.
- On Device tab now uses the placeholder "Search local models" instead of
  "Search Unsloth models".
- The picker's "Eject model" button uses medium font weight.

* Studio: drop unused group/eject marker class on the eject control

* Studio: make the inline eject control valid HTML

The eject shortcut was a focusable span (role/tabIndex) nested inside the
trigger button. A button's content model forbids focusable descendants, so
make it a plain decorative span (aria-hidden, no role/tabIndex) that keeps
the mouse shortcut. Keyboard and screen-reader users eject via the picker's
"Eject model" button.

* Studio: disable the inline eject shortcut on touch devices

On touch (no hover) the red eject icon and title tooltip never reveal, so
tapping the loaded pill could unload the model with no visible affordance.
Add [@media(hover:none)]:pointer-events-none so taps fall through to the
trigger and open the picker; touch users eject from the picker instead.

---------

Co-authored-by: shimmyshimmer <shimmyshimmer@users.noreply.github.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
2026-06-30 16:52:03 +02:00
Daniel Han
e8945cab46
Whole-document context for RAG chat attachments (#6693)
* Add whole-document context mode to RAG chat attachments

Thread-attached files are injected in full when they fit a token budget,
instead of only top-K retrieved chunks, so the model reads the entire file
for summarize/reason-over-document requests. Oversized files fall back to
top-K retrieval so the context window is never blown. KB and project
corpora are unchanged (still retrieval).

- core/rag/store.py: all_chunks_for_scope returns every completed-document
  chunk for a scope, ordered document-then-index, joined with filename.
- core/rag/tool.py: whole_document_context renders the chunks as the same
  <chunk> blocks + citation source-map retrieval produces, returns None
  when empty or over budget.
- core/inference/tools.py: build_rag_autoinject tries whole-document first
  for thread scopes, falls through to search_for_autoinject otherwise.
- core/rag/config.py: THREAD_WHOLE_DOC + WHOLE_DOC_MAX_TOKENS (env-tunable).
- tests/test_rag_whole_document.py: store ordering, whole-doc render +
  budget cutoff, auto-inject whole-doc vs top-K fallback, KB never whole-doc.

* Add scanned-PDF OCR fallback to RAG ingestion

A PDF page with no extractable text layer (a scanned or image-only page)
previously ingested as empty, so image PDFs were invisible to retrieval and
whole-document context. Such pages are now rendered and transcribed by the
loaded vision model during ingestion, so they become searchable and readable
like any other page. This restores OCR for the RAG document flow without a
separate extraction pipeline.

- core/rag/parsers.py: render_pdf_pages renders whole pages (1-based) to PNG.
- core/rag/captioner.py: factor the shared vision call into _vision_complete;
  add _ocr_one + ocr_pages (transcribe rendered pages, OCR_MAX_PAGES bound).
- core/rag/ingestion.py: _ocr_scanned_pages runs right after parse, replacing
  text on near-empty PDF pages. No-op when OCR is off, no page is scanned, or
  no vision model is loaded (degrades like figure captioning).
- core/rag/config.py: OCR_SCANNED, OCR_MIN_CHARS, OCR_MAX_PAGES, OCR_DPI,
  OCR_TIMEOUT_S, OCR_MAX_TOKENS (env-tunable).
- tests/test_rag_ocr_fallback.py: page render, ocr_pages gating + cap, scanned
  PDF end-to-end OCR into chunks + whole-doc, born-digital skips OCR, disabled
  leaves the page empty.

* Broaden OCR prompt to figures/tables and guard against repetition runaway

The OCR prompt now asks the vision model to also transcribe text inside figures,
diagrams, charts and tables, so labels and table cells on scanned pages are
indexed rather than skipped. Verified on real documents that this does not
regress plain-text transcription.

Some vision models loop on sparse images (e.g. a title-only cover) and emit the
same line hundreds of times. _collapse_runaway caps any run of identical
consecutive lines so a pathological page cannot flood the index; legitimate
short repeats (a label appearing a few times) survive. Applied in ocr_pages.

* Restrict whole-document injection to thread attachments only

whole_document_context resolved the combined project+thread scope, so a project
chat (the frontend sends both thread_id and project_id) injected the entire
project corpus in full, contradicting the design that project and KB corpora stay
retrieval-only. A large project corpus could also push the total over budget and
drop a small thread attachment back to top-K.

Resolve the thread scope alone in whole_document_context, and in
build_rag_autoinject only enter whole-doc mode when a thread attachment is present
and no KB is selected (a KB pick is exclusive: search that corpus). Project
sources and KBs keep top-K retrieval. Adds regression tests for the mixed
project+thread payload, the budget isolation, and KB precedence.

* Address review: keep project retrieval, harden budget + OCR guards

Follow-up to the 8-reviewer pass on the whole-document + OCR work.

- Preserve project grounding in project chats. The thread-scope-only fix made
  whole-doc exclusive of retrieval, so a thread attachment silently dropped the
  project corpus for that turn. build_rag_autoinject now whole-docs the thread
  attachment AND retrieves the project sources top-K, merged under one citation
  numbering via tool.render_sources. KB selection stays exclusive.
- Budget: a NULL/zero token_count no longer bypasses the cap (length-based
  fallback in _row_token_count), so a malformed huge doc can't inject in full.
- OCR runaway guard: _collapse_runaway now also caps each distinct line at a
  generous total across the page (not just consecutive), bounding the
  interleaved/alternating loops weak models emit; blank-line floods collapse too.
- OCR: warn when a scanned PDF exceeds OCR_MAX_PAGES (pages past the cap stay
  untranscribed) instead of silently dropping them.
- Document the known limits: OCR'd pages have no PDF highlight regions; vision
  models need a micro-batch >= image tokens (Gemma-family) or the server aborts.
- Tests for project-retrieval composition, NULL-token budget, and interleaved
  runaway; drop the now-superseded exclude-project test.

* Add OCR toggle to RAG retrieval settings

Make scanned-PDF OCR user-controllable per upload instead of only via the
RAG_OCR_SCANNED config default. The retrieval settings panel gains an OCR
scanned pages switch (persisted in localStorage, on by default); the chosen
value is read fresh at upload time and sent with each document upload.

Backend: the three upload routes accept an optional ocr form field and pass it
through start_ingestion to _ocr_scanned_pages, which now treats None as use the
config default and an explicit bool as an override. The on/off policy lives only
in _ocr_scanned_pages now, so ocr_pages no longer re-checks the config (that
double gate would have blocked a per-upload ocr=True while the default was off).

Tests cover both override directions (force on while config off, force off while
config on).

* Add "Describe figures & charts" toggle with chart-aware captions

Surface RAG figure captioning as a user control and make it actually useful for
graphs and plots. The figure detection already clustered vector drawings and
raster images into regions and rendered them, but captioning was off by default,
had no UI, and used a thin generic prompt.

Accuracy: the caption prompt now asks for chart type, axis titles and units,
legend or series, salient trends and readable values, and table columns, while
forbidding invented numbers. The token budget is configurable (CAPTION_MAX_TOKENS)
and captions pass through the same runaway guard as OCR so a looping vision model
cannot flood the index.

Control: a per-upload caption override threads from the three upload routes through
start_ingestion and _run, with the on/off policy single-sourced in _run (caption
self-gating removed from caption_images, mirroring the OCR change) so a force-on
override works when the config default is off. The frontend adds a "Describe
figures & charts" switch in the retrieval settings, persisted in localStorage and
sent with each upload. Default on; it is a no-op without a vision model and bounded
to CAPTION_MAX_IMAGES figures per document.

Tests cover the new caption_images contract, the runaway guard on captions, the
chart-aware prompt and token budget (and that OCR keeps its own prompt and budget),
and both override directions end to end through ingestion.

* Generalize figure understanding: transcribe-first prompt + high-DPI tiling

Make figure/chart description work across any visual and any model strength, not
just a strong VLM on simple figures. Two changes, validated by a recall benchmark
on authoritative documents (ResNet/Attention papers, USDA, UN UDHR).

1. Transcribe-first caption prompt. The caption now asks the model to transcribe
   every visible label verbatim (titles, axis labels and units, legends, every
   box/node/arrow label, table cells, equations) and then add a one-line summary,
   instead of only describing the figure. Transcription is the most model-robust
   visual task, so weak models that cannot reason about a chart still recover its
   labels.

2. High-DPI tiling of figure pages. Figure-bearing pages are rendered as an
   overlapping grid of high-DPI tiles (plus a full-page pass for context); each
   tile is transcribed, then merged and de-duplicated. This keeps small diagram
   labels legible and covers every sub-figure without relying on exact region
   detection, which previously missed sub-figures and small labels.

Supporting changes: figure render DPI 130 -> 200 with a clip margin so edge labels
are not lost; vision calls are deterministic (temperature 0) so transcription does
not randomly drop labels; the repetition guard now applies to captions too. New
config knobs: FIGURE_DPI, FIGURE_MARGIN_FRAC, FIGURE_TILE_ROWS/COLS, FIGURE_TILE_
OVERLAP, FIGURE_FULLPAGE, CAPTION_MAX_PAGES, larger CAPTION_MAX_TOKENS, and
CAPTION_MAX_IMAGES as a per-document tile budget.

Measured figure context recall (per-label, dense academic figures):
  Qwen2.5-VL: 0.50 -> 0.83 (overall 0.81 -> 0.94)
  Gemma-4-E2B (weak): ~0 with loops -> 0.83 (overall 0.91)
Born-digital text and scanned-page recall are unchanged (no regression).

parsers gains _figure_boxes (shared detection), pages_with_figures, and
render_pdf_figure_tiles; captioner gains merge_page_captions and a temperature
parameter; ingestion routes figure captioning through the tiled path.

* Fix RAG review issues: whole-doc budget pre-check, figure gating, empty re-ingest, vision auth

Whole-document context now runs a cheap token-sum pre-check (store.scope_token_estimate)
before hydrating every chunk's text, so an attachment that cannot fit the budget is
rejected without loading the whole corpus into memory. The estimate mirrors
all_chunks_for_scope's filter and the per-row token-count fallback exactly.

Ingestion skips all figure work (PDF rasterization and detection, not just the caption
call) unless a vision model is loaded, so a text-only deployment pays nothing. When OCR
is enabled, scanned/image-only pages are excluded from figure tiling since OCR already
transcribes them whole, avoiding double vision work and overlapping index entries; a
scanned figure page is still tiled when OCR is off.

start_ingestion no longer dedupes forever to a prior ingest that produced zero chunks
(e.g. a scanned PDF uploaded before a vision model was loaded): the empty record is
dropped and the content is re-ingested.

Vision OCR and caption requests now send the backend Authorization header, so they
match the chat endpoint and do not 401 under direct-stream (--api-key) mode.

Adds tests for the budget estimate, scanned-page exclusion, the vision-model gate, the
empty re-ingest path, and the auth-header passthrough.

* Trim RAG vision-ingestion comments and docstrings

Tighten the verbose multi-line docstrings and comments added across the RAG vision
ingestion work (captioner, config, parsers, ingestion, store, tool, build_rag_autoinject,
the RAG tests, and the chat-store/upload-hook frontend toggles) to one or two lines while
keeping their intent. No code changed: verified comment/docstring-only against the prior
commit, and the RAG test suite still passes.

* Fix figure-tiling exclusion and client dedupe for re-ingestable docs

Figure tiling now excludes only the pages OCR actually transcribed, not every
text-less page. _ocr_scanned_pages returns the set of pages it OCR'd, and _run passes
that to pages_with_figures as exclude_pages (replacing the ocr_on-keyed min_text_chars
heuristic). A scanned page that OCR skipped (past OCR_MAX_PAGES, or whose OCR returned
empty) is no longer dropped from captioning, so a chart on such a page still gets a
caption.

The document panel's upload dedupe no longer skips re-selecting a file whose only
matching doc completed with zero chunks. Such a doc is re-ingestable (e.g. a scan
attached before a vision model loaded), and the backend re-ingests on the same content
hash, so the client must let it reach the backend; healthy or still-indexing docs are
still skipped. The SSE complete frame's chunk count is recorded on the doc so the
check is exact.

Adds a regression test for the un-OCR'd scanned figure page and updates the
pages_with_figures test to the exclude_pages interface.

* Address review findings: whole-doc budget guard, job numChunks, dead code, upload cap

whole_document_context now treats a non-positive max_tokens as "never inject" instead
of injecting the whole corpus unbounded, so RAG_WHOLE_DOC_MAX_TOKENS=0 tightens rather
than disables the budget (the real off switch stays RAG_THREAD_WHOLE_DOC=0).

The job-status endpoint and get_job_status now expose num_chunks (joined from the
document), and the upload hook threads it through the SSE-fallback completion paths
(reconcile + poll). Previously a document that finished via the connection-cap fallback
had no chunk count client-side, so the re-ingest dedupe wrongly treated it as empty and
re-uploaded it. IndexJob/JobEvent gain the field and the untyped cast is dropped.

Removes the dead render_pdf_figures function (superseded by the tiling path), its test,
and the unused FIGURE_MARGIN_FRAC config knob.

Adds an upload size cap (RAG_MAX_UPLOAD_BYTES, default 200 MB; 413 on exceed with the
partial file cleaned up) so a pathological file can't drive unbounded parse + vision
work. render_pdf_figure_tiles clamps rows/cols to >= 1 (no ZeroDivisionError on a
misconfigured grid). Captioning progress is reported after OCR so the bar is monotonic.
sqlite connections set busy_timeout=5000 so a long figure/scan ingest holding its
connection doesn't make a concurrent ingest/read fail with "database is locked".

Adds tests for the non-positive budget, the zero-grid clamp, job-status num_chunks, and
the oversize-upload rejection.

* Extract PDF text as layout-aware Markdown via pymupdf4llm

parsers._pdf now extracts each PDF page as Markdown with pymupdf4llm.to_markdown
(page_chunks=True) instead of flat page.get_text("text"), so tables, headings and lists
keep their structure in the indexed chunks and retrieve far better (a table's cells stay
associated with their row instead of flattening into a token stream). Gated by
RAG_PDF_MARKDOWN (default on); falls back to plain PyMuPDF text when the toggle is off,
pymupdf4llm is missing, extraction fails, or a page yields no Markdown. The scanned-page
OCR and figure-tiling passes operate on rendered pixels and are unaffected; docx/html/txt
keep their existing extractors.

The preview-highlight locator already strips Markdown punctuation when building anchors;
it now also splits anchor tokens on pipes so a Markdown table row still anchors to the
raw PDF word stream.

Declares pymupdf4llm as a studio/RAG dependency (was only transitively present via the
data-designer plugin). Adds parser tests (Markdown table reaches the page text, the
plain-text fallback, the missing-lib fallback) and a locator test for table-pipe anchoring.

* Pin pymupdf4llm to 0.3.4 so the package scan does not pull onnxruntime

The lockstep pymupdf4llm 1.27.x line makes pymupdf-layout a hard dependency,
which in turn pulls onnxruntime (plus numpy/networkx/protobuf). The security-audit
pip scan-packages job resolves requirements --with-deps, so adding pymupdf4llm to
no-torch-runtime.txt and studio.txt surfaced onnxruntime's un-baselined CRITICAL
finding and flipped the hf-stack shard from pass to fail.

pymupdf4llm 0.3.x keeps pymupdf-layout behind an optional [layout] extra, so a plain
install resolves to pymupdf + tabulate only and never touches onnxruntime. 0.3.4
requires pymupdf>=1.27.1, satisfied by our pinned pymupdf==1.27.2.3, and to_markdown
(page_chunks=True) produces equivalent layout-aware Markdown on real PDFs (verified on
the Attention, ResNet and USDA documents). Production already installs these files
--no-deps, so onnxruntime was never shipped at runtime; this only fixes the scanner.

The parser test now asserts Markdown markup (heading or table pipes) rather than table
pipes specifically, since 0.3.4 emits a heading but not a pipe table on the tiny
borderless synthetic fixture; both markers are absent from the plain-text fallback.

* Fix RAG whole-doc review findings

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

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

* Address RAG whole-doc review follow-ups

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

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

* Address RAG review follow-up edge cases

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

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

* Reserve image budget for whole-document RAG

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

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

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-30 15:55:23 +02:00
Daniel Han
43d3caf38b
Studio: imatrix GGUF option and FP8/NVFP4 compressed export in the export UI (#6729)
* Studio: wire imatrix GGUF option and FP8/NVFP4 compressed export into the export UI

GGUF export gains an importance-matrix toggle. When enabled it auto-downloads the
upstream Unsloth imatrix for the base model (or uses a custom path), which unlocks
the IQ low-bit quants iq2_xxs, iq2_m, iq3_xxs and iq4_xs. Merged export gains an
FP8 / NVFP4 compressed-tensors precision selector that runs llm-compressor for vLLM.

Backend threads imatrix_file through routes -> orchestrator -> worker -> export_gguf
(both the local save and the hub push), and maps the new compressed format_type
values onto the fp8/nvfp4 save_method, reporting the "<dir>-<suffix>" sibling output
directory. Frontend adds the imatrix Switch on the GGUF card and a merged precision
picker on the merged card, threaded through the export runtime store.

Depends on unslothai/unsloth#6706 (save.py imatrix_file and compressed-tensors
export) and unslothai/unsloth-zoo#839 (quantize_gguf imatrix flag).

* Studio export: guard imatrix/compressed against older unsloth builds and force imatrix for IQ quants

Addresses review feedback on the export wiring:
- GGUF: pass imatrix_file only when set, so a plain no-imatrix export (e.g. Q4_K_M) no
  longer fails with an unexpected-keyword error against an unsloth build that predates the
  imatrix_file parameter. When imatrix is requested but unsupported, return a clear
  upgrade message instead of a TypeError.
- Merged: gate FP8/NVFP4 compressed-tensors export on the installed unsloth actually
  supporting it, returning a clear message rather than a cryptic save_method failure.
- Frontend: IQ quants (iq2_xxs, iq2_m, iq3_xxs, iq4_xs) are imatrix-only, so force the
  imatrix on when one is selected and lock the toggle, instead of submitting an IQ quant
  with no imatrix that llama.cpp would reject.

Extends the backend tests for the new capability guards and the conditional kwarg wiring.

* Studio: upload compressed merged models to the Hub without recompressing

For an FP8/NVFP4 Hub export the model is already produced locally in the "<dir>-<suffix>"
output. Uploading it directly with HfApi.upload_folder (mirroring export_base_model) avoids
re-running the expensive compressed-tensors quantization a second time inside
push_to_hub_merged, which for NVFP4 also re-runs calibration and risks OOM. Falls back to
push_to_hub_merged when there is no local compressed output to reuse.
2026-06-30 03:41:02 -07:00
Nilay
32f28b2180
Studio: keep "Fine-tuned" compare label clear of the floating top right controls (#6755)
* fix header overlap

* fix

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-30 11:26:33 +01:00
oobabooga
0ace5cc436 Coordinate diffusion with training VRAM and guard its cache deletes 2026-06-29 22:20:08 -03:00
OrbisAI Security
220ff5aaba
fix: CVE-2026-54290 security vulnerability (#6736)
Automated dependency upgrade by OrbisAI Security

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-29 15:28:06 +01:00
Lee Jackson
f7d509e1f2
fix: remove sidebar update dev override (#6746) 2026-06-29 16:17:16 +02:00
Michael Han
1069b28c43
Studio: name the missing extractor when a Recipes file upload fails (#6642)
* Studio: name the missing extractor when a Recipes upload fails

A missing optional dependency (pymupdf4llm for PDF, mammoth for DOCX) was
reported as a generic "Text extraction failed", which gives the user nothing to
act on. Catch ImportError and surface the package name instead.

* Studio: narrow missing extractor error handling

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

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

---------

Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-29 16:12:44 +02:00
Michael Han
11469a60fe
(feat) Add project names to studio training runs (#6512)
* (feat) Add project names to studio training runs to avoid models being overwritten when doing similar training runs

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

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

* Update studio/frontend/src/features/export/export-page.tsx

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update studio/frontend/src/features/export/export-page.tsx

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

* Update studio/frontend/src/features/export/export-page.tsx

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* better project name sanitization, removed duplicated project name normalization

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

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

* implement checkpoint scanning utilities and tests for base model inference

* [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

* Guard project_name against null and use leading important modifiers

* Fix/adjust training project names for PR #6512

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

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

* Fix/adjust training project names for PR #6512

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

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

* Address project-name review feedback

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

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

* Show project names in training recents

* Keep GGUF export directories source-specific

---------

Co-authored-by: NZ-Linix <nz-linix@outlook.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: NZ-Linix <linus.ordowski@outlook.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-06-29 16:06:36 +02:00
Lee Jackson
f80e66ea34
studio: keep chat header below dialogs (#6745) 2026-06-29 16:03:57 +02:00
Lee Jackson
07578eab60
Fix on-device locations dialog layout (#6743)
* Fix on-device location path overflow

* Remove redundant native path tooltip
2026-06-29 15:50:55 +02:00
Wasim Yousef Said
755da2f155
Speed up Studio desktop startup (#6742)
* Speed up Studio desktop startup

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

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

* Address Studio startup review findings

* Keep orphaned run cleanup before readiness

* [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-06-29 15:27:39 +02:00
Lee Jackson
2f8521ed54
Fix compare adapter selection (#6411) 2026-06-29 12:57:48 +02:00
Daniel Han
cbb2eaa729 Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats
Codex review on the native engine arg builder:

- build_sd_cpp_command emitted --width/--height unconditionally, so an
  img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
  the input. width/height are now Optional (None = unset): an image-conditioned
  run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
  the size from the input image (set_width_and_height_if_unset); a plain txt2img
  run with unset dims keeps the prior 1024x1024 default; explicit dims are always
  honored. width/height are read only by the builder, so the type change is local.

- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
  that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
  explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
  and emits the flag for any explicit value != 1.

Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)
2026-06-29 10:36:44 +00:00
Daniel Han
4d87ff4b22 Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs
Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.

Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.
2026-06-29 10:33:09 +00:00
Matt Van Horn
54b95fbcc8
fix(studio): show local file path tooltip for Hub-tab local models (#6715)
Local models in the Studio Hub tab (Custom folders, LM Studio, and
Local models sections) did not reveal their on-disk path on hover,
unlike the Fine-tuned rows which already do. Each of these rows maps
over a LocalModelInfo with a required path, so pass tooltipText built
from the model name and path via a small shared localPathTooltip
helper, matching the existing FT-row tooltip format.

Refs #6382

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-29 11:12:47 +01:00
pre-commit-ci[bot]
7be5ec4fdd [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-06-29 05:47:15 +00:00
Daniel Han
6c8b0d146e Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening
- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
  timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
  extract through a per-member containment check (Zip-Slip guard); expanduser the
  --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
  the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
  installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
  sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
  crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.
2026-06-29 05:46:45 +00:00
oobabooga
bebf1ca0c2 Gate Images picker to GGUF across all format-filter selections 2026-06-29 02:21:49 -03:00
Nilay
b56d24ea3e
Studio: cascade user message deletion to include assistant reply (#6720)
* cascade user message deletion to include assistant reply

* Fix comment typo in delete-thread-message

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-28 04:16:03 -07:00
Daniel Han
8cc0fe3790 Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output
Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.
2026-06-28 06:19:08 +00:00
Daniel Han
0ad814a452
Revert "feat: add GPU-aware model filtering and For You section- Add fit filt…" (#6722)
This reverts commit a636693019.
2026-06-27 17:48:16 -07:00
Daniel Han
4c72e09480
Studio: stop handing CI/user secrets to downloaded llama.cpp binaries (#6696)
* Studio: stop handing CI/user secrets to downloaded llama.cpp binaries

The macOS prebuilt path installs llama.cpp from the unslothai/llama.cpp
fork's latest (unpinned, mutable) release and then executes the
downloaded llama-server / llama-quantize binaries during install-time
validation. binary_env() built that child environment from a full
os.environ.copy(), so a compromised or tampered prebuilt would inherit
every secret in the process: HF_TOKEN and the workflow GitHub tokens in
CI, and HF / cloud credentials for end users running install.sh /
setup.sh.

We publish prebuilts daily, so pinning a release tag is not workable.
Instead, neutralise the impact: these binaries have no reason to read any
token, so strip secret-bearing variables (exact names plus
TOKEN/SECRET/PASSWORD/CREDENTIAL/PRIVATE_KEY/API_KEY markers) before
handing the env to a downloaded binary. The installer's own GitHub and
Hugging Face API calls read os.environ directly, so authentication and
release-API rate limiting are unaffected; PATH, LD_LIBRARY_PATH,
DYLD_LIBRARY_PATH and CUDA/ROCm vars are preserved. One change covers the
install-time validation path for all six macOS workflows and end users.

Follow-up (separate, sequenced): publish build-provenance attestations
from the fork's prebuilt workflows and verify them in CI, so a forged
release is rejected rather than merely starved of secrets.

* Strip KUBECONFIG, SSH_AUTH_SOCK, and PASSPHRASE-marked vars from binary env

Extend the deny-list per PR review: KUBECONFIG and SSH_AUTH_SOCK are
credential pointers/capabilities a downloaded binary never needs, and a
PASSPHRASE marker catches SSH_PASSPHRASE / GPG_PASSPHRASE. Tests updated.

* Studio: also scrub proxy/index env vars and URL-embedded credentials before running prebuilt binaries

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

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

* Scope mlx-ci secrets to the install + download commands for PR #6696

Drop the ambient step-level env block and pass GH/GITHUB/HF tokens only
on the installer and GGUF-download commands, so the directly invoked
llama-quantize / llama-server smoke runs see no secrets. The installer
still reads tokens from os.environ for the releases API and probe fetch.

* Trim verbose comments around the secret-env scrubber for PR #6696

Comment-only: condense the block comments added across this PR. Logic
unchanged (comment_tools.py check confirms code-only signature equal).

* Redirect HOME / cache pointers to an empty dir for prebuilt binaries (PR #6696)

Address Codex P2: stripping token env vars still let a tampered binary
read on-disk token stores (~/.cache/huggingface/token, ~/.aws/credentials,
~/.config/gh) through $HOME and the cache/config pointers. Point HOME plus
the HF / XDG / Windows home pointers at a single empty throwaway dir for
the downloaded-binary env. Defense in depth: a binary resolving the real
home via getpwuid is out of scope and needs OS sandboxing.

* Close residual credential-probe gaps for PR #6696

Address the latest Codex review:
- Strip token-only URL userinfo too (scheme://ghp_token@host), not just
  the user:pass form.
- Redirect HOMEDRIVE/HOMEPATH alongside USERPROFILE so a Windows binary
  cannot reconstruct the real profile from %HOMEDRIVE%%HOMEPATH%.
- Drop explicit credential-file pointers (NETRC, PIP_CONFIG_FILE,
  DOCKER_CONFIG, GIT_CONFIG_GLOBAL) that live outside HOME.
- Probe ldd with a secret-free env: linux_runtime_dirs ran ldd on the
  untrusted prebuilt with the inherited os.environ, and ldd may execute
  the binary, so it could observe HF_TOKEN/GITHUB_TOKEN during the probe.

Factored the shared scrub into secret_free_environ().

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

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

* Separate token-bearing install from binary smoke; drop CI command files (PR #6696)

Address the two P1s in the latest review:
- mlx-ci: GitHub bakes secrets into the run-script text, so inline token
  assignments in a step that later runs the prebuilt let a tampered binary
  read them from the script. Split into a token-bearing install + download
  step that never launches a binary, and a secret-free smoke step that runs
  llama-quantize / llama-server.
- secret_free_environ now drops the GitHub Actions command files
  (GITHUB_ENV, GITHUB_PATH, GITHUB_OUTPUT, GITHUB_STEP_SUMMARY, BASH_ENV) and
  the smoke step unsets them, so a tampered prebuilt cannot inject PATH/env
  into the later token-bearing MLX steps.

* Run the prebuilt smoke last, after all token-bearing steps (PR #6696)

Address the P1 workspace-poisoning vector: even with no secrets in its env,
a tampered prebuilt could edit the checkout or installed modules, and the
later HF_TOKEN MLX steps would then execute that poisoned code on push
builds. Move the prebuilt install + smoke to the end of the job so the
untrusted binary runs after every token-bearing step, leaving nothing for it
to corrupt. The MLX GGUF reload uses a source-built llama-cli, not this
prebuilt, so nothing depends on the earlier position.

* Trim comments around the secret-env scrubber and prebuilt CI steps (PR #6696)

Comment-only: condense the security-rationale block comments and merge the
duplicated prebuilt-step description in mlx-ci. Logic unchanged
(comment_tools.py check confirms the code-only signature is equal; install
suite still passes).

* Authenticate the GGUF export release-API lookup with the read-only GITHUB_TOKEN (PR #6696)

* Rename env scrubber off the secret-named identifier CodeQL flags as a clear-text sink (PR #6696)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-27 05:21:05 -07:00
Daniel Han
98a01e70cd
Studio: restore tensor parallelism for vision/mmproj GGUFs (#6659)
* Studio: restore tensor parallelism for vision/mmproj GGUFs

#6416 disabled --split-mode tensor for any GGUF that ships an mmproj projector to
dodge a GGML_ASSERT crash (#6415) seen on an older llama.cpp build with consumer
Blackwell (sm_120). The blanket skip silently dropped tensor_parallel=true for
every multimodal/MTP GGUF (e.g. Qwen3.6-35B-A3B-MTP); on hardware where the model
fits on one GPU the load then collapsed to a single GPU. mmproj + --split-mode
tensor works on current builds (verified end to end on B200/sm_100), so the skip
was disabling a working configuration.

Make the vision skip self-healing per binary:
- attempt tensor for vision models by default
- skip upfront only on a binary already seen to abort on tensor + mmproj this
  session (_vision_tensor_split_aborts), recorded when such a launch crashes at
  startup (_record_vision_tensor_split_abort). Process scoped, so a studio update
  re-probes the new build. The route-level layer-split fallback stays the net.
- add _select_gpus(min_gpus=...) so a downgraded tensor request can keep multiple
  GPUs instead of collapsing to one (default 1, no behavior change).

Add tests/test_tp_vision_regression.py: an AST allowlist guard over the
tensor_parallel drop sites (which would have flagged #6416), plus cache and
_select_gpus coverage. No GPU required.

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

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

* Studio: address review on vision tensor-parallel self-healing

Three fixes from the PR review:

- Record a vision-tensor abort only after every startup retry fails. The first
  version cached the binary on the first spawn crash, which on every build
  (including capable ones) is the benign --fit step abort that the existing
  --fit off retry resolves. That poisoned the cache so the next vision load in
  the same process skipped tensor. Recording now happens at the post-retry
  failure block (after fit-off, flash-attn-off and MTP-drop), so a binary that
  actually works is never cached.

- Gate the record on the tensor/mmproj crash signature: a hard signal fault
  (_is_signal_crash) with no non-tensor cause (_output_has_nonprojector_diagnostic
  excludes OOM and unknown-arch), so an OOM, bad extra args, or MTP/flash-attn
  crash no longer marks an otherwise capable binary incompatible.

- Preserve the multi-GPU request on the cached downgrade. The vision gate now
  raises _layer_min_gpus to the visible GPU count and threads it through the
  layer-split GPU selection (_select_gpus min_gpus and the subset loops), so a
  downgraded tensor request still spreads across GPUs instead of collapsing to a
  single card the model happens to fit.

Verified two vision+tensor loads in one backend process both tensor-split across
4 GPUs (the benign fit abort no longer poisons the cache). Tests updated.

* Studio: harden vision tensor-parallel self-healing (review round 2)

Address the second review round on the vision/mmproj tensor-parallel fix:

- Preserve vision on the first load: a --split-mode tensor + --mmproj
  GGML_ASSERT now raises so the route-level tensor->layer fallback retries
  layer split with the projector intact, instead of stripping --mmproj and
  silently loading text-only (which returned success and skipped the fallback,
  losing vision on the first load until the next cached load).

- Symmetric multi-GPU preservation: the pooled-VRAM tensor downgrade now raises
  _layer_min_gpus from the usable tensor GPUs like the vision downgrade, so it
  no longer collapses a multi-GPU request to a single card.

- Base the layer fallback minimum on usable GPUs: _select_gpus caps min_gpus to
  the count of cards with usable VRAM, so a downgrade never forces a nearly-full
  card in (or trips --fit) just to hit the count.

- Re-probe after in-app updates: key the per-binary abort cache on (path, mtime)
  like _capability_cache, so POST /api/llama/update swapping the binary in place
  (no backend restart) re-probes the new build instead of inheriting the old
  build's abort.

- Bump _layer_min_gpus for a known-bad vision binary independent of the tensor
  drop, so the route fallback's layer retry (tensor already off) still spreads
  across GPUs.

Adds deterministic non-GPU regression tests for each.

* Studio: gate cached-vision layer minimum on the current tensor request

The cached-vision _layer_min_gpus bump fired for every later vision load on a
binary recorded as tensor+mmproj-incompatible, including loads that did not
request tensor parallelism. A plain non-tensor vision load that fits on one card
would then grab every GPU just because an earlier TP attempt aborted in the same
backend process.

Re-tie the bump to the current tensor request (back inside the tensor-drop
guard), so only a downgraded tensor request preserves the multi-GPU spread; a
non-tensor vision load minimizes device count as before.

* Studio: preserve GPU count + confirm assert on vision tensor fallback

Third review round on the vision/mmproj tensor-parallel fix:

- Preserve multi-GPU on the first tensor->layer fallback. The route-level retry
  runs tensor-off, so the in-function downgrades can't see the original tensor
  request and a fits-on-one-card model loaded the first successful fallback on a
  single GPU. The GGUF load closure now passes preserve_multi_gpu_on_layer (the
  toggle asked for tensor, this attempt is layer) and load_model raises
  _layer_min_gpus for it, so the downgrade still spreads across GPUs.

- Cap the auto-context layer loops to usable GPUs. They bypass _select_gpus, so a
  raised _layer_min_gpus could force a nearly-full card into the subset (or trip
  --fit). They now start from _auto_min_gpus, capped to the GPUs with usable VRAM.

- Confirm the tensor/mmproj assert before caching. Recording (and the layer-retry
  raise) now require the ggml assert marker via _is_tensor_split_assert, not the
  bare-signal predicate shared with the projector-incompat branch, so a corrupt
  or too-new projector that SIGSEGVs independent of split mode is no longer cached
  as tensor/mmproj-incompatible.

Adds deterministic non-GPU regression tests for each.

* Studio: extend multi-GPU fallback to extra/env tensor + overhead-aware cap

Fourth review round on the vision/mmproj tensor-parallel fix:

- Preserve multi-GPU fallback for all tensor requests, not just the UI toggle.
  Tensor can also be requested via --split-mode tensor in extra args or an
  inherited LLAMA_ARG_SPLIT_MODE=tensor env; the fallback retries those too, so
  the preserve_multi_gpu_on_layer hint now keys off _effective_tensor_parallel
  (the same check the fallback uses), comparing the overall request against the
  current attempt instead of only request.tensor_parallel.

- Cap the auto-context layer fallback to GPUs that can pay the per-device layer
  overhead. The cap counted any card with positive usable VRAM, so a nearly-full
  GPU with a few MiB free stayed eligible and could be exposed to llama.cpp and
  OOM. It now mirrors _select_gpus: a card counts only if usable VRAM exceeds the
  per-device pipeline overhead.

Adds deterministic non-GPU regression tests for both.

* Studio: match the #6415 split-axis assert + replay layer-preserve hint

Fifth review round on the vision/mmproj tensor-parallel fix:

- Narrow the tensor/mmproj crash signature. _is_tensor_split_assert matched any
  GGML_ASSERT/GGML_ABORT, so an unrelated invariant a corrupt GGUF or projector
  trips with --mmproj present could be cached as tensor/mmproj-incompatible. It
  now matches the specific #6415 warmup assertion
  (GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) in ggml-backend-meta),
  whose split-axis signature is inherent to tensor splitting. A reworded future
  assert just re-crashes-then-falls-back (vision preserved via layer split)
  instead of poisoning the cache for other models.

- Persist the layer-preserve hint for respawns. A successful tensor->layer
  fallback committed _last_load_kwargs without preserve_multi_gpu_on_layer, so
  _respawn_if_dead replayed only --split-mode layer + tensor_parallel=False and a
  mid-session respawn of a fits-on-one-card model came back single-GPU. The hint
  is now in the replay snapshot, so recovery keeps the multi-GPU placement.

Adds deterministic non-GPU regression tests for both.

* Studio: tighten comments on the vision tensor-parallel fix

Make the comments and docstrings added by this PR succinct: collapse the
multi-line block comments in llama_cpp.py / inference.py to one or two lines,
trim the verbose test docstrings (the names and assert messages already carry the
intent), and shorten the module docstring. No code changes; verified comment-only
with scripts/comment_tools.py check --strip-docstrings.

* Studio: cache vision tensor abort only on the split-axis token

_is_tensor_split_assert also accepted any GGML_ASSERT/GGML_ABORT from
ggml-backend-meta, but that file holds many asserts, so an unrelated
scheduler/projector/model invariant on an --mmproj launch could cache the binary
as tensor/mmproj-incompatible and make later compatible vision models skip tensor
parallelism. Match the GGML_BACKEND_SPLIT_AXIS_* token itself (unique to the
#6415 warmup assert), not the source file name.

* Studio: don't leak the httpx test stub into later tests

The regression module stubbed httpx via sys.modules.setdefault, which installs
the lightweight stub even when real httpx is present but not yet imported. The
stub then persists for the whole pytest process, so provider/HF tests collected
later (importing httpx or huggingface_hub.errors) got a module missing
HTTPError/Response. Mirror the neighboring llama_cpp helper tests: import real
httpx first and only fall back to a stub on ImportError.

* Studio: latch the #6415 tensor-split abort on the first spawn, key it per model

The self-heal recorded the --split-mode tensor abort only in the post-retry
failure block, after the flash-attn-off retry. But SPLIT_MODE_TENSOR requires
flash_attn, so the flash-off retry can't run tensor and its output no longer
carries the warmup split-axis assert (ggml-backend-meta :541). The record
therefore never fired on the real reproducer and the crash loop repeated on
every load (reported by oobabooga on #6659).

Latch instead on the first spawn that shows the signal crash + split-axis
marker: record it, kill the process, and raise straight to the route's layer
fallback, skipping the futile flash-attn/MTP retry ladder for this crash.

The crash is a tensor-split geometry limit (e.g. MQA n_head_kv=1 splitting to
GGML_BACKEND_SPLIT_AXIS_0), not a vision/mmproj property: it reproduces without
--mmproj and even single-GPU tensor. So drop the vision/mmproj scoping, rename
_vision_tensor_* -> _tensor_split_*, and key the session cache on
(binary, mtime, model) rather than (binary, mtime) so one model's abort no
longer skips tensor for every other model on the same build.

Regression tests updated to pin the early-spawn record, the per-model cache,
and that an unrelated ggml-backend-meta assert is not treated as the marker.

* Studio: reload on explicit tensor-off after a multi-GPU layer fallback

When a tensor load is downgraded to layer but kept multi-GPU to honor the
tensor request (preserve_multi_gpu_on_layer, the geometry-cache gate, or the
budget downgrade), the server reports tensor_parallel=False with --split-mode
layer stored. A later Apply that explicitly turns the tensor toggle off then
matched the loaded state and deduped to already_loaded, so Studio kept the
fallback's all-GPU CUDA_VISIBLE_DEVICES placement instead of re-selecting
normal placement (a single GPU for a model that fits on one card).

Latch a _layer_preserves_tensor_intent flag in load_model whenever a tensor
request is downgraded to layer with the multi-GPU floor raised
(_layer_min_gpus > 1), clear it when tensor stays on or on unload, and force a
reload in _request_matches_loaded_settings when the user explicitly turns the
tensor toggle off while that flag is set. An Apply that does not touch the
toggle still dedupes, so a working multi-GPU layer server is not churned.

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

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

* Studio: address reviewer.py findings on the tensor-split self-heal

P1 (dedup): tensor intent can be dropped via extras, not only the toggle. An
explicit llama_extra_args=["--split-mode", "layer"] matches the stored fallback
extras, so _request_matches_loaded_settings deduped to the preserved all-GPU
placement instead of reloading. Now reload when layer_preserves_tensor_intent
and the user explicitly drops tensor via the toggle OR via extras
(_effective_tensor_parallel of the explicit extras is false).

P1 (downgrade symmetry): the len(tp_gpus) < 2 compute-buffer downgrade cleared
tensor_parallel without raising _layer_min_gpus, unlike the budget and geometry
downgrades. GPUs below tensor's replicated compute-buffer reserve can still take
layer split's lower overhead, so keep the multi-GPU request (len(gpus) >= 2) and
let _select_gpus cap unusable cards.

P2 (cache key): key the tensor-split abort cache on st_mtime_ns, so a binary
replaced in place within the same second after an abort is re-probed instead of
inheriting the stale entry.

P2 (test hygiene): load routes/inference.py via importlib in the regression
tests instead of importing the routes package, which runs routes/__init__.py and
pulls in every router (e.g. python-multipart). Added regression coverage for the
extras-off reload, the compute-buffer multi-GPU preservation, and the same-second
nanosecond cache invalidation.

* Studio: record the tensor-split abort on the Windows CRT abort exit too

The first-spawn split-axis latch only recorded when _is_signal_crash matched
(POSIX signal or 0xC0000000+ NTSTATUS). On MSVC builds GGML_ASSERT terminates
through the CRT abort() path with exit code 3, which is neither, so the cache
never filled on Windows and every later load of the same bad binary/model
repeated the tensor crash before falling back to layer.

The split-axis marker is definitive, so accept either a signal crash or the
Windows abort() exit (3) when the marker is present. Add _is_abort_exit and a
unit test, and assert the early latch honors it.

* Studio: fix UnboundLocalError on --fit-on fallback, reload backend fast path

Two follow-ups from review on the tensor-split self-heal:

UnboundLocalError: _layer_min_gpus was initialized inside the GPU-selection try.
If NVML probing or GGUF/mmproj sizing raised, the except path logged "using
--fit on" and fell through to the command builder, where the new
self._layer_preserves_tensor_intent = _layer_min_gpus > 1 then raised, turning a
safe --fit-on layer fallback into a hard load failure. Bind _layer_min_gpus
before the try so the except path always has it.

Backend fast path: _request_matches_loaded_settings forces a reload when a
preserved tensor->layer fallback gets an explicit tensor-off request, but
load_model's own _already_in_target_state still matched the tensor-off/layer
settings and short-circuited, so the placement re-selection never ran. Mirror
the guard there: reload when layer_preserves_tensor_intent and the request drops
tensor intent. The flag clears on that reload, so there's no loop.

Added regression coverage for both.

* Studio: testable tensor-split record decision; skip futile fit-off retry

Follow-ups from a deeper review of the tensor-split self-heal:

Extract the record decision into _should_record_tensor_split_abort(rc, output)
(marker AND (signal crash OR Windows abort)) and call it from the early latch.
The combined boolean was only covered by source-inspection substring checks, so
an or->and typo would silently stop recording on Windows (CRT abort exit 3 is
not a signal) with every test still green. Add a behavioral test over the
POSIX / Windows / NTSTATUS / clean-exit / SIGKILL / no-marker matrix.

Skip the --fit off retry inside _spawn_and_wait when the crash already shows the
split-axis marker: that abort is fit-independent, so the retry just warms up and
crashes a second time before the latch records it. Skipping it lets the caller
latch immediately and corrects the latch comment.

Also clarify the dedup-guard comments (toggle read from model_fields_set vs
extras via _effective_tensor_parallel without env; the backend fast path is
intentionally broader and only ever forces a reload).

* Studio: don't reload-loop tensor-off requests under env tensor

The preserved-fallback reload guard fired on the raw tensor toggle, ignoring
LLAMA_ARG_SPLIT_MODE=tensor. For an env-driven tensor user, an explicit
tensor_parallel=false request then forced a reload that re-engaged tensor via
the env and re-created the same preserved layer fallback, so every /load
reloaded -- bypassing the env-downgrade matching that exists to avoid exactly
this loop.

Gate the guard on the env-aware effective tensor state: reload only when an
explicit toggle/extras change leaves _effective_tensor_parallel (which consults
the env) off. If the env still forces tensor, fall through to the existing
env-downgrade match, which dedupes instead of looping. Added a regression test
with LLAMA_ARG_SPLIT_MODE=tensor set.

* Studio: tighten comments and test docstrings on the TP self-heal

Condense the verbose comments and test docstrings added across the review rounds
into fewer, succinct lines without changing their intent: the early-latch and
downgrade-site rationale, the cache/key and helper docstrings, the dedup-guard
comments, and the per-test docstrings. No code changes (AST-verified comments
and docstrings only); tests and lint unchanged.

* Studio: clear preserved tensor flag on diffusion; carry it across non-drop reloads

Two follow-ups on the preserved-fallback machinery:

Diffusion: the DiffusionGemma path early-returns from load_model before the
command builder that sets/clears _layer_preserves_tensor_intent, so the flag
from a prior tensor->layer fallback leaked onto a later diffusion load and
forced needless reloads of the diffusion server on tensor-off/extra Applies.
Clear it when starting diffusion.

Settings reload: the preserve hint was recomputed only from the new request, so
a reload for an unrelated setting (e.g. max_seq_length) with the tensor toggle
omitted dropped a preserved multi-GPU layer placement back to one GPU. Carry
llama_backend.layer_preserves_tensor_intent into the hint when the request is
not an explicit tensor-off/extras-off drop, so a fitting model stays multi-GPU.

Added regression tests for the diffusion clear, the carry-forward, and the
updated tensor-intent computation.

* Studio: gate the preserve carry-forward on the same model being loaded

The tensor-intent carry-forward read llama_backend.layer_preserves_tensor_intent
without checking it belonged to the model being loaded. On a direct model switch
(load B without an explicit /unload of A), the flag is still set from A's
downgrade (it isn't reset until B's load_model reaches the command builder, after
the route reads it), so a plain load of B got preserve_multi_gpu_on_layer=True
and was spread across all GPUs even though it fits on one and the user never
requested tensor for it. The backend dedup doesn't have this leak (it checks
model_identifier first); the leak was only in the route hint.

Extract the decision into _carry_preserved_tensor_intent(preserved, same_model,
explicit_drop) and gate it on the backend still holding the same model. Add a
behavioral truth-table test (catches a `not` inversion and a missing same-model
guard) and tighten the compute-buffer downgrade test to bound its source window.

* Studio: match the HF quant too when carrying preserved tensor intent

The same-model guard on the preserve carry-forward compared only model_identifier,
which is variant-agnostic for HF repos. A later load of the same repo with a
different gguf_variant (which already bypassed dedupe on the variant mismatch)
was treated as the same model, so a request that omits tensor settings inherited
the prior variant's preserved intent and forced multi-GPU layer placement for a
quant that never requested tensor. Also require the loaded hf_variant to match for
HF repos (local direct-file loads already differ by model_identifier path). Added
a regression test for the variant guard.

* Studio: match the loaded GGUF by path too when carrying preserved tensor intent

A local directory holding multiple GGUF variants keeps one variant-agnostic
model_identifier (the directory) while config.gguf_file selects the file, so the
same-model guard let variant B inherit variant A's preserved tensor->layer
fallback and forced B onto multi-GPU. Mirror _already_in_target_state's identity
logic: match by resolved path when both sides have a local file, else by HF
variant. #6659

* Studio: let implicit same-settings reloads dedupe after a preserved fallback

The backend _already_in_target_state mirror forced a reload on ANY effective
tensor-off request once a tensor->layer fallback was preserved. In the HF
auto-pick / local-directory flows the route-level dedup is skipped, so an
identical /load with tensor omitted reached this guard and reloaded every time
even without an explicit drop. Thread the route's preserve_multi_gpu_on_layer
decision in so only an explicit drop reloads; implicit carry-forward dedupes. #6659

* Studio: only an explicit tensor/split-mode change drops preserved intent

The explicit-drop test treated request.llama_extra_args is not None as a drop,
so a same-model reload that merely added an unrelated pass-through arg (e.g.
--top-k 20) without touching the tensor field or --split-mode disabled the
carry-forward and collapsed a fitting model back to one GPU. A drop now requires
an explicit tensor_parallel field change or a non-tensor --split-mode override,
via a shared _is_explicit_tensor_drop helper used by both the already-loaded
dedup and the load carry-forward so the two readers agree. #6659

* Studio: treat an explicit clear of extras as a tensor drop

When tensor intent was extras-driven (--split-mode tensor) and fell back to a
preserved layer split, a later request that explicitly clears extras
(llama_extra_args=[]) but omits tensor_parallel left the empty list with no
split-mode override, so the carry-forward kept the model pinned multi-GPU instead
of returning to normal layer selection. _is_explicit_tensor_drop now also counts
an explicit empty-list clear as a drop, while an unrelated extra (--top-k) or
inherit (None) still carries the preserved intent. #6659

* Studio: don't treat the UI's tensor_parallel echo as a tensor drop

The Studio frontend always sends tensor_parallel and copies the /load response's
resolved value back into its state, so after a tensor->layer fallback every
ctx/settings reload carries tensor_parallel=false even though the user never
changed it. Keying the drop on the field (or on an empty extras clear) collapsed
the preserved multi-GPU placement on the next reload. A fallback also always
stores --split-mode layer, never a tensor split mode, so a clear never wipes
tensor intent. _is_explicit_tensor_drop now drops only on an explicit non-tensor
--split-mode override; the bare field echo, an empty clear, an unrelated extra,
and inherit all keep the preserved placement, and --split-mode tensor /
tensor_parallel=true re-engage tensor. #6659

* Studio: match the resolved config.identifier when carrying tensor intent

The same-model guard for the carry-forward compared the raw request id, but
ModelConfig.from_identifier normalizes it (adds the unsloth/ prefix for a
shorthand, fixes repo-id case) before load_model stores config.identifier. So a
ctx/settings reload using the shorthand id missed the match, dropped
_carry_preserved_tensor_intent, and could collapse a preserved multi-GPU layer
placement to one GPU. Compare against config.identifier (what the backend stores),
keeping it symmetric with _already_in_target_state. #6659

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-27 01:52:18 -07:00
oobabooga
0607351a81 Merge remote-tracking branch 'origin/main' into image-generation
# Conflicts:
#	studio/backend/routes/models.py
2026-06-26 23:23:02 -03:00
oobabooga
ae7dc8937a Fix toast id type in load-progress resume 2026-06-26 23:18:36 -03:00
oobabooga
9c6f852a43 Fix image-generation GPU eviction, load lifecycle, and Images picker gating 2026-06-26 23:18:36 -03:00
Daniel Han
b11966b2db
studio: list the full local model catalog from /v1/models (#6519)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-26 20:42:06 -03:00
Daniel Han
9451aef51e
studio: return a clean model id from the OpenAI API instead of the local .gguf path (#6518)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-26 16:07:53 -03:00
Avaya Aggarwal
cb274484a6
Add GGUF --tensor-parallel CLI option (#6561)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-06-26 15:30:10 -03:00
Lee Jackson
1d4c4d4a76
Merge branch 'main' into image-generation 2026-06-26 18:22:14 +01:00
Daniel Han
2ef394137a
Studio: harden background consumer loops and streaming paths against silent UI freezes (#6653)
* Studio: harden the data-recipe and inference consumer loops against pump death

Follow-up to #6643. The same single-unsupervised-consumer pattern the training
pump had lives in two sibling loops, with the same failure mode: one bad event
kills the only thread that updates the in-memory state every UI surface reads,
while the worker subprocess keeps running.

- data_recipe JobManager._pump_loop: a malformed worker log line that makes
  parse_log_message raise no longer kills the pump. Guard _handle_event, the
  queue read, and the worker-exit finalize, and broaden _drain_queue so a drain
  error still finalizes the job instead of leaving it wedged "active" (which also
  leaked the workflow-scoped API key until its 24h expiry).
- inference InferenceOrchestrator._dispatcher_loop: guard the routing body so a
  malformed response or a mailbox put error can't kill the dispatcher and hang
  every in-flight generation (callers key liveness on the subprocess, not on
  this thread).

Adds regression tests for both.

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

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

* Studio: extend consumer-loop hardening to RAG, hub, auth, and stream-reader paths

Continuation of the data-recipe and inference pump hardening: the same
"background producer updates in-memory state that a single unsupervised
consumer surfaces to the UI" pattern shows up in several more Studio paths,
each able to silently freeze a UI surface while the worker keeps running.

RAG ingestion SSE (core/rag/ingestion.py):
- job_events polled the queue with a blocking get and never noticed client
  disconnect or a dead worker, so a closed tab or a producer that died
  without emitting a terminal event left the stream hanging. It now polls
  with a timeout, emits heartbeats, ends on terminal job status, caps idle
  time, and always pops the job registry in finally.
- Added _reap_finished_jobs() and call it from start_ingestion so finished
  job state does not accumulate.

Startup reconcile (storage/rag_db.py, main.py):
- reconcile_orphaned_ingestion_jobs() marks ingestion jobs (and their
  documents) that were left non-terminal by a previous crash as failed, so
  the UI does not show jobs stuck "running" forever after a restart. Wired
  in at startup next to cleanup_orphaned_runs().

Hub download watcher (hub/services/download_lifecycle.py):
- _watch() could leave a job pinned "running" if finalize raised. Body is
  now guarded: on failure it logs and sets the job to error, and always
  invalidates the hf cache scan in finally.

External provider stream (core/inference/external_provider.py):
- read timeout was None (no stall ceiling); set to 300s so a wedged
  upstream surfaces as an error instead of an indefinitely hung stream.

Auth store (auth/storage.py):
- Enable WAL + busy_timeout on the auth DB so token validation (read on
  every request) and login writes stop serialising on the rollback journal.
  Matches studio_db / rag_db / providers_db.

Login rate limiter (routes/auth.py):
- _LOGIN_IP_BUCKETS could grow unbounded under spoofed-IP traffic; cap it
  and prune stale buckets, mirroring the per-account bucket handling.

Training progress SSE (routes/training.py):
- Break promptly on client disconnect instead of waiting for the next
  yield to fail on a closed socket, matching the export / data-recipe SSE
  routes.

llama-server stdout drain (core/inference/llama_cpp.py):
- Broaden the drain guard so an unexpected decode/read error logs at debug
  and stops the drainer cleanly instead of escaping the thread.

Frontend stream readers (chat-api.ts, rag-api.ts):
- Wrap the SSE read loops in try/finally + reader.cancel() so early return
  ([DONE]), thrown errors, and consumer aborts release the reader lock
  instead of holding it until GC.

Tests:
- test_training_progress_stream_nan: fake request now implements the async
  is_disconnected() the route polls, matching the other SSE route fakes.

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

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

* Studio: address Codex review feedback on the consumer-loop hardening

Four follow-ups from the automated review, all on code this PR introduced:

- Data-recipe pump (manager.py): a queue read that keeps raising an error
  outside the read's narrow catch set (e.g. a broken queue pipe after the
  child died) hit the `continue` guard and skipped the dead-worker finalize
  below, spinning forever and leaving the job wedged "active" with its
  workflow key unretired. On a read failure, fall through to finalize when
  the worker is no longer alive. Added a regression test.

- RAG ingestion SSE (ingestion.py): the 5-minute idle cap could end the
  stream while the job was still pending/running (a large document spends
  minutes in embedding/storing with no per-batch progress event). The route
  then sends [DONE], and the client treats a no-terminal-frame end as
  completion, marking the document indexed mid-ingestion. Drop the idle cap:
  while the worker is alive and non-terminal we keep heartbeating; the stream
  ends only on terminal DB status, the None sentinel, or client disconnect.

- Login rate limiter (auth.py): the per-IP path pruned but then added the
  new IP unconditionally, so a spoofed-source-IP spray kept _LOGIN_IP_BUCKETS
  unbounded and made every new IP pay a full-dict prune scan. Gate the add on
  the cap, mirroring the account path.

- Hub download watcher (download_lifecycle.py): if finalize raised before it
  reaped (proc.wait) and dropped the worker (e.g. an I/O error draining
  stderr), the crash path published a terminal state while the live Popen
  stayed registered and kept writing the cache, and the terminal set_job let
  claim() admit a retry on the same repo. Terminate + drop the worker before
  setting the terminal state.

* Studio: keep login throttling working when the per-IP bucket dict saturates

Review follow-up. The previous cap fix skipped creating a bucket for a new IP
once _LOGIN_IP_BUCKETS was full, returning ip_fails=0. Under a sustained spray
that also fills the account dict, every failure from such an IP then looked
first-seen and _login_blocked had no bucket to enforce, so the cap effectively
disabled throttling once saturated.

Bound the dict with a FIFO eviction instead: if the IP is new and the dict is
full, reclaim expired buckets (rate-limited so a burst of distinct IPs can't
make each failure an O(n) sweep) and, if still full, evict the oldest-inserted
IP. The new IP always gets a real bucket, so a saturating (e.g. spoofed
X-Forwarded-For) spray stays throttled while memory stays bounded. Added a
regression test that saturates the dict and asserts a later IP is still blocked.

* Studio: address Codex review (RAG queue lifecycle, stream error, orphan chunks)

Three follow-ups on the Phase 6 changes:

- RAG ingestion SSE (ingestion.py): job_events removed the per-job queue in its
  finally on ANY exit, including an early client disconnect while the worker is
  still running. That dropped the worker's later events (the queue is the only
  one _emit writes to) and made a reconnect find no queue and receive only
  [DONE], which the client treats as completion. Only drop the queue on a
  terminal exit (None sentinel / terminal DB status); leftover terminal queues
  are still swept by _reap_finished_jobs. Added queue-lifecycle tests.

- External provider stream (routes/inference.py): once the 300s read timeout can
  fire, the stream's except path failed the monitor but ended without an error
  frame or [DONE], so the chat client saw a bare EOF and saved the timed-out
  answer as a successful partial with no error. Emit an SSE error frame (and
  [DONE]) on stream failure so the client surfaces it.

- RAG startup reconcile (storage/rag_db.py): marking a half-ingested document
  failed left its chunks/fts/vec rows intact, and retrieval filters by scope not
  status, so a failed document could still be retrieved and cited. Purge the
  document's chunks when reconciling it to failed (the doc row stays for
  re-ingest).

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

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

* Studio: release the remaining SSE stream readers (training, data-recipe, export)

reviewer.py follow-up. The chat and RAG SSE readers were wrapped in
try/finally + reader.cancel(), but the other three readers built on the same
response.body.getReader() pattern were left without it: streamTrainingProgress,
streamRecipeJobEvents, and streamExportLogs leak the ReadableStreamDefaultReader
lock (held until GC) when the consumer aborts, returns early, or a parse/callback
throws. Wrap each in try/finally + reader.cancel() (export already had a
try/catch, so it only needed the finally). All five frontend SSE readers now
release the reader symmetrically.

* Tighten resilience comments and docstrings

Condense the verbose explanatory comments and internal-helper docstrings added
in this branch to shorter, clearer forms. Comment/whitespace only; verified no
code changed via AST diff. No behaviour change.

* Studio: keep chunks for completed docs during ingestion reconcile

Startup reconciliation flips orphaned (non-terminal) ingestion jobs to failed and
purges the document's chunks so a failed source can't be retrieved. But it dropped
the chunks unconditionally, so a document the worker had already committed as
'completed' before the crash (only its job row left non-terminal) lost every chunk
while still reporting 'completed'. That leaves an empty source that retrieval can't
return and dedup (status != 'failed') blocks from re-ingest.

Only purge chunks when the document UPDATE actually transitions it to failed; an
already-completed document keeps its chunks. Adds reconcile regression tests for
both the completed-doc and genuine in-flight-orphan cases.

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

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

* Studio: drop a finished RAG job's queue when the client disconnects

job_events kept the per-job queue until it consumed the None sentinel, so a UI
that stops on the terminal event (its reader.cancel aborts the stream before
[DONE]) left the queue registered until the next _reap_finished_jobs sweep; a
batch of uploads followed by idling retained them all.

_run writes the terminal DB status before emitting the terminal event, so on
generator exit, drop the queue when the job's DB row is already terminal (worker
done, nothing to resume) and keep it only while the worker is still running. Adds
a disconnect-after-terminal-event regression test.

* Remove stray async task output files committed by mistake

* Studio: harden login IP throttle and end progress stream on disconnect

Two Codex review items:

Login per-IP throttle: when the per-IP bucket dict saturated, FIFO eviction could
drop a still-hot (blocked) bucket, so an IP could flood the dict with distinct
(or spoofed) source IPs to push out its own bucket and retry as first-seen. Stop
evicting hot buckets; a new IP that can't fit now shares a bounded overflow
counter that still trips the per-IP threshold, so a saturating spray stays
throttled and no live counter is reset.

Progress SSE: on client disconnect the polling loop only broke and fell through
to the unconditional final 'complete' frame, so a buffered or proxying consumer
could read a still-active run as completed. Return from the generator instead.

Adds regression tests for both (spray cannot reset a hot bucket; disconnect while
active emits no complete frame).

* Studio: shard the login overflow counter and stop cancelling chat stream after [DONE]

Two Codex review items:

Login throttle overflow: the single shared overflow counter meant that once a
saturating spray pushed it past the per-IP threshold, _login_blocked returned 429
for every new unbucketed source IP, before credentials were checked -- a global
login denial. Shard the overflow into a fixed array of counters keyed by hash(ip),
so a hot shard only throttles the IPs that map to it while a single source's
repeated failures still concentrate in one shard and stay throttled. Memory stays
bounded and no live bucket is evicted. Adds a regression test that a hot overflow
shard does not block an unrelated IP.

Chat stream: the reader.cancel() in the SSE finally fired even after a natural
[DONE]/EOF. The backend finalizes its api-monitor entry right after yielding the
sentinel (the local pass-through finishes after the last yield), so a client
cancel there can be observed as a disconnect and mark a completed request as
cancelled. Track natural completion and only cancel on an early/abnormal exit.
(No frontend unit test: the Studio frontend has no test harness.)

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

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

* Studio: give prep-timeout test fakes an is_disconnected method

The progress stream now ends on client disconnect (await request.is_disconnected()
before falling through to the terminal frame). After merging that into the
prep-timeout tests added later on main, their _FakeRequest/_ReconnectRequest must
provide is_disconnected or the generator raises AttributeError under CI.

* Studio: keep the login overflow throttle when bucket capacity frees up

_login_blocked only consulted the per-IP overflow shard while the bucket dict was
still at capacity. If a slot freed before the 60s window expired (e.g. another
IP's successful login calls _clear_login_bucket), a source counted in a hot shard
stopped being blocked and its next failure got a fresh per-IP bucket, resetting
the throttle the overflow path exists to preserve. Always max in the IP's shard
(shards are empty outside saturation, so it is a no-op in the common case). Adds a
regression test that a hot source stays throttled after a bucket frees.

* Studio: clear a login IP's overflow throttle on successful login

_clear_login_bucket reset the per-IP and per-account buckets on a successful
login but not the overflow shard, so after the dict saturated and an IP was
counted in overflow, a later successful login left those entries behind and the
next failed attempt could immediately return 429.

Store overflow entries as (timestamp, ip) so a source is throttled by its own
count within the shard (also removing cross-IP collateral within a shard), and
drop just that IP's entries in _clear_login_bucket. Adds a regression test that a
successful login clears the overflow throttle.

* Studio: bound the login overflow shard memory under high-cardinality spray

The per-IP overflow tracked failures in a time-pruned deque of (timestamp, ip)
tuples, so a spoofed-X-Forwarded-For spray of distinct one-off IPs grew memory and
the per-check scan with request cardinality for the whole window -- undermining
the bucket cap that exists to bound memory. Replace each shard with a fixed-
capacity dict (ip -> [count, window_start]): O(1) lookups, and when a shard is
full a one-off IP evicts the lowest-count entry (Space-Saving) so memory is hard-
bounded while a persistent attacker keeps a high count and is never evicted. Adds
a regression test that shards stay within the per-shard cap under a 5000-IP spray.

* Studio: purge chunks for already-failed docs during ingestion reconcile

The reconcile chunk-purge was gated on the documents UPDATE actually flipping a
non-terminal doc to failed. A doc the worker had already marked 'failed' before
the crash (job row left non-terminal) was not re-flipped, so its committed chunks
were kept and stayed retrievable/citable, since retrieval filters by scope not
status. Purge chunks whenever the document is not 'completed' (failed, in-flight,
or gone), preserving the completed-doc carve-out. Adds a regression test.

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

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

* Studio: don't inherit an evicted IP's count onto a new overflow source

When a full overflow shard evicted the lowest-count entry, the new source
inherited that count (Space-Saving base + 1). If a shard was saturated with hot
entries, an unrelated new IP could land at/over the threshold and be 429'd after a
single attempt -- cross-IP collateral despite the per-source-isolation intent.
New entries now start clean at count 1; the only cost is that a heavy hitter that
is the lowest-count entry in a fully saturated shard can briefly reset, which is
preferable to blocking a bystander. Adds a regression test.

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

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

* Studio: carry overflow failures into a new IP bucket on transition

_login_blocked took max(per-IP bucket, overflow shard) rather than combining them,
so a source could log (threshold-1) failures in overflow during saturation and,
once a bucket slot freed, another (threshold-1) in a fresh bucket within the same
window -- roughly doubling the per-IP limit. When a saturated-era IP first gets a
real bucket, migrate its windowed overflow count into that bucket (and drop the
overflow entry) so the combined failures throttle at the intended limit. Adds a
regression test.

* Studio: reconcile a completed doc's orphaned job to completed, not failed

When a crash left an ingestion job non-terminal after its document was already
committed as completed, reconcile marked the job failed. After restart the upload
UI has no in-memory SSE queue and falls back to getJob(), which treats a failed
job as an indexing failure and removes/toasts a document that is actually
searchable. Mark the job completed (keeping its chunks) when its document is
completed. Extends the completed-doc reconcile test to assert the job status.

* Studio: clamp the overflow failure count migrated into a login bucket

A saturated source could accrue an unbounded overflow count, then materialize
one deque entry per recorded failure when a bucket slot freed, allocating an
arbitrarily large deque under the login lock. Only at-or-above the per-IP
threshold matters for blocking, so cap the count there at the record and take
sites; the migration is now bounded without weakening the limit.

* Studio: keep the RAG job stream alive on a transient status read

The heartbeat poll read the job row unguarded; a momentarily-locked DB would
raise out of job_events, which the SSE route turns into a terminal error frame,
and the UI drops a document whose worker is still running. Treat a failed status
read as non-terminal: heartbeat and retry, and keep the queue so a reconnect can
resume.

* Studio: set busy_timeout before journal_mode on the auth DB

Switching journal_mode needs a lock, so if a refresh-token write already holds
one, journal_mode=WAL raises SQLITE_BUSY and the shared try leaves the
connection on SQLite's default zero lock wait. Set busy_timeout first so the
switch waits instead of failing.

* [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-06-26 03:31:33 -07:00
Matt Van Horn
e9c6364e1e
feat: improve Unsloth Studio chat title generation quality (#6697)
* feat: improve Unsloth Studio chat title generation quality

* fix: address self-review (guard echoed role labels before punctuation stripping)

* Address title generation review feedback

Consolidate the echo guard into a single leading-label check (now also
covering base and lora) and drop the post-punctuation duplicate that
could never match a colon once punctuation is stripped. Swap the
slice-based first-assistant lookup for an indexed find to avoid copying
the messages array, and note the brace counter's assumptions in the
test helper.

---------

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Co-authored-by: Daniel Han <unslothai@gmail.com>
2026-06-26 01:55:59 -07:00
Daniel Han
7f45635280
Studio: auto-shut-down an exposed first-run instance if the admin password is never changed (#6651)
* Studio: set the admin password before exposing it on the network

On first run Studio seeds the default `unsloth` admin with a random
bootstrap password and embeds it into index.html (window.__UNSLOTH_BOOTSTRAP__)
so the local user can change it without typing it. A request with no Origin
header counts as same-origin, which is what a normal top-level GET sends, so
the page hands out the password to whoever loads it. That is harmless on the
default 127.0.0.1 bind, but `--secure` (public Cloudflare tunnel) and
`--host 0.0.0.0` (raw port reachable on the network) would serve the plaintext
admin password to remote visitors during the bootstrap window.

Fix this at the source: when launching a network-exposed web UI, prompt the
operator in the terminal for a real admin password (with confirmation) before
the socket binds or the tunnel opens, and persist it via update_password (which
clears must_change_password and deletes the .bootstrap_password file). After
that there is no bootstrap secret to leak. Non-interactive launches can supply
it via UNSLOTH_STUDIO_ADMIN_PASSWORD. The masked reader echoes '*' per
character and works on Linux, macOS, and Windows (PowerShell/cmd). Loopback
binds, --api-only (no web UI), and Colab are unaffected.

As defense in depth, the index handler now embeds the bootstrap object only for
a direct local navigation: same-origin AND a loopback TCP peer with no
proxy/tunnel forwarding headers (cf-ray, cf-connecting-ip, x-forwarded-for,
x-forwarded-host, x-real-ip, forwarded). Colab stays exempt. This keeps the
password off the wire even when the prompt is skipped (no TTY and no env var).

Adds unit coverage for the prompt/confirm/decision logic, an integration test
that provisioning clears the bootstrap state, and regression tests for the
local-direct gate (loopback/IPv6/mapped/localhost peers, LAN/public peers,
missing client, each forwarding header, spoofed XFF, and the Colab exemption).

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

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

* Studio: fail fast on an explicitly empty admin-password env var

resolve_admin_password_source treated UNSLOTH_STUDIO_ADMIN_PASSWORD="" like
the var was unset and fell back to the bootstrap backstop. Treat any set value
(including empty) as the env source so it reaches the minimum-length guard and
refuses to expose the server instead of silently keeping the seeded password.

* Studio: apply repo kwarg-spacing format to the secure-admin-password files

* Studio: drop the pre-exposure password prompt; keep the local-direct gate

Per review, the blocking prompt added friction for --secure / 0.0.0.0 first-run
launches without extra security: the local-direct injection gate in main.py
already keeps the bootstrap password off the network for any remote request.
Remove the prompt module and its tests; the gate plus the existing
must_change_password first-login flow are the fix.

* Studio: shut down an exposed first-run instance if the admin password is never changed

The local-direct gate keeps the seeded bootstrap password off the network, but
it stays a valid credential until first login changes it. For an exposed web UI
(--secure / 0.0.0.0, not --api-only, not Colab), arm a daemon timer: if the
password is still the seeded one after the deadline (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT,
default 3600s, 0 disables), print a message and shut Studio down via the existing
graceful-shutdown path; if it was changed, leave Studio running.

* Studio: revert the local-direct injection gate; keep the 1-hour auto-shutdown

Per maintainer decision, keep the first-run auto-fill behavior unchanged (the
bootstrap password still seeds the login form for convenience) and rely on the
exposed-instance auto-shutdown to bound the window: an exposed web UI that never
changes the seeded admin password is torn down after UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT
(default 1h). Restores studio/backend/main.py and its origin test to upstream.

* Studio: render the bootstrap-timeout shutdown message with a human duration

The message hardcoded 'minute(s)' via timeout//60, so a sub-minute timeout
(e.g. a 30s test value) printed 'within 1 minute(s)'. Add _format_duration so
it reads '30 seconds' / '1 minute 30 seconds' / '60 minutes' as appropriate.
The default 3600s still renders '60 minutes'.

* Studio: drop stale local-direct gate reference from bootstrap_timeout docstring

The gate was reverted (timer-only), so the module docstring should not describe
a main.py gate that no longer exists.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-26 01:27:27 -07:00
Leo Borcherding
c86165e735
Studio Colab: opt-in shareable Cloudflare tunnel link (#6684)
* Studio Colab: add opt-in shareable Cloudflare tunnel link

colab.start(cloudflare=True) opts in to a free Cloudflare quick tunnel and
shows a trycloudflare.com link above the proxy iframe, reachable from any
device. Default OFF: bare start() keeps the in-tab Colab-proxy behavior.
run_server suppresses the tunnel on Colab by design, so colab.py starts it
directly via cloudflare_tunnel.start_studio_tunnel(); failures degrade to
the Colab proxy only.

* Studio Colab notebook: surface opt-in cloudflare=True in start cell

* Studio Colab: reskin shareable Cloudflare link to match the proxy banner

Retrofit _shareable_link_html to reuse the original Colab proxy banner skin
from show_link (white card, black border, Unsloth gem, black Open button)
instead of the plain dark box, so the shareable Cloudflare link gets the same
prominent 'Ready!' treatment.

* Studio Colab: address review feedback on Cloudflare tunnel

- try/finally around tunnel start + embed + keepalive so a KeyboardInterrupt
  while the tunnel is starting or the iframe is rendering tears it down instead
  of orphaning the cloudflared process (Gemini review).
- Publish the directly-started tunnel URL onto app.state.cloudflare_url via a new
  _publish_cloudflare_url helper so /api/health advertises it; otherwise the
  frontend's API examples fall back to the unreachable raw server_url (Codex P2).
  _stop_cloudflare_tunnel now also clears it so health stops showing a dead tunnel.
- Notebook: make cloudflare=True a replacement for start(), not an addition, since
  start() blocks and the second call would never run if both are left in (Codex P2).

* Studio Colab: gate Cloudflare tunnel on auth + honor opt-out in run_server

- Refuse to open the Cloudflare tunnel while the admin still holds its seeded
  bootstrap password. While requires_password_change is true the server injects
  that password into same-origin index GETs, and a public tunnel request counts
  as same-origin, so sharing the link would leak admin access. New
  _bootstrap_password_pending() gate (fails safe) blocks the tunnel and tells the
  user to change the password first, then re-run start(cloudflare=True) (P1).
- Pass cloudflare=False into run_server so the opt-out holds even when Colab
  detection fails; this helper is now the sole owner of the tunnel decision,
  preventing run_server from opening a tunnel on the 0.0.0.0 bind by default (P2).

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

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

* Studio Colab: drop duplicate tunnel link log and simplify start cell guidance

* Studio Colab: validate /api/health identity before reusing or tunneling a port

* Studio Colab: condense verbose docstrings and comments

---------

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-06-26 00:56:23 -07:00
Daniel Han
1396c01253
Fix offline checkpoint load/export: "tokenizer is weirdly not loaded" (#6554)
* Fix offline checkpoint load/export failing with "tokenizer is weirdly not loaded"

Loading a fine-tuned checkpoint with no internet (e.g. a Studio export) crashed
with "Unsloth: The tokenizer is weirdly not loaded? Please check if there is one."

For a LoRA adapter the loader reassigns model_name to the base model repo id and
only keeps the local checkpoint dir as tokenizer_name when it contains
tokenizer_config.json, tokenizer.json AND special_tokens_map.json. Modern
tokenizers (e.g. Gemma) store special tokens inside tokenizer_config.json and
omit special_tokens_map.json, so tokenizer_name fell back to the base repo id.
The tokenizer/processor loads in vision.py then hit the Hub with no
local_files_only, so with no network they failed (AutoProcessor) or hung for
minutes (AutoTokenizer) even though every file was already cached.

loader.py: keep the local checkpoint dir as tokenizer_name when it has a
tokenizer config plus the actual tokenizer files (tokenizer.json / tokenizer.model
/ vocab files); special_tokens_map.json is no longer required.

vision.py: compute an effective local_files_only (explicit kwarg plus the
HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE env vars, mirroring loader.py and
diffusion.py) and thread it through every AutoConfig, AutoProcessor,
AutoTokenizer and the manual VLM processor fallback, including the
hf_hub_download in that fallback (which now prefers a local file). When a load
fails and no offline env var is set, retry against the local cache. The retry
forces HF offline mode because local_files_only alone does not stop
AutoProcessor / AutoTokenizer from issuing a /api/models request during class
resolution. The final error now explains the offline/cache cause instead of the
misleading "weirdly not loaded" message.

studio export: probe Hub reachability once per checkpoint load and pass
local_files_only when offline so exports use the local checkpoint dir / cache
instead of hanging or crashing with no internet.

Online behavior is unchanged: the new flags default to off and the retry only
runs after a network related failure.

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

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

* Address review: safer offline forcing, cached fallback config, proxy-aware probe

Follow-up to the offline checkpoint load fix, addressing review feedback:

- vision.py: only flip the process-wide HF offline flag when offline is actually
  requested (local_files_only / env) or after a real network failure, never
  pre-emptively while we might be online. The flip is now guarded by a lock +
  depth counter so nested or concurrent windows restore the flag correctly
  (no stale value).
- vision.py: guard the get_auto_processor fallback so a network error there
  returns None and the local-cache retry still runs instead of escaping.
- vision.py: in the manual VLM processor fallback, read tokenizer_config.json
  via hf_hub_download(..., local_files_only=...) so a cached repo-id config is
  still resolved offline and the model-specific image/video tokens are restored.
- studio export: make the reachability probe proxy aware (probe the configured
  HTTP(S) proxy egress, honour NO_PROXY, use the endpoint port) so a proxy-only
  setup is not wrongly marked offline; allow UNSLOTH_OFFLINE_PROBE=0 to disable.
- studio export: run the audio/vision type-detection probes inside the
  forced-offline window when offline, so their config/tokenizer reads hit the
  local cache instead of waiting out connection timeouts.

Online behavior remains unchanged.

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

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

* Address review: gate offline retry, safer tokenizer_name pop, skip audio net probe offline

- vision.py: only force the process-wide HF offline flag on the tokenizer
  retry when offline was requested or the captured primary error is actually
  network related, so a permanent tokenizer error no longer toggles global
  offline mode for other concurrent loads.
- loader.py: always pop tokenizer_name out of kwargs and let a caller-supplied
  value win, avoiding a "multiple values for keyword argument 'tokenizer_name'"
  TypeError when it is also passed explicitly downstream.
- model_config.py / export.py: add local_files_only to detect_audio_type so the
  raw requests.get tokenizer_config fetch is skipped offline (it ignores the HF
  offline flag), and pass it from the export probe.

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

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

* Address review: classify LocalEntryNotFoundError as offline-related

huggingface_hub's LocalEntryNotFoundError subclasses FileNotFoundError, so the
"not isinstance(cur, FileNotFoundError)" guard in _is_offline_related_error was
swallowing it and it could never be recognised as offline, despite being listed
in the network error types. It means "not in cache and the Hub is unreachable",
which is genuinely offline. Capture the class into an isinstance-checkable tuple
(empty, hence a no-op, if the import is unavailable) and exclude it from the
FileNotFoundError guard, so a real offline failure now triggers the local-cache
retry while a plain missing-file error still propagates.

* Address review: require merges.txt for BPE, status-gate HTTP errors, isolate local-only audio cache

- loader.py: a local dir with vocab.json but no merges.txt (and no tokenizer.json)
  is not a loadable BPE tokenizer, so do not treat it as self-sufficient; require
  merges.txt alongside vocab.json in both gate blocks, otherwise fall back to the
  base model tokenizer as before.
- vision.py: _is_offline_related_error no longer buckets every HfHubHTTPError /
  requests HTTPError as offline. HTTP errors are judged by status code: only a
  transient 5xx triggers the forced local-cache retry, while 401/403 (auth/gated)
  and 404 (missing) propagate as the real error instead of being masked. Hard
  signals (connection/timeout/OfflineModeIsEnabled/LocalEntryNotFoundError) still
  classify as offline.
- model_config.py: include local_files_only in the audio-detection cache key so a
  local-only (offline) negative result cannot be reused by a later online probe,
  which would otherwise route an audio model through the text loader until restart.

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

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

* Address re-review: fix studio test stubs, force offline env in probe window, drop redundant retry

- studio/backend/tests/test_vision_cache.py: the three _detect_audio_from_tokenizer
  stubs were called with the new local_files_only kwarg and raised TypeError, failing
  Backend CI. Add local_files_only to the stub signatures and add a test that a
  local-only negative does not poison a later online audio probe.
- export.py: the type-detection probe window now also sets HF_HUB_OFFLINE /
  TRANSFORMERS_OFFLINE env vars (saved/restored), not just the in-process flag.
  transformers_version._load_config_json / _check_tokenizer_config_needs_v5 gate
  their urllib fetches on the env vars, and is_vision_model may spawn a subprocess
  that inherits os.environ but not the in-process flag; without the env vars a
  probe-detected offline export could still block on a network timeout.
- vision.py: only retry the processor load when the first attempt was online and
  failed with a network error. When local_files_only was already requested the first
  attempt was forced offline, so the previous retry just repeated identical failing
  work before the last-resort path.
- model_config.py: correct the _audio_detection_cache type annotation to the 3-tuple
  key (name, token_fingerprint, local_files_only).

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

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

* Address review: thread-safe probe-offline env window, clear error for local dir without config

- export.py: guard the HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE mutation in
  _force_offline_probe_window with a lock + depth counter (mirrors _force_hf_offline),
  so concurrent / nested export probes only flip on first entry and restore on last
  exit. This prevents overlapping export requests from permanently poisoning those
  env vars or restoring a stale value.
- vision.py: in the VLM processor fallback, when tokenizer_name is a local directory,
  read its tokenizer_config.json directly and raise a clear FileNotFoundError if it is
  absent, instead of handing the local path to hf_hub_download (which would treat it as
  a repo id and raise a confusing HFValidationError / RepositoryNotFoundError).
  hf_hub_download is now only used for actual repo ids.

* Address review: classify raw socket.gaierror DNS failures as offline

Add the platform-specific getaddrinfo / DNS-resolution wording to the offline
detection list in _is_offline_related_error so a bare socket.gaierror (an OSError
subclass) is recovered from the local cache: "Name or service not known" and
"Temporary failure in name resolution" (Linux) and "nodename nor servname
provided" (macOS). Genuine non-network OSErrors (disk full, permission denied)
and plain FileNotFoundError still propagate.

* Address review: retry degraded VLM offline, force offline for text export + patch-tokenizer fallback

- vision.py: a degraded VLM processor (text-only, no image_processor) whose manual
  fallback fails offline used to be kept, so image inputs broke even with cached
  files. _construct_vlm_processor_fallback now returns its failure error;
  _acquire_processor surfaces it, and the caller retries forced-offline when the
  result is None OR a degraded VLM and the failure was network related, keeping the
  original result if the retry is not strictly better (never regress). The retry is
  still gated on an online first attempt + offline-related error so a permanent
  error never flips the global offline flag.
- vision.py: wrap the patch_tokenizer except-branch AutoTokenizer.from_pretrained in
  the same forced-offline-on-network-error pattern as the primary / last-resort
  loads, so an offline export where patch_tokenizer raises does not hang or fail.
- export.py: force HF offline around the two FastLanguageModel loads (text and SNAC)
  when the probe detected offline. Their text tokenizer path (load_correct_tokenizer
  -> AutoTokenizer) does not forward local_files_only, so without this a text export
  could still contact the Hub. Added a small _offline_window_if helper reused by the
  probe and load windows.

* Consolidate offline loading into one entry-point decision

Decide offline once per entry point instead of at every HF call site. The
prior approach threaded local_files_only into ~15 scattered config / tokenizer
/ processor / weight loads, each wrapped in its own try-online, classify-error,
retry-forced-offline dance, which is what kept surfacing "another call site you
missed", "another error shape misclassified", and global-flag thread-safety in
review.

FastLanguageModel / FastModel / FastBaseModel.from_pretrained now share an
@_offline_aware_load decorator: when offline (explicit local_files_only kwarg or
HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE env) it sets local_files_only and runs the
whole load inside one _force_hf_offline() window so every nested HF call inherits
it; when online it runs normally and, only if the load fails with a genuinely
network-related error, retries once forced-offline. The online path is unchanged
(no probe added) and 401 / 403 / 404 / permanent errors still propagate.

Centralise the offline helpers in loader_utils.py as the single source of truth
(shared by loader.py, re-exported from vision.py, and reused by the Studio
exporter):
- _force_hf_offline now sets the HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE env vars
  AND the in-process huggingface_hub / transformers flags, refcounted under one
  lock so nested / concurrent windows restore correctly. Setting the env vars
  covers env-gated urllib probes and spawned subprocesses too.
- _get_effective_local_files_only, _is_offline_related_error (unchanged
  classifier, retains the 5xx-vs-4xx, LocalEntryNotFound and gaierror handling),
  _offline_aware_load, and _resolve_checkpoint_tokenizer_name.

loader.py: wrap both entry points; drop the two duplicated env-var fallback
blocks and the two byte-identical local-tokenizer-gate blocks (now
_resolve_checkpoint_tokenizer_name).

vision.py: drop the per-site force_offline params and the three retry gates
(processor, patch_tokenizer fallback, last-resort). They now just surface the
underlying error so the single entry-point safety net retries forced-offline. A
network fallback error now takes precedence over a permanent primary error so the
offline retry still fires when the manual VLM fallback needs cached repo files.

studio/backend export.py: reuse the unified core _force_hf_offline (env + flags)
and drop the duplicate probe-window primitive; the snac / text branches no longer
need their own window. model_config.py: also gate the raw requests.get audio
fallback on the HF offline env vars so it is covered even without the kwarg.

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

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

* Address 10-reviewer P1 findings: vision cache split, PEFT offline, retry OOM

Split the Studio vision-detection cache by local_files_only, mirroring the audio
cache fix. is_vision_model / _is_vision_model_uncached / _raw_config_has_vision_config
/ load_model_config now thread local_files_only, the cache key includes it, and the
exporter passes it. Offline detection also skips the transformers-5 network
subprocess and stays on the local cache, so an offline negative can no longer be
keyed under the online entry and poison a later online probe. Adds a regression
test mirroring the audio poison test.

Forward local_files_only to both PeftModel.from_pretrained adapter-attach sites in
loader.py so a cached remote LoRA adapter resolves from the local cache under
explicit local-only / offline loads (defence-in-depth alongside the forced-offline
window).

_offline_aware_load: run the forced-offline retry OUTSIDE the except block and
collect + empty the device cache first. An except-scoped exception keeps its
__traceback__, which pins the failed attempt's frame locals (a partially loaded
model) until the block exits; loading the model again while that copy is still
alive could OOM a large VLM. Letting the except block close drops the traceback so
the partial load is freed before the retry reallocates.

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

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

* Address Codex review: env-offline cache key + rebuild HF sessions in offline window

Key the Studio audio and vision detection caches on the EFFECTIVE offline state
(local_files_only OR the HF offline env vars), not just the kwarg. detect_audio_type
and is_vision_model both skip the remote fetch / network subprocess when
HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE is set even with the default
local_files_only=False, so the result reflects offline; storing it under the online
(False) key let an env-offline negative poison a later online lookup once the env var
was cleared. Both now compute effective_offline once and use it for the cache key and
the downstream call. Adds a regression test for the env-offline dimension.

_force_hf_offline now rebuilds huggingface_hub's cached sessions on enter and exit
(best-effort _reset_hf_sessions). On hub 0.x the offline adapter is baked into the
per-thread requests.Session at creation, so flipping the constant alone leaves an
already-cached online session able to hit the network inside the window (and an
offline one stuck offline after restore); resetting forces the next get_session() to
match the current flag. On hub 1.x offline is checked dynamically per request, so
reset_sessions does not exist and the helper is a safe no-op.

The third review point (release the failed load before retrying) was already fixed in
af0f58a: the forced-offline retry now runs outside the except block and frees the
device cache first, so the failed attempt's traceback-pinned partial model is
released before the retry reallocates.

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

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

* Align Studio _env_offline parsing with the canonical offline helper

model_config._env_offline gates the raw requests.get tokenizer-config fallback in
detect_audio_type and the audio/vision detection cache keys, but it only accepted
unstripped "1"/"true"/"yes". unsloth's offline helpers (loader_utils._env_says_offline
and the from_pretrained env fallback) accept the canonical set {1,true,yes,on} after
strip + lowercase, so HF_HUB_OFFLINE=on or HF_HUB_OFFLINE=" 1 " was treated as offline
by the loaders but online here, leaving the raw network fetch reachable while
"offline". Use the same strip + lowercase {1,true,yes,on} set. Adds parsing tests.

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

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

* Fix lint: drop dead offline-helper re-exports from vision.py

The import-hoist verifier (scripts/verify_import_hoist.py) flagged vision.py's
re-export block as HOISTED-IMPORT-UNUSED blockers: it imported eight offline
helpers from loader_utils but only used three internally
(_get_effective_local_files_only, _is_offline_related_error, _offline_aware_load).
The other five were imported purely to preserve `from unsloth.models.vision import
X`, but nothing imports four of them from vision, and loader.py already imports
_resolve_checkpoint_tokenizer_name straight from loader_utils.

Import only the three names vision.py actually uses, and point the Studio exporter
at the canonical source (from unsloth.models.loader_utils import _force_hf_offline)
instead of re-exporting it through vision. loader_utils stays the single source of
truth; no behaviour change.

* Address Opus review: chain probe errors, unify env-offline, status-less HTTP

Chain the original AutoConfig/PeftConfig probe exception into the combined
RuntimeError in both FastLanguageModel.from_pretrained and FastModel.from_pretrained
(`raise RuntimeError(combined_error) from (autoconfig_exc or peft_exc)`). The probes
caught every Exception and stringified it, so the re-raised RuntimeError had no
__cause__/__context__ and _is_offline_related_error could not classify it -- the
network-down-but-cached auto-retry never fired for these entry points. With the
cause chained, the decorator sees a ConnectionError/LocalEntryNotFoundError/5xx and
retries forced-offline from cache; a permanent cause (404 / bad config) is still not
offline-classified and propagates without a wasted retry.

Unify the third offline-env parser: studio/backend/utils/transformers_version._env_offline
now uses the canonical {1,true,yes,on} + strip + lowercase set (matching
loader_utils._env_says_offline and model_config._env_offline), so HF_HUB_OFFLINE=on
or " 1 " no longer leaks the direct urllib metadata fetches to the network.

_is_offline_related_error: a status-less HTTP error (no response / unparseable code)
now falls back to the network-wording check instead of being dropped, so a transient
HTTP failure with clear "couldn't connect" wording is treated as offline. HTTP errors
with a real status code still decide by code (4xx propagates, 5xx is offline).

* Condense offline-loading code comments, drop dead helper, dedupe import for PR #6554

* Add unit tests for offline-loading helpers for PR #6554

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

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

* Guard load cleanup with try/finally and add retry-contract tests for PR #6554

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

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

* Add gc.collect retry-step test for PR #6554

* Tighten offline-loading comments and docstrings for PR #6554

* Raise the both-config-failed error before model-type lookup so offline retry fires for PR #6554

* Prefer offline cause for retry and bound export reachability probe for PR #6554

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

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

* Skip remote mapper while offline, harden text-load cleanup, and stop stacked offline retries for PR #6554

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

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

* Surface VLM fallback offline errors, probe offline before export version activation, and restore progress bars across retries for PR #6554

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

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

* Restore offline env after export version activation so the persistent worker re-decides per load for PR #6554

* Classify socket.gaierror and urllib URLError as offline by type for PR #6554

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

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

* Probe offline around export load preflights and never offline-retry TLS failures for PR #6554

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

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

* Force in-process offline for export preflights, verify proxy egress in probe, and skip caching offline version negatives for PR #6554

* Snapshot offline constants before forcing env and require local processor files for VLM checkpoints for PR #6554

* [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-06-25 23:16:53 -07:00