Commit graph

500 commits

Author SHA1 Message Date
Daniel Han
a35fbe22ea
Studio: serialize non-streaming responses once and pool the proxy client (#6393)
* Studio: serialize non-streaming responses once and pool the proxy client

Two safe latency wins on the OpenAI/Anthropic-compatible endpoints that leave
the streaming generation paths untouched (they keep Connection: close and
max_keepalive_connections=0 so a client disconnect still stops GPU decode).

1. Non-streaming responses used JSONResponse(content=model.model_dump()), which
   builds a dict and then re-runs json.dumps. Serialize once with
   model.model_dump_json() via a small _model_json_response helper. The body is
   byte-identical (nulls preserved), about 3x faster to encode in a microbench.

2. The non-streaming completions and embeddings proxies built a fresh
   httpx.AsyncClient per request. Route them through one pooled client
   (core/inference/llama_http) closed on shutdown; streaming generation keeps
   its own per-request close-only client. About 5x faster per call to the
   local llama-server in a microbench.

The existing API-monitor tests for the non-streaming completions, embeddings
and passthrough paths now patch nonstreaming_client instead of httpx.AsyncClient
to match the pooled client, so they stay deterministic.

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

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

* Studio: make the pooled non-streaming client per event loop

Review follow-up on the shared httpx client. It was a single module-global
instance, which has two lifecycle problems the per-request client did not:

1. After aclose() in lifespan shutdown, nonstreaming_client() kept handing back
   the closed client, so a second lifespan in the same process (repeated
   TestClient, embedded restart) failed with "client has been closed".

2. An httpx client binds its transport to the loop it first runs on, so reuse
   from another loop could raise "Event loop is closed".

Hold one client per running loop in a WeakKeyDictionary, recreate when missing
or closed, and close all on shutdown. Single-loop production is unchanged.

* [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-17 22:38:02 -07:00
Daniel Han
9a966adf51
Studio: trim serving-log noise and surface llama-server engine stats (#6377)
* Studio: trim serving-log noise and surface llama-server engine stats

Studio prints one structured line per HTTP request, so the SPA's polling and
per-invalidation fan-out bury the lines that matter.

- Dedup identical successful GETs within a short window (default 300ms,
  UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS) so a burst logs once. The dedup key
  includes the query string, so distinct query-driven GETs are not collapsed.
  Runs after the response is sent, so it adds no request latency; mutations,
  non-2xx, and loading polls are untouched.
- Collapse pure-liveness polls (/api/health, /api/auth/status,
  /api/inference/status, /api/inference/monitor) to a longer heartbeat
  (default 10s, UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS). The API monitor
  console polls /monitor every 1.5s while open.
- Translate llama-server's Prometheus /metrics into a periodic vLLM-style
  engine_stats line (generation/prompt throughput and requests in flight) from
  a daemon poller, gated on UNSLOTH_STUDIO_ENGINE_STATS. Throughput uses
  llama-server's predicted_tokens_seconds / prompt_tokens_seconds gauges, with
  a tokens_predicted_total / prompt_tokens_total counter-delta fallback; it does
  not use n_decode_total (which counts llama_decode() calls, not tokens). No KV
  field is emitted, since llama.cpp does not expose kv_cache_usage_ratio.
  --metrics is added only when probe_server_capabilities reports the binary
  supports it, so older/custom binaries still load. The poller keeps retrying
  through transient scrape failures (stop() drives shutdown) and a malformed
  sample cannot crash its thread.
- api_monitor.append_reply: once the preview cap is reached, skip the per-chunk
  re-concat (avoids O(n^2) on long generations) while still recording the "..."
  truncation marker for a reply that lands exactly on the cap.
- unsloth studio --verbose and unsloth studio run --verbose both restore every
  per-request log; --verbose before a subcommand is rejected with guidance
  (matching --secure / --parallel). run --verbose still forwards --log-verbose
  to llama-server, preserving the pre-existing pass-through verbosity.

* [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-17 05:37:57 -07:00
Daniel Han
0e69614878
studio: deterministic VRAM auto-fit for GGUF (MTP reserve, compute buffer, total-based budget) (#6312)
* studio: reserve MTP draft VRAM in GGUF auto-fit

Auto-fit advertised a context (for example ~110k for the Qwen3.6-27B MTP
GGUF) that fit on paper but OOMed mid-generation or during tool calls once
MTP speculative decoding was active. The MTP draft path's VRAM was reserved
as a flat 5% of total VRAM, which tracks neither of the two real costs: the
MTP head keeps its own attention KV cache that grows with context, and the
speculative verification buffer grows with --spec-draft-n-max. On the
hybrid Mamba/attention Qwen3.6 models the main KV is small, so auto-fit
happily kept a near-native context while the draft path pushed the load
over budget at runtime.

Replace the flat fraction with a byte-accurate, context- and n_max-aware
reserve sized from GGUF dims: draft KV from nextn_predict_layers and the
attention dims at f16 (llama.cpp's MTP draft context uses f16 KV regardless
of the main cache type), plus a verify buffer per embedding-unit per draft
token. The reserve is evaluated per candidate context inside the fit binary
search and added to every pin/fit check, including the tensor-parallel
planner and its even-split decision. Coefficients were calibrated against
llama-server VRAM measurements on the Qwen3.6-27B MTP GGUF (RMS 14 MiB).

The flat fraction remains as a fallback when GGUF dims are unavailable, so
non-MTP loads are unchanged. The budget now also engages when the user wires
MTP through extra args (--spec-type draft-mtp, including chains), reads the
effective draft depth from --spec-draft-n-max or the legacy --draft-max with
extras taking precedence over the first-class field, reserves a separate
drafter's weights when supplied via --model-draft/--spec-draft-model/-md,
and mirrors _build_speculative_flags so it never reserves for MTP the launch
resolver will not emit (needs a head/drafter and a binary that supports
--spec-type mtp).

Adds tests/test_mtp_vram_budget.py.

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

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

* studio: total-based VRAM budget + deterministic compute-graph buffer

Build on the byte-accurate MTP reserve with three changes that make the
GGUF auto-fit budget deterministic across architectures and recover usable
context, especially for MTP models on a single tight card.

1. Total-based budget. Cap GPU occupancy at a fraction of TOTAL VRAM rather
   than a fraction of FREE VRAM, and raise the fraction from 0.90 to 0.95:

       budget = free - (1 - 0.95) * total      (per GPU, summed for a pool)

   The reserve is now absolute (a fixed slice of the card) instead of
   shrinking as the GPU fills, so a partly-used GPU keeps a constant cushion
   for compute/CUDA/verify buffers instead of over-promising context and
   spilling to CPU at runtime. _get_gpu_memory() reads memory.total alongside
   memory.free; _fit_context_to_vram, _select_gpus and the load_model pool
   loops thread the totals through. Multi-GPU layer-split pools
   sum(free_i - 0.05*total_i); tensor mode reserves per device.

2. Deterministic compute-graph buffer. Replace the flat 5 GB/device tensor
   reserve (a magic constant that over-reserved about 8x on a 27B model) with
   _estimate_compute_buffer_bytes, sized from GGUF dims and the launch flags:

       out = n_vocab * n_ubatch * 4            # vocab-width output buffer
       act = 4 * n_embd * n_ubatch * 4         # activation scratch
       pipeline_per_device = act + out * (n_parallel - 1)
       tensor_per_device   = 2*act + out * n_parallel

   The buffer is context-independent and scales with --parallel (serving
   slots), not with how the model is split across GPUs. It is now reserved in
   BOTH multi-GPU paths (layer split folds one buffer into the pooled
   footprint; tensor mode reserves it per device). The flat 5 GB stays only as
   a fallback when vocab/embedding dims are unavailable. Calibrated against
   llama-server measurements (parallel 1/2/4/8 give 36/492/1388/3220 MiB on a
   single GPU; about 600 MiB/device tensor); the estimate is a small upper
   bound.

3. GGUF parsing. Read vocab size (tokenizer tokens array length) and
   feed_forward_length for the compute-buffer estimate.

Effect on the Qwen3.6-27B MTP Q6_K case (MTP on): a single 32 GB card at
about 31 GB free advertises f16 23k to 64k, q8_0 44k to 115k, q4_0 82k to
200k; 2x 24 GB tensor mode recovers the full 262k window for f16 (was about
134k). Validated on hardware: 1x 32 GB f16 at 64768 loads at 29.3 GB / 120
t/s; 2x 23 GB tensor f16 at 262144 loads at 22.2 GB/device / 98 t/s; both
within 0.4% of the estimate. Adds test_compute_buffer.py and updates the
KV/context-fit/MTP-budget tests for the 0.95 constant and the new budget.

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

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

* studio: tighten comments in the VRAM auto-fit changes

Condense the docstrings and inline comments added by this PR (internal backend
helpers): drop restated-signature docstrings, fold multi-line block comments to
one or two lines, and remove notes that just repeat the code. No behavior change
(AST-verified comment/docstring-only via comment_tools.py); the backend test
suite is unchanged and green.

* studio: address review findings in the VRAM auto-fit budget

Five fixes from a parallel-reviewer pass on this PR; all confirmed against the
real functions and covered by new tests.

- Tensor mode now honors the total-based VRAM cap. _plan_tensor_parallel took
  total_by_idx and budgets each GPU at free - (1-frac)*total, mirroring the
  layer-split paths; previously it fit against raw free and could spend the 5%
  safety cushion on a partly-used multi-GPU box (reproduced ~3.3 GB over).

- Draft K and V cache types are parsed and accounted independently. A one-sided
  override (e.g. --cache-type-k-draft q4_0, V left f16) no longer applies the
  small quant to both axes and under-reserves the f16 axis. The embedded-head
  formula sizes per axis; the separate-drafter path uses the heavier type so it
  never under-reserves.

- The compute-graph buffer honors a user --ubatch / --ubatch-size / -ub override
  (parsed and threaded into every _estimate_compute_buffer_bytes call and the
  tensor planner); it previously always assumed the 512 default, under-reserving
  up to ~8x at --ubatch 4096.

- GPU ranking uses the usable budget (free - (1-frac)*total) instead of raw free
  in _select_gpus and both auto-context subset loops, so a more-used large card
  no longer outranks a less-used small card that has more usable room.

Adds regression tests for each (tensor total cap, ubatch reserve scaling, split
K/V no-under-reserve, --ubatch parser, usable-ranking GPU selection). Full
targeted backend suite green (321 passed).

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

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

* studio: gate tensor-parallel admission on usable VRAM budget

The tensor-parallel GPU admission filters still used raw free VRAM after
the total-based budget landed, an asymmetric fix: a partly-used large card
can clear the per-device compute-buffer reserve on raw free while its usable
budget (free - (1-frac)*total) does not, so the planner admitted it and the
even split could emit a near-zero weight slice for a GPU that should have
been excluded.

- _plan_tensor_parallel: admit GPUs by usable budget, not raw free (move the
  _usable helper above the filter).
- load_model: admit the tensor set by _gpu_usable, and downgrade to layer
  split when the pooled usable budget cannot hold weights plus per-device
  compute buffers (the planner can only floor the context, not stop an
  overcommitted launch).

Adds regression tests: planner drops a GPU whose usable budget is below the
reserve, and a source-level check that load_model admits on the usable
budget and carries the pooled-weight downgrade.

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

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

* studio: size the MTP reserve for the user's overriding drafter

A user --model-draft passed in extra_args is appended last and wins at the
llama-server launch, but the VRAM budget preferred Studio's auto-detected
drafter (mtp_draft_path or extras), so a larger custom drafter was
under-reserved. Flip the precedence to extras-first, matching the draft-depth
(n_max) resolution two lines above. Adds a source-level regression test.

* studio: account for MTP reserve in tensor gate, restore 2-col GPU probe

Two issues found by re-review of the prior fix:

- The tensor-parallel capacity gate only checked the model weights against the
  pooled budget, not the MTP reserve. A separate-drafter MTP load whose weights
  fit but weights + drafter do not could still launch overcommitted in tensor
  mode. Add the non-shrinkable MTP reserve (drafter weights + floor draft KV, or
  the flat 2 GiB fallback when dims are unavailable) to the gate.

- The nvidia-smi probe was switched to a three-column query (index,free,total)
  for the total-based budget but required exactly three columns, so a driver or
  mock returning the legacy two-column "index,free" was dropped and the probe
  fell through to the real GPUs. Accept two columns (total 0) and treat an
  unknown total as the legacy free*fraction in _select_gpus.

Tests: tensor gate asserts the MTP term is included; _get_gpu_memory parses both
two- and three-column output; the existing two-column GPU-detection mocks pass
again.

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

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

* studio: keep the VRAM cushion in tensor planning when GPU totals are unknown

_plan_tensor_parallel fell back to raw free VRAM when a GPU's total was
unavailable (a two-column nvidia-smi probe reporting total 0), while
_select_gpus and the load_model ranking both fall back to free*fraction. That
let tensor planning spend the 5% cushion the rest of the fit preserves and
over-advertise context in exactly that path. Align the fallback to
free*_CTX_FIT_VRAM_FRACTION. Updates the no-totals planner test expectations
(now free*frac) and adds a regression test that total 0 keeps the cushion.

* studio: honor LLAMA_ARG_* env overrides and HF draft flags in the VRAM budget

The budget parsed llama-server flags only from the request's extra_args, but the
child process inherits Studio's full environment (child_env_without_native_path_secret
copies os.environ), and llama-server honors LLAMA_ARG_* env vars for the same
options. So a service-level override the child acts on was invisible to the fit,
which could then advertise a context/GPU set that OOMs at load.

- _extra_args_n_ubatch: fall back to LLAMA_ARG_UBATCH (drives the compute buffer;
  an unseen 4096 vs the 512 default under-reserves ~8x).
- _extra_args_mtp_draft_path: also recognize the HF draft-repo flags
  (--spec-draft-hf/-hfd/-hfrd/--hf-repo-draft) and fall back to
  LLAMA_ARG_SPEC_DRAFT_MODEL / LLAMA_ARG_SPEC_DRAFT_HF_REPO. An HF repo isn't a
  local file so it can't be sized, but recognizing it routes to the flat reserve
  instead of mis-sizing Studio's auto/embedded drafter.
- _extra_args_draft_cache_types: fall back to
  LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V per axis.

CLI extra_args win over env (they are appended last at launch). Each parser takes
an injectable env for deterministic tests. Adds env-fallback and HF-flag tests.

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

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

* studio: review polish - drop non-flag --ubatch, harden GPU probe, document buffer

Non-blocking items from a second review pass; no behavior change in the common path:

- _extra_args_n_ubatch: drop --ubatch; the binary only accepts --ubatch-size/-ub,
  so parsing --ubatch implied support it does not have (it would over-reserve for a
  launch that fails on the unknown flag).
- _get_gpu_memory: skip a malformed nvidia-smi line instead of letting one bad line
  raise and drop the whole NVIDIA probe to the torch fallback.
- _estimate_compute_buffer_bytes: document that the per-slot output-buffer model
  assumes a small n_outputs_max (chat decode); it would under-count for
  embeddings / --logits-all / reranking, which Studio does not run on this path.

* studio: honor LLAMA_ARG_SPEC_TYPE when deciding the MTP reserve

_extra_args_requests_mtp only checked extra_args, but the child inherits Studio's
env and llama-server honors LLAMA_ARG_SPEC_TYPE. So a service-level
LLAMA_ARG_SPEC_TYPE=draft-mtp would run MTP while the fit skipped the draft
reserve and could advertise a context/GPU set that OOMs at load. Recognize the
env value (CLI still wins). Completes the env-override coverage alongside ubatch,
draft model, and draft cache types. Adds an env regression test.

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

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

* studio: reserve VRAM for non-MTP model-based draft modes too

The draft reserve only engaged for MTP. A user passing a non-MTP model-based
draft mode (--spec-type draft-simple / draft-eagle3) with a --model-draft loads
a separate draft model whose weights + KV consume GPU memory, but the fit
reserved nothing and could OOM at load. Engage the existing drafter reserve for
those modes when extras (or LLAMA_ARG_SPEC_TYPE) name a drafter; ngram-* load no
model and are unaffected. Purely additive (reserves where there was none).
Adds parser + gate tests.

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

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

* studio: floor quantized embedded MTP draft KV at f16; fix two test issues

Address PR review feedback (three findings):

1. Quantized embedded MTP draft KV was underpriced. The embedded head is a
   single draft layer, so llama.cpp cannot amortize quantized-KV overhead over
   many layers the way the main model does: a quantized draft KV (e.g.
   --spec-draft-type-k q4_0) actually fits LESS context than f16, not more
   (ggml-org/llama.cpp#24102, where a collaborator recommends f16 for the draft
   KV). Pricing q4_0 at 0.5625 of an element (~28% of f16) under-reserved, so a
   quantized override could advertise a context that shrinks or OOMs at load.
   Floor the embedded draft KV bytes-per-element at f16 (quantized types priced
   as f16, f32 still its full 4 bytes). The separate multi-layer drafter, where
   quantization does amortize, keeps the user's real type.

2. test_load_model_reserves_for_non_mtp_draft_modes asserted an exact one-line
   source substring that pre-commit black wrapped across lines, breaking CI.
   Strip whitespace before matching so the check survives any line-wrapping.

3. test_compute_buffer.py installed a partial httpx stub via setdefault that, if
   collected before test_kv_cache_estimation.py, leaked into sys.modules without
   HTTPError/Response and could break the transformers introspection tier by
   collection order. Adopt the sister file's pattern: only stub when real httpx
   is absent, and include the full symbol set.

Updates the affected draft-KV tests to assert the f16 floor.

* studio: guard httpx stub in test_mtp_vram_budget too

test_mtp_vram_budget.py installed a partial httpx stub via setdefault that, like
test_compute_buffer.py before it, lacked HTTPError/Response and could leak into
sys.modules ahead of tests that need huggingface_hub/transformers, breaking the
introspection tier by collection order. Apply the same guard used by
test_kv_cache_estimation.py: only stub when real httpx is absent, with the full
symbol set.

* studio: per-device layer-split reserve, effective spec-type, drafter weights, KV restore

Address PR review feedback (four findings in the auto-fit budget):

A. Reserve the per-device layer-split overhead. A layer (pipeline) split allocates
   a fixed per-device overhead (CUDA context + per-device compute scratch) on every
   participating GPU, beyond the slot-scaling compute buffer that is conserved across
   the split. Measured ~0.9 GB/device on the Qwen3.6-27B GGUF (b9625), independent of
   --parallel: layer-split TOTAL VRAM grew +894 MiB (parallel=8) / +946 MiB
   (parallel=1) per extra GPU, ~linear to +2.6 GB at 4 GPUs. The fit folded a single
   compute buffer for all subset sizes, so a k-GPU layer split was short by
   ~(k-1)*0.9 GB and could pin a context that fits the pool on paper but OOMs a device.
   Reserve (k-1) * _PIPELINE_PER_DEVICE_OVERHEAD_MIB per subset in the layer-split fit;
   k=1 adds nothing, so single-GPU sizing (and the validated benchmark rows) is unchanged.

B. Track the effective --spec-type. _extra_args_requests_mtp returned true on the
   first MTP-ish --spec-type and consulted LLAMA_ARG_SPEC_TYPE even when a CLI
   --spec-type was present, contrary to llama.cpp (last CLI value wins; a CLI flag
   overrides the env). So `--spec-type draft-mtp --spec-type ngram-mod` or a non-MTP
   CLI value with a stale MTP env over-reserved a drafter the launch won't load
   (shrinking context / selecting extra GPUs). Route both detectors through a new
   _effective_spec_type helper.

C. Keep known drafter weights in the fallback reserve. When a separate drafter's KV
   metadata can't be sized, _estimate_mtp_overhead_bytes returned None and discarded
   the drafter's known weight bytes, falling back to the flat 5% reserve; a drafter
   larger than that cushion could launch over budget and OOM. Reserve the known
   weights even when KV sizing fails (None only when nothing is known).

D. Restore quantized KV on tensor->layer-split downgrade. The tensor attempt drops a
   quantized KV cache (tensor mode aborts on it). When the GPU-count or capacity gate
   then downgrades to layer split -- which supports quantized KV -- the dropped type
   was lost and the launch used f16, using more VRAM and shrinking context. Remember
   the dropped type and restore it on downgrade (the launch re-emits it from the var).

Adds regression tests for each.

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

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

* studio: per-device overhead in GPU pin, skip CPU draft, gate env spec-type

Address PR review feedback (three follow-up findings):

F1. Reserve the per-device layer-split overhead in the pin path too. The earlier
    per-device reserve was added to the auto-context fit loops but not to
    _select_gpus, which the explicit-ctx and file-size-only paths use to PIN GPUs
    with -ngl -1 (no --fit fallback). A 2+ GPU pin within ~1 GiB/extra-GPU of the
    budget could OOM a device at load. Add a per_device_overhead_bytes arg to
    _select_gpus so a k-GPU pin must hold model + (k-1)*overhead; pass the pipeline
    overhead at both pin call sites. Single-GPU pins are unchanged.

F2. Don't charge a CPU-offloaded drafter against the GPU budget. A user passing
    --spec-draft-ngl 0 or --spec-draft-device none/cpu keeps the separate draft
    model's weights + KV on CPU, but the budget still charged the full drafter GGUF
    size, auto-reducing context or downgrading GPU selection. Detect the CPU-offload
    flags and drop the separate drafter (and its flat fallback) from the budget; an
    embedded head follows the main -ngl and is unaffected.

F3. Consult LLAMA_ARG_SPEC_TYPE only when it can reach the child. llama-server's CLI
    args override env, and _build_speculative_flags emits a --spec-type/--spec-default
    for every UI mode except "off". So a stale MTP env on a non-MTP model (auto mode)
    made the fit reserve MTP that the emitted --spec-default disables, shrinking
    context / picking extra GPUs. Gate the env consult on "no user --spec-type and UI
    mode off"; the MTP-model auto path still engages via Studio's own detection.

Adds regression tests for each.

* studio: drafter budget precedence and --spec-default in effective spec-type

Two spec-precedence fixes surfaced by an independent multi-reviewer pass:

R3. Size the drafter the launch actually loads. _mtp_draft_for_budget consulted
    LLAMA_ARG_SPEC_DRAFT_MODEL (via _extra_args_mtp_draft_path's env fallback)
    before Studio's resolved mtp_draft_path, but _build_speculative_flags emits
    --model-draft mtp_draft_path, which overrides the env at launch. With a stale
    (smaller) env drafter, the budget under-reserved and could OOM. Order the
    budget by what actually launches: CLI extras --model-draft (appended last,
    wins), then Studio's emitted mtp_draft_path (when MTP engages and the user
    doesn't own --spec-type), then the env drafter.

R4. Treat --spec-default as a CLI spec override in _effective_spec_type. It only
    recognized --spec-type, so extras=["--spec-default"] with LLAMA_ARG_SPEC_TYPE=
    draft-mtp fell through to the env and over-reserved MTP, even though the CLI
    --spec-default overrides the env to a non-MTP default. Recognize it as a CLI
    spec flag (resolves to "default", non-MTP) that suppresses the env fallback.

Adds regression tests for each.

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

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

* studio: refine MTP draft reserve (parallel slots, last-wins, KV cushion, ranking)

Address PR review feedback (five follow-up findings, all edges of this session's
earlier MTP/auto-fit changes):

G1. Price the separate drafter's KV per --parallel slot. _mtp_draft_kv_bytes called
    the drafter's _estimate_kv_cache_bytes with the default n_parallel=1, but the
    drafter is served under the main model's slot count; a sliding-window drafter
    (Gemma) grows KV per slot and was under-reserved. Thread n_parallel through the
    draft KV / overhead estimate and the fit closure.

G2. Honor last-wins for the draft-offload flags. _extra_args_draft_offloaded_to_cpu
    returned True on the first CPU value, so --spec-draft-ngl 0 --spec-draft-ngl -1
    (final = GPU) wrongly dropped the drafter reserve while the server kept it on
    GPU -> OOM. Decide on the final value of each flag only.

G3. Keep the flat cushion when only the drafter weights could be sized. The weights
    fallback installs mtp_overhead_fn, which made callers drop the flat MTP reserve,
    leaving the still-unsized draft KV with no cushion. Keep the flat fraction on in
    that weights-only case, on top of the byte-accurate weights.

G4. Rank auto/cap GPU subsets by the active budget fraction. The ranking used a
    hard-coded 0.95 while the fit tests _pin_fraction (lowered by the flat MTP
    reserve); on mixed-total GPUs that could order subsets differently and pick a
    worse plan. Rank with the same fraction the fit uses.

G5. Keep the embedded-head flat reserve under a draft CPU-offload flag. F2's
    not-_draft_on_cpu guard also dropped the reserve for an embedded MTP head, which
    is part of the main model and stays on GPU regardless of --spec-draft-ngl. Only
    suppress the flat reserve for a CPU-offloaded separate drafter (no embedded head).

Adds regression tests for each.

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

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

* studio: keep GPU on non-integer total, keep tensor flat reserve for weights-only

Two review findings:

- _get_gpu_memory dropped a whole GPU when nvidia-smi reported a non-integer
  memory.total ("N/A" on some drivers / MIG / vGPU): index, free and total were
  parsed in one try/except that skipped the line on any ValueError, so the GPU
  vanished from the probe and the load could silently spill to CPU. Parse index
  and free (required) first, then total separately, defaulting to 0 (the fit then
  uses the free*frac path for that GPU). Adds N/A and bad-free test cases.

- Tensor planning skipped the flat MTP reserve for a weights-only drafter (file
  size known, KV unsizable): the capacity gate used the byte floor whenever
  mtp_overhead_fn was set, so it reserved only the drafter weights and no draft
  KV. Tensor mode has no --fit valve, so that could overcommit and OOM. Keep the
  flat reserve (never below the byte floor) in the weights-only case too, mirroring
  the layer-split _mtp_kv_unsized handling. Adds a regression test.

(A third suggestion -- fold --batch-size into the compute-buffer reserve -- was
checked on hardware and declined: -b 8192 -ub 512 used identical VRAM to the
default at -c 64000, so the logical batch does not size the graph buffer; the
estimate correctly uses the physical micro-batch.)

* studio: budget the main KV from LLAMA_ARG_CACHE_TYPE env when Studio emits none

The child inherits LLAMA_ARG_CACHE_TYPE_K / LLAMA_ARG_CACHE_TYPE_V, but Studio
emits --cache-type-k/-v only when the param or extras set the type. When neither
does, a heavier env type (f32) reaches the child while the auto-fit budget
assumed the f16 default, under-reserving the main KV and risking OOM at the
advertised context. This is the one main-KV axis that lacked the env-aware
handling the other axes already have (spec-type, draft model, draft cache type,
ubatch).

load_model now adopts the heavier of the two env types when it exceeds f16 (only
f32 does), and the launch re-emits it so child and budget stay byte-consistent.
Quantized env types are <= f16 and remain safely over-reserved by the default,
so they are left untouched (no change). A single value is used because the
budget's KV estimate has one cache_type_kv knob, matching parse_cache_override's
existing key/value collapse.

Adds _env_main_cache_type_for_budget plus regression tests covering f32 adoption,
the K/V heavier-of collapse, quantized/unknown no-ops, and the load_model source
precedence.

* studio: budget tensor parallel when LLAMA_ARG_SPLIT_MODE env selects it

Studio emits --split-mode tensor only on its tensor branch; the default
layer-split path emits nothing and resolve_tensor_parallel consults only extras.
The child inherits LLAMA_ARG_SPLIT_MODE, so a tensor env on a layer-split plan
silently runs the child tensor-parallel (heavier per-device compute buffer)
while the budget reserved only the layer-split per-device overhead, under-
reserving on multi-GPU.

load_model now flips the plan to tensor when extras do not set a split mode and
the env selects tensor, so Studio plans, reserves, and emits tensor consistently.
The flip is one-directional (guarded on not tensor_parallel and no extras
split-mode) so an existing tensor plan is never downgraded and extras keep
precedence. Other env modes (layer/row/none) are not a runtime-heavier surprise
and are left untouched.

Adds _env_split_mode_is_tensor plus unit and load_model source-level tests.

* studio: reconcile inherited llama.cpp env with the budgeted launch decision

Addresses a review pass over the VRAM auto-fit work. The budget now sizes the
right amount, but the child process inherits LLAMA_ARG_* env (see
child_env_without_native_path_secret), and a few axes could still run the child
in a mode Studio neither chose nor budgeted.

Mixed known/unknown GPU totals over-advertised the pooled layer-split budget.
_pool_budget_mib pooled free and total separately, so an unknown-total GPU
(MIG/vGPU/N/A) contributed its full free with no cushion when mixed with
known-total GPUs (~(1-frac)*free over-advertise, about 500 MiB in a two-GPU
case). It now sums each GPU's own usable budget, and the layer-split fit calls
take that as an absolute budget (budget_frac=1.0, total_mib=None) so the fit and
the footprint check agree. All-known-total pools are unchanged.

LLAMA_ARG_SPLIT_MODE=tensor survived a tensor-to-layer downgrade. The downgrade
only stripped CLI extras, so the inherited env still ran the child tensor while
Studio budgeted layer split. When the final decision is layer split, a non-layer
inherited split mode (and any paired LLAMA_ARG_TENSOR_SPLIT) is now cleared from
the child env.

Inherited quantized LLAMA_ARG_CACHE_TYPE_K/_V crashed tensor mode. Tensor mode
aborts on a quantized KV cache; Studio drops a quantized cache_type_kv for the
tensor attempt but the inherited env reached the child anyway. When the final
decision is tensor split, a quantized cache-type env is now cleared so the child
uses the tensor-safe default that was budgeted.

Env-derived cache budget no longer mutates the emitted launch flags. An env-only
main KV type now informs the budget only; it is not re-emitted, so an asymmetric
K=f32,V=f16 env reaches the child as set instead of being rewritten to symmetric
--cache-type-k/-v f32.

Adds source-level regression tests for all four and confirms the documented
single-GPU/tensor/pipeline numbers are byte-identical before and after.

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

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

* studio: tighten comments in the VRAM auto-fit code

Compress the verbose docstrings and inline comments added by this work to
succinct 2-4 line versions, drop restated/obvious ones, and cut duplicated
rationale across the two tensor-downgrade branches. Keeps the non-obvious intent
(env-inheritance precedence, the #24102 embedded-draft floor, the per-device
overhead and pool-budget rationale) while removing roughly 120 lines of comment
text from llama_cpp.py. Also trims the few longest test-comment blocks; concise
per-test scenario notes are left intact.

No logic change: verified with comment_tools.py check --strip-docstrings (code-
only signature unchanged vs the prior commit) and the full backend suite still
passes (824).

* studio: lock in env-drafter engagement for the separate-draft reserve

A review suggested an env-provided LLAMA_ARG_SPEC_DRAFT_MODEL would skip the
draft reserve and OOM. It does not: the gate's _extra_args_mtp_draft_path(extra_args)
call defaults env=None, which consults os.environ, so an env-only drafter still
sets _user_draft_via_extras and is sized via _env_draft_for_budget. Add a source
guard that the gate keeps the env-inclusive form (not extras-only env={}) and a
behavioral test mirroring the reviewed scenario, so a future cleanup can't
regress it. No production change.

* studio: carry the unsized MTP reserve and env split/offload into tensor planning

Addresses a review pass over the multi-GPU and env-inheritance paths.

Tensor planner dropped the unsized draft-KV cushion. When a separate drafter has
known weights but unreadable KV metadata, _plan_tensor_parallel receives a
non-None weights-only mtp_overhead_fn and applied the flat 2 GiB reserve only for
the no-fn case, so its binary search spent the unsized-KV cushion on context and
over-advertised. Add mtp_flat_reserve_bytes (subtracted from the pooled budget and
the even-split check), and pass it from load_model whenever _mtp_kv_unsized. The
layer path and the tensor pre-gate already kept this cushion.

Stale LLAMA_ARG_TENSOR_SPLIT survived in tensor mode. When the planner picks an
even split it emits no --tensor-split, so an inherited tensor-split env reached the
child and overrode the budgeted split. The layer downgrade branch cleared it; the
tensor branch now does too.

Env-only draft CPU offload was ignored. _extra_args_draft_offloaded_to_cpu checked
extras but not LLAMA_ARG_N_GPU_LAYERS_DRAFT, so an env-offloaded drafter was still
charged GPU budget and under-advertised context. It now consults that env (the
device flag has no env), called with env=os.environ.

Layer-split compute buffer had no fallback when GGUF dims are missing. The estimate
returns 0 then, so the layer path folded no buffer while the tensor path falls back
to the flat reserve. Use the flat reserve for the layer path too (a safe upper
bound, since the tensor buffer >= the layer one).

All four are gated on conditions the documented benchmarks don't hit; the
single-GPU/tensor/pipeline reconfirm numbers are byte-identical, and the full
backend suite passes (830) with regression tests for each fix.

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

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

* studio: share the env-aware tensor decision across load and dedup matchers

A review pass found the inherited-LLAMA_ARG_SPLIT_MODE=tensor flip lived only
in load_model, so the two duplicate-load matchers disagreed with it.

Consolidate the decision into _effective_tensor_parallel (extras + toggle, then
flip on when extras set no split mode and the child inherits a tensor split
env). load_model, the backend matcher (_already_in_target_state) and the route
matcher (_request_matches_loaded_settings) now all call it. Before, an env-driven
tensor server compared against resolve_tensor_parallel (env-blind) in both
matchers, so a follow-up load that should dedup was seen as a mismatch and the
healthy server was needlessly killed and reloaded.

Also finish the tensor cache-type handling: when the tensor attempt drops a
quantized KV it now re-adopts a heavier inherited env cache type (f32) for the
budget, mirroring the initial adoption; and the two layer-split downgrades clear
_cache_type_from_env so the restored quantized type is actually re-emitted rather
than left to a stale inherited env.

All gated on inherited env the documented benchmarks don't set; the single-GPU,
tensor and pipeline reconfirm numbers are byte-identical, and the full backend
suite passes (832) with unit + source regression tests for the shared helper and
the route matcher.

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

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

* studio: complete the env-aware tensor/spec handling across all paths

A second review pass found the env-aware tensor/MTP handling was applied
asymmetrically: some paths inherited LLAMA_ARG_* env, others didn't. Three real
follow-ups, plus a small consolidation so the env semantics live in one place.

1. Tensor fallback ignored the inherited tensor env. load_with_tensor_fallback
   computed its retry gate with the env-blind resolve_tensor_parallel, so an
   env-only tensor load (toggle off, no --split-mode extra) that crashed on a
   tensor-incompatible GGUF re-raised instead of retrying layer split. It now
   uses the env-aware decision; and since the inherited env would otherwise
   re-engage tensor on the retry (CLI args persist, the env does too), the retry
   forces --split-mode layer (CLI wins over env) so it can't re-crash.

2. Duplicate-load matchers looped reloads after a tensor->layer downgrade. Both
   matchers compared the env-expanded tensor decision against the loaded server,
   but load_model may downgrade tensor to layer (capacity/buffer) and scrub the
   child env. The still-set parent env then made every identical request look
   like a mismatch, killing and reloading a healthy layer server. Add
   _tensor_parallel_matches_loaded, which only lets an inherited tensor env raise
   a match against a server that actually launched tensor; a downgraded server
   matches the same request (an identical load would downgrade the same way).

3. MTP binary-capability fallback leaked an inherited LLAMA_ARG_SPEC_TYPE. When
   the binary lacks MTP, _emit_mtp degraded but emitted no spec flag, so an
   inherited LLAMA_ARG_SPEC_TYPE=draft-mtp still reached the child and attempted
   MTP the gate had budgeted off. It now emits --spec-default (CLI wins over env)
   like the sibling no-head / non-MTP fallbacks.

Consolidation: moved _env_split_mode_is_tensor / _effective_tensor_parallel into
llama_server_args.py (with the new _tensor_parallel_matches_loaded) so the
lightweight tensor_fallback module can share them without importing llama_cpp;
llama_cpp re-exports them for back-compat.

All gated on inherited env the documented benchmarks don't set; the single-GPU,
tensor and pipeline reconfirm numbers are byte-identical, and the full backend
suite passes (883) with regression tests for each fix.

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

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

* studio: budget the heavier axis of asymmetric --cache-type-k/-v extras

A review pass found the explicit-extras counterpart of the env cache-type fix.
load_model adopts the heavier inherited LLAMA_ARG_CACHE_TYPE_K/_V env for the
reserve, but the explicit-extras path used resolve_cache_type_kv, which collapses
both axes to one last-wins value. So extras such as
--cache-type-k f32 --cache-type-v f16 (lighter axis last) budgeted f16 for both
axes while the child allocates f32 on K, over-advertising context and
re-opening the OOM path this PR closes.

Add parse_cache_override_per_axis (keeps the K/V last-wins values apart) and
_extra_args_main_cache_type_for_budget (the heavier of the two by bytes/elem),
and budget from it. The user's extras are appended last and win per axis at the
child, so this only raises the reserve; the emitted command and the asymmetric
child cache are unchanged, and the common single-axis / symmetric cases resolve
to the same type as before.

Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical, and
the full backend suite passes (892) with per-axis parser and heavier-axis budget
regression tests.

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

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

* studio: fix tensor-safety masking and strip inherited HF drafter selectors

A review pass found two more env/extras edge cases on the speculative and tensor
cache paths.

Tensor-safety could miss a quantized axis. The previous change budgets the
heavier-by-bytes cache type, but that masks a quantized axis paired with a
heavier one: --cache-type-k f16 --cache-type-v q4_0 resolves to f16, so the
tensor-safety block did not fire and the q4_0 axis survived into tensor mode,
which aborts on quantized KV. Test each explicit --cache-type-k/-v axis (not just
the budget type) so any quantized axis drops the cache for the tensor attempt.

Inherited HF drafter selectors were not stripped. _extra_args_mtp_draft_path
treats --spec-draft-hf / -hfd / -hfrd / --hf-repo-draft as drafter selectors, but
_SPEC_FLAGS only stripped the local --model-draft selectors, so on an inherited-
extras Apply a stale HF drafter survived and last-wins-overrode Studio's
re-derived spec choice. Add the HF aliases to _SPEC_FLAGS. The per-drafter tuning
knobs (--spec-draft-type-*, -ngld, --spec-draft-device) are intentionally left in
place: the VRAM budget reads them via the same parsers the child honors, so they
stay consistent on inherit, and stripping them would silently move a CPU-offloaded
drafter back onto the GPU.

A third flagged item -- that the HF draft env var should be LLAMA_ARG_HFD_REPO --
was a false positive from a stale manpage; the bundled binary's common/arg.cpp
sets LLAMA_ARG_SPEC_DRAFT_HF_REPO for --spec-draft-hf, which the code already
uses, so it is left unchanged.

Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical, and
the full backend suite passes (899) with regression tests for both fixes.

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

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

* studio: preserve asymmetric cache on tensor downgrade and skip CPU-drafter reserve

A review pass found two more tensor-path edges, one a regression from the
per-axis cache change.

Tensor-to-layer downgrade collapsed asymmetric cache extras. The per-axis
tensor-safety check strips an asymmetric --cache-type-k/-v (tensor rejects
quantized KV), but the downgrade restored only the scalar heavier type, so a
layer fallback silently rewrote --cache-type-k q4_0 --cache-type-v f16 to
symmetric f16/f16 even though layer split supports the original. Save the
original extras before the tensor strip and restore them verbatim (minus the
user --split-mode) on both downgrade points; the budget still uses the heavier
scalar, the child gets the real asymmetric cache. Before the per-axis change this
case happened to survive (last-wins was f16, untouched), so this restores that.

Tensor mode reserved GPU VRAM for a CPU-offloaded drafter. The layer path drops
the flat MTP reserve when the only drafter is a separate CPU one with no embedded
head, but the tensor capacity gate and planner still charged it, under-advertising
context. Gate the tensor reserve on the same condition via _mtp_reserves_gpu.

Reconfirm numbers (single-GPU table, tensor, pipeline) are byte-identical (both
fixes are gated on conditions the benchmarks don't hit), and the full backend
suite passes (901) with regression tests for each.

* studio: drop now-unused llama_server_args imports from llama_cpp

The refactor re-pointed load_model and the matchers off resolve_tensor_parallel /
resolve_cache_type_kv and moved the env split-mode helper into llama_server_args,
leaving those three names imported but unused in llama_cpp. The repo's import-hoist
safety-net lint blocks that, so drop them; the env split-mode test now imports
_env_split_mode_is_tensor from its real home (llama_server_args).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-17 03:10:22 -07:00
alkinun
14188d6f45
Add API server monitor in Studio (#5558)
* Add Studio API activity monitor

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

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

* Mark disconnected monitor streams cancelled

* Fix monitor lifecycle, parsing, and clock issues for PR #5558

Backend:
- Finalize the monitor entry as error on the chat-completions validation
  reject paths (addresses the codex P2 comment around line 2337). Adds a
  small _reject helper next to the api_monitor.start so each early-raise
  in the GGUF tool-passthrough, non-empty-messages, GGUF vision/audio,
  PIL image decode, and text-only model branches no longer leaves the
  entry stuck running until eviction.
- _monitor_openai_chunk is now defensive about malformed shapes
  (non-list choices, non-dict choice/delta) so a misbehaving provider
  chunk does not raise into the streaming generator and abort the
  user's response.
- _monitor_openai_sse_line accepts both data:value and data: value per
  SSE spec; previously the single-space-only prefix silently dropped
  tokens from compact emitters.
- openai_completions stream now keeps a residual buffer between
  aiter_bytes chunks so a data: line whose newline lands in the next
  TCP frame is reconstructed before parsing, instead of being dropped
  by the per-chunk splitlines call.

api_monitor:
- duration_ms is derived from time.monotonic anchors and clamped at 0,
  so NTP / manual clock steps no longer produce negative durations.
- finish() and fail() are idempotent: a second call (for example [DONE]
  arriving after the generator's finally block already ran) no longer
  moves finished_at or finished_monotonic.
- set_usage only derives total_tokens when no authoritative total has
  been recorded, so a later partial-usage chunk does not clobber a
  provider-reported total from an earlier chunk.
- _trim guards limit < 3 so the slice cannot underflow.

Tests:
- Adds regression coverage for the new idempotency, total preservation,
  monotonic clock, and _trim guard behaviour.

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

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

* Mark monitored stream failures correctly

* Fix completions monitor failures

* Fix responses stream monitor cleanup

* Fix audio input API monitor lifecycle

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

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

* Fix API monitor review follow-ups

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

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

* Fix completions monitor usage accounting

* Fix API monitor cancellation and usage gaps

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

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

* Move API monitor into settings

* Fix API monitor review followups

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

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

* Lazy load API monitor details

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

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

* Fix API monitor review issues

* Finalize passthrough monitor on clean EOF

* Finalize monitor on chat validation rejects

* Fix API monitor review follow-ups

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

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

* Fix API monitor follow-up review issues

* Monitor embeddings and tool-call replies

* Cover tool-call monitor replies and cancellations

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

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

* Scope API monitor entries by subject

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

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

* Finalize monitor entries on cancellation

* fix: address PR 5558 CI failures

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

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

* fix: update provider proxy test stub

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-16 20:26:13 +01:00
Daniel Han
58c2ec1ebd
Studio: Xet-primary model downloads with automatic HTTP fallback on stall (#6372)
* Studio: add shared Xet-primary download helper with HTTP stall fallback

Xet is the fast default transport in huggingface_hub, but a stalled Xet
transfer hangs with no progress and no exception, and a blocked native thread
cannot be killed. The safetensors inference path already recovers (subprocess
watchdog + respawn with HF_HUB_DISABLE_XET=1); the GGUF and training paths do
not. Add a reusable helper that the in-process paths can adopt.

utils/hf_xet_fallback.py:
- DownloadStallError (moved here from core/inference/orchestrator.py, which now
  imports it; behavior unchanged, still a RuntimeError subclass).
- get_hf_download_state / start_watchdog: a no-progress watchdog built on the
  sparse-aware hub.utils.hf_cache_state helpers; fires only while a .incomplete
  is present and the on-disk byte total is unchanged for stall_timeout.
- hf_hub_download_with_xet_fallback: cached files short-circuit; otherwise the
  download runs in a spawn child (own process group) supervised by the watchdog.
  On a stall it kills the child, makes the partial safe for HTTP via
  prepare_cache_for_transport, and respawns once with HF_HUB_DISABLE_XET=1. Cancel
  and deterministic errors (auth/missing/disk) propagate without a fallback.

Tests cover the watchdog state machine, the transport decision logic, and a
regression lock that HF_HUB_DISABLE_XET is honored in a fresh interpreter.

* Studio: route GGUF Chat-Mode downloads through the Xet->HTTP fallback

The GGUF load path (_download_gguf main+shards, _download_companion_gguf for
mmproj/MTP) called a bare blocking hf_hub_download with no recovery, so a Xet
stall hung the Chat-Mode load with no fallback. Route those three calls through
hf_hub_download_with_xet_fallback: Xet stays primary, HTTP is used only if Xet
stalls, per-file so finished shards stay cached. The existing _cancel_event is
threaded through, the Cancelled sentinel is preserved, and companions stay
best-effort (a terminal stall is swallowed to None). Cached files short-circuit
in the helper with no subprocess, so the fast path is unchanged.

The two offline mmproj tests are repointed from huggingface_hub.hf_hub_download
to the new call boundary (the helper) since the download now goes through it.

* Studio: recover a stalled training model-load via Xet->HTTP respawn

Training runs in a spawn subprocess and FastModel.from_pretrained downloads
internally, so the download cannot be wrapped per-file like GGUF. Instead the
worker now watches the HF cache during the model-load phase (emitting
model_load_started / model_load_completed and a stall event), and the parent
recovers a stall by terminating the worker and respawning it once with
HF_HUB_DISABLE_XET=1.

worker.py: set HF_HUB_DISABLE_XET=1 before any HF import when the parent passes
disable_xet (respawn), and wrap trainer.load_model with start_watchdog.

training.py: plumb disable_xet through the config; track the model-load window;
on a first-load stall arm a one-shot respawn (handled on the exiting pump thread,
so no pump self-join) that preserves the DB run row (history is not duplicated)
and re-runs the load over HTTP. A second stall, or a stall outside model-load,
surfaces as a normal error. W&B init happens after model-load, so a pre-load
respawn cannot duplicate it; the dataset is re-formatted in the new worker.

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

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

* Studio: add gated test-only fault-injection hook for the Xet stall path

UNSLOTH_HF_XET_FORCE_STALL=1 makes the Xet download attempt write a partial
blob and hang, so the no-progress watchdog and the HTTP fallback can be
exercised end to end against a real repo (never set in production). Used to
verify recovery on real models: a forced Xet stall on a 5.37GB Qwen3.5-35B-A3B
shard triggered the watchdog and the HTTP retry downloaded the correct file
(sha256 verified).

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

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

* Studio: tighten Xet-fallback comments and consolidate its tests

Trim docstrings and inline comments across the Xet->HTTP fallback code to
the non-obvious why (spawn-not-thread, killpg-not-getpgid, the sparse-partial
HTTP-resume hazard); drop comments that merely restate the code. Verified
comment-only with an AST signature check.

Merge the three helper-level test files (watchdog, transport policy, and the
HF_HUB_DISABLE_XET regression lock) into tests/test_hf_xet_fallback.py, and
prefer the real structlog over a bare stub so test collection order cannot
leak an incomplete module to others that log at import.

Full backend suite: 3455 passed, 14 pre-existing flash-attn failures only.

* [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-16 06:17:54 -07:00
Wasim Yousef Said
048f34e8f2
Fix GGUF variant file selection (#6342)
* Fix GGUF variant resolution

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

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

* Address GGUF variant review feedback

* Harden GGUF endian filtering

* Address GGUF endian review comments

* Mirror GGUF endian filter in local resolver

* Fix GGUF route import test stub

* Apply GGUF endian filtering across load paths

* [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-16 12:42:58 +02:00
Nilay
0ac1fb5d9e
CLI: fix --local-dataset being parsed as a string instead of a list (#6357)
* fix CLI dataset path resolution

* enhance list type check

---------

Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com>
2026-06-16 02:32:26 -07:00
Daniel Han
5a38447b25
Studio: omit --threads when unset so llama.cpp picks physical cores (#5894)
* Studio: omit --threads when unset so llama.cpp picks physical cores

Studio passed --threads -1 when no thread count was set. The intent was physical cores, but an explicit --threads -1 makes llama.cpp's arg parser resolve it to hardware_concurrency() (every hyperthread), which contends on the memory bus and slows CPU and hybrid decode (a user saw about 60-75 fall to about 20-30 tok/s under CPU offload). Leaving --threads unset keeps n_threads at -1, which llama.cpp resolves to physical cores via common_cpu_get_num_math(). Omit the flag when unset; still pin it for an explicit override and the Windows full-offload OpenMP cap.

* Studio: drop inherited LLAMA_ARG_THREADS when omitting --threads

Omitting --threads relies on llama.cpp resolving physical cores via common_cpu_get_num_math. But the child inherits os.environ and llama.cpp also reads --threads from LLAMA_ARG_THREADS, which routes through the arg handler and maps <=0 to hardware_concurrency. So an ambient LLAMA_ARG_THREADS would silently override the physical-core default. Scrub it from the child env only when we omit the flag.
2026-06-16 01:32:19 -07:00
Daniel Han
cef7dcf160
Studio: improve logging for dynamic transformers version switching (#6108)
* Studio: log transformers version-switching decisions and stop swallowing MLX activation failures

Two logging gaps in dynamic transformers version switching (issue #6103):

1. get_transformers_tier returned a tier with no trace of why. Add an
   info log at each decision point naming the model and the trigger
   (which substring matched, or which config check fired), so a model
   landing on the wrong tier is diagnosable.

2. The MLX fast-path in run_training_process activated the transformers
   version inside a bare 'except Exception: pass', silently swallowing
   failures while the non-MLX path reports them. A missing or broken
   version venv (e.g. Gemma-4 needing 5.5.0) left no trace and only a
   confusing downstream crash. Extract a small _activate_transformers_version_or_warn
   helper that logs a warning on failure while keeping the non-fatal
   fall-through, and call it from the MLX path.

Adds tier-selection logging tests and helper warn/silent tests.

* Studio: clarify path-prepend log, warn on venv version mismatch, log per-package install progress

Completes the remaining logging items of #6103 in studio/backend/utils/transformers_version.py:

- activate_transformers_for_subprocess: the early "Activated transformers X.X.X" line was misleading because at that point only the venv directory has been prepended to sys.path, not imported. It now says it prepended the venv to sys.path and notes the loaded version is confirmed later by "Subprocess loaded transformers ...".
- _venv_dir_is_valid: a detected version mismatch is logged at warning instead of info, since it immediately triggers a full venv wipe and reinstall that should be visible in the logs.
- _ensure_venv_dir: log each package as it starts installing with an N/M progress counter, so a slow runtime install is not mistaken for a hang (pip/uv output is piped and only surfaced on error).

Adds tests covering all three behaviours; pre-existing unused imports are left untouched.

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

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

* Studio: make tier log-capture tests independent of import order

The new issue #6103 caplog assertions in test_transformers_version.py
relied on the module-level sys.modules.setdefault("loggers", stub)
winning the import race. In a full backend pytest run another module
(for example test_log_filter_no_truncation, collected earlier) imports
the real loggers first, so the setdefault is a no-op and
transformers_version.logger becomes a structlog/stdout logger that
caplog cannot capture -- the tier, activation, venv-mismatch and
install-progress log assertions then fail even though the line was
emitted.

Bind a real stdlib logger to transformers_version.logger for the
duration of each test via an autouse fixture, so the module logs through
logging and caplog captures them regardless of collection order.

* Studio: log local checkpoint tier decisions and warn on MLX inference activation

- get_transformers_tier: the local config.json fast path returned a tier
  without logging it, so local checkpoints stayed opaque while HF ids were
  traceable. Log each decision there too, with a caplog regression test.
- inference worker: the MLX path swallowed _activate_transformers_version
  failures with a bare except, the same gap issue #6103 fixed for training.
  Warn instead, keeping the non-fatal fall-through.

---------

Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-15 23:31:43 -07:00
Daniel Han
e73a89ff82
Studio: warn when a GPU model silently loaded on CPU (#6339)
* Studio: warn when a GPU model silently loaded on CPU

llama-server can serve HTTP 200 while running a model entirely on CPU when its GPU backend fails to init, so Studio could run a GGUF on CPU without saying so (#5807 / #5106 / #5830). The silent-CPU warning already exists but stopped firing on current llama.cpp because _classify_gpu_offload keyed only on the dropped 'model buffer size' lines. Add a shared classify_gpu_offload_lines (offloaded N/M counts, GPU model-buffer markers excluding _Host, device_info disconfirm-only) and delegate to it so the warning fires again. Log-only: no install or load behavior changes.

Pure classification of already-captured startup log lines, run once after load; no new subprocess, no slowdown.

* Studio: key the CPU-offload warning on the main model, not a draft

With MTP/speculative decoding llama-server logs 'offloaded N/M layers to GPU' twice: once for the main model and once for the small draft model. The old scan returned True on any non-zero count, so a drafter that fits on GPU while the main GGUF runs on CPU suppressed the warning (the Qwen3.6-27B-MTP case). Decide on the line with the most layers (the main model) instead, so a drafter cannot mask a main model on CPU.
2026-06-15 23:06:44 -07:00
Hua
9413e72802
Fix the libaray path for probe_server_capabilities() (#5797)
* Fix the libaray path for probe_server_capabilities()

even when running something as simple as `./llama-server --help`,
the binary still requires correct LD_LIBRARY_PATH to work - or it
returns merely an "error while loading shared libraries":
"libllama-server-impl.so: cannot open shared object file: No such file or directory"

For a local installation with no LD_LIBRARY_PATH specifically set,
the probe_server_capabilities() run of `./llana-server --help` should
share the same libaray resolution logic as start_llama_server().

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

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

* Adjust and readd comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-15 14:55:41 +01:00
Daniel Han
f297593a7c
Studio: keep training from failing when a namespace-package shadows unsloth (#6269)
* Studio: guard the training import against a namespace-package shadow of unsloth

A directory named unsloth without an __init__.py on sys.path (a stray checkout,
a partial clone, or a polluted PYTHONPATH) makes the path finder return a
namespace package, so the worker's 'from unsloth import FastLanguageModel' fails
with 'cannot import name ... (unknown location)'. A normal site-packages install
always wins this race, so only source/editable installs are exposed. Before the
import, drop the offending sys.path entries, bind the real packages, then
restore sys.path so other modules on those entries keep importing. It is a
no-op when unsloth already resolves to a real package.

* Studio: import unsloth before unsloth_zoo in the namespace-shadow guard

_ensure_real_packages imported the requested names in argument order, so a
unsloth_zoo namespace shadow made it import unsloth_zoo directly before
unsloth. That skips unsloth.__init__ -> _gpu_init, which runs its ROCm and
Windows bitsandbytes fixes before its own import unsloth_zoo, so the recovery
path could import zoo with those guards skipped and fail on the ROCm/Windows
cases those fixes handle.

Import parent-first via reversed(names) so unsloth is imported first and pulls
in the real unsloth_zoo after _gpu_init has run; the later cached import is a
no-op. This also covers the case where only unsloth_zoo is shadowed: the bad
sys.path entry is still dropped and unsloth owns the zoo import. Detection,
sys.path pruning, shadow-cache clearing, and restoration are unchanged.

Add tests/test_namespace_shadow_guard_pr6269.py: CPU-only subprocess scenarios
(only zoo shadowed, both shadowed, only unsloth shadowed, healthy no-op, real
package absent, multiple shadow entries) that assert unsloth imports before
unsloth_zoo and that sys.path is restored.

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

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

* Studio: restore sys.path even if invalidate_caches fails in the shadow guard

Move importlib.invalidate_caches() inside the try/finally so a failure there
still restores sys.path, and tighten the import-order comment. Add a test that
forces invalidate_caches to raise and asserts sys.path is restored.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-15 06:06:48 -07:00
Wasim Yousef Said
9d7740a82f
Rename chat artifacts copy to canvas (#6298)
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-15 13:38:39 +02:00
Daniel Han
6c919bba82
Studio: arm the VRAM-settle wait after the startup orphan reaper (#6315)
On a restart the constructor reaps the previous run's orphaned llama-server, but the driver does not reclaim that VRAM synchronously. _kill_orphaned_servers now returns the number of processes it killed, and __init__ arms _last_kill_monotonic when that count is positive, so the first load_model waits for VRAM to settle before ranking GPUs by free memory instead of pinning the model onto the smaller card. The compute-graph / auto-fit reserve is handled separately in #6312.
2026-06-15 04:22:12 -07:00
Daniel Han
ca0528d1f8
Studio: Bypass Permissions (skip confirmation, disable tool sandbox) (#5895)
* Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls

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

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

* Fix race in tool-call confirmation gate

* Studio: gate built-in tool calls and harden the confirmation handshake

The Allow / Always allow / Deny controls only lived in the fallback tool
card, but the built-in tools (web search, python, terminal, code
execution, image generation) render with their own components and so
never showed the buttons. Those calls paused after tool_start with no way
to approve them, hanging until the 1 hour timeout. Only MCP tools, which
use the fallback renderer, actually worked.

Render the controls for every tool card by wrapping each registered tool
component (and the fallback) in thread.tsx with a shared
ToolConfirmationControls, so the gate applies uniformly.

Also make the handshake robust:
- The gate keys on a per-call approval_id minted by the backend and
  echoed in tool_start, instead of session_id alone, so a stale or
  concurrent confirmation can no longer resolve the wrong call.
- The approval slot is registered before tool_start is yielded, closing
  the race where a fast click or an auto "Always allow" could reach the
  backend before the waiter existed.
- The frontend resolves with the same session id the request was sent
  with (plus the approval_id), fixing the new-thread mismatch where the
  confirmation targeted a different session than the blocked stream.
- The confirm endpoint returns {resolved}; the UI keeps the buttons and
  shows a retry hint until the backend confirms a match, instead of
  hiding them on a failed or mistargeted post.
- The gate runs after the disabled-tool and duplicate-call checks, so a
  call that will not execute is not put up for approval. A denied call is
  still excluded from duplicate detection, so re-issuing and approving it
  works.
- "Always allow" is scoped per session to match the backend gate.

Add backend tests for the approval registry, the SSE no-deadlock
handshake, and the loop integration (allow, deny, disabled, duplicate,
re-issue after deny).

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

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

* Move "Confirm tool calls" to the Tools section

* Studio: add Bypass Permissions (skip confirmation, disable tool sandbox)

Adds an opt-in Bypass Permissions toggle next to Confirm tool calls. When on,
no tool call shows a confirmation prompt and the python/terminal sandbox is
disabled: safety checks, command blocklist, and resource limits are skipped.
Secret env vars are still stripped and HOME stays repointed at the session
workdir. Default off keeps current behavior, and it takes precedence over
Confirm tool calls. Enabling it requires accepting a warning each time.

* Studio: harden Bypass Permissions secret handling and fix Anthropic tool path

Follow-up to the Bypass Permissions feature. Addresses the review findings:

- Anthropic /v1/messages 500: declare bypass_permissions on
  AnthropicMessagesRequest so tool requests that omit the field default to
  False instead of raising AttributeError (extra='allow' does not set absent
  attributes).
- /proc parent-env leak: stripping the child env did not stop a same-uid
  bypassed child from reading /proc/<parent>/environ to recover the
  tool-executing process's unfiltered secrets. Clear PR_SET_DUMPABLE on that
  process before the first bypass exec so its /proc entries become root-owned.
  Hardening is fail-closed: if prctl is denied, bypass execution is refused
  rather than run with the parent environ still readable. Mitigation, not a
  full boundary; documented in the code.
- Broker/capability vars: strip SSH_AUTH_SOCK, SSH_AGENT_PID, GPG_AGENT_INFO,
  GNUPGHOME, KUBECONFIG, DOCKER_HOST so a bypassed tool cannot use the
  operator's live agents.
- Credential-bearing URL values: drop any env var whose value embeds URL
  userinfo (scheme://user:pass@ and token-only scheme://token@) regardless of
  the variable name. Benign proxy/index URLs without credentials are kept, so
  proxy-only and internal-index setups still work in bypass mode.
- Windows temp isolation: repoint TEMP and TMP (not just TMPDIR) at the
  per-session sandbox dir.
- Frontend: stop persisting bypassPermissions; a reload now starts with the
  sandbox/confirmation bypass off and requires re-accepting the warning dialog.

Adds regression tests for each finding in test_bypass_permissions.py.

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

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

* Studio: strip cred-location env vars (HF_HOME etc.) in Bypass Permissions

Repointing HOME did not stop SDKs auto-reading cached creds via vars that
point at the real home/cache/config: HF_HOME (startup always sets it; token
lives under $HF_HOME/token), HF/XDG cache roots, NETRC/BOTO_CONFIG/
PIP_CONFIG_FILE, and Windows HOMEDRIVE/HOMEPATH. Drop those, and repoint
USERPROFILE/APPDATA/LOCALAPPDATA at the per-session workdir. Adds regression
tests.

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

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

* Studio: lock in bypass HF token resolution with an end-to-end test

The drop-based fix relies on the whole HF_HOME/XDG fallback chain being
removed so huggingface_hub resolves under the repointed HOME. Add a test
that sets HF_HOME and XDG_CACHE_HOME at a real cache and asserts the
resolved token path lands under the workdir, not the operator's cache.

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

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

* Studio: strip npm _auth, MYSQL_PWD, and BASH_ENV from bypass env

Three more credential vectors dodged the bypass scrubber: NPM_CONFIG__AUTH
(npm _auth, base64 so no URL userinfo and no AUTH marker), MYSQL_PWD (markers
use PASSWD, not PWD, since PWD is the cwd var), and BASH_ENV (bash -c sources
it for non-interactive shells, so a startup file can re-export stripped
secrets). Add an AUTH marker, the exact MYSQL_PWD name, and drop BASH_ENV plus
PGPASSFILE. Adds regression tests incl. an end-to-end check that a bypass
terminal call does not source BASH_ENV.

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

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

* Studio: extend bypass env scrubber and enforce confirm precedence in loops

From a parallel review pass over the bypass changes:
- Drop more credential-location vars in _build_bypass_env: npm/yarn/git/cargo/
  rclone config pointers (NPM_CONFIG_USERCONFIG, NPM_CONFIG_GLOBALCONFIG,
  YARN_RC_FILENAME, GIT_CONFIG_GLOBAL, GIT_CONFIG_SYSTEM, CARGO_HOME,
  RCLONE_CONFIG) and the GIT_ASKPASS/SSH_ASKPASS auth helpers.
- Enforce confirm_tool_calls AND NOT bypass_permissions inside the safetensors
  and GGUF tool loops, not just at the route, so a direct internal caller
  passing both flags never prompts.
- Soften the toggle hint: environment secrets are stripped, but bypassed code
  can still read files and credentials on the machine (no overclaim that keys
  stay hidden).
Adds regression tests for the new names and the loop-level precedence.

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

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

* Studio: add GGUF loop test for bypass-over-confirm precedence

The safetensors loop precedence is covered behaviorally; the GGUF loop needs a
live llama-server so add an AST guard asserting its _needs_confirm gate
references both confirm_tool_calls and bypass_permissions, matching the other
llama_cpp source-inspection tests.

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

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

* Studio: add red Bypass Permissions badge in the composer

When Bypass Permissions is on, show a persistent red pill in the composer
tool-pill row (like the Search/Code pills), matching Claude Code's always-
visible bypass indicator. Clicking it turns bypass off, mirroring the other
composer toggles. Enabling still goes through the settings toggle + warning
dialog. Adds a data-variant=danger style for the destructive-colored pill.

* Studio: show Bypass Permissions badge in the Thread composer too

The empty-state and active Thread render their own composer (thread.tsx),
not shared-composer, so the badge only appeared in the split layout. Mirror
the red dismissible pill in ComposerAction so it shows in every composer.

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

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

* Studio: keep the Bypass Permissions badge visible when the composer is collapsed

The Thread composer only renders the pill row when expanded, so the active-mode
badge vanished on the default (collapsed) empty state. Render it before the
expand gate (it returns null when bypass is off) so the red indicator always
shows while bypass is on.

* Studio: make the Bypass Permissions confirm button a solid red button

The destructive button variant is a subtle 10% tint that read as bare red text
next to the outlined Cancel. Force the solid destructive fill (the variant's
class loses to the tint through AlertDialogAction's Slot merge, so use the !
override the codebase already uses for this case) and shorten the label to
'I understand' so it fits the small dialog's two-column footer.

* Studio: add Bypass Permissions to the composer + More menu

Adds a 'Bypass Permissions' entry to the composer plus-menu (under More by
default) in both composers, so it can be toggled without opening Run settings.
Enabling routes through the same danger warning dialog; disabling is immediate.
A shared BypassPermissionsMenuItem keeps the two composers in sync.

* Studio: harden bypass env scrubber for IMDS opt-out and connection strings

Two gaps in the Bypass Permissions secret scrubber:

- The broad AWS_ prefix also dropped AWS_EC2_METADATA_DISABLED, a non-secret
  opt-out. Removing it re-opens the IMDS instance-role credential path that the
  operator explicitly disabled, so a bypassed boto/AWS-CLI call could recover
  cloud creds. Keep that flag (and AWS_EC2_METADATA_V1_DISABLED) via a keep-list
  while still stripping the real AWS credential vars.
- Azure App Service connection strings (SQLCONNSTR_/CUSTOMCONNSTR_/...,
  WEBSITE_CONTENTAZUREFILECONNECTIONSTRING) and values like Password=/AccountKey=
  /SharedAccessKey= slipped past the name and URL-only value classifiers. Add
  CONNSTR/CONNECTIONSTRING name markers and a connection-string value matcher.

* Studio: let Bypass Permissions suppress the confirm-tool-calls guards

The confirm-vs-bypass precedence (confirm and not bypass) was applied at the
loop call sites but not at the earlier request guards, so a client sending
confirm_tool_calls + bypass_permissions together was rejected (stream=true
required / unsupported for external or Anthropic tools) before the precedence
took effect. Gate all four confirm guards on not bypass_permissions so both
flags together proceed with the gate suppressed, matching the documented rule.

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

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

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-15 04:04:22 -07:00
Anmol Mishra
9f694ab750
fix(studio): Windows GGUF cancel hang + CPU spinlock overhead (#5692) (#5749)
* fix(studio): Windows GGUF cancel hang + CPU spinlock overhead (#5692)

Two fixes for Windows-native GGUF inference via llama-server:

**Issue 1 — GPU/CUDA Hang on Stream Cancellation:**
- Add `Connection: close` header to all httpx requests proxying to
  llama-server, preventing Keep-Alive from masking downstream socket
  closure.
- Introduce `_await_disconnect_then_close` background watcher that
  polls `request.is_disconnected()` every 100ms and calls
  `resp.aclose()` immediately when the client disconnects. This runs
  alongside the existing cancel-POST watcher and covers client aborts
  that never reach the /cancel endpoint (tab close, proxy aborts,
  Colab, mobile navigation, etc.).
- Change all StreamingResponse `Connection: keep-alive` headers to
  `Connection: close`.

**Issue 2 — High CPU Spinlock & KV Cache Backup Overhead:**
- Set OMP_WAIT_POLICY=PASSIVE and OMP_NUM_THREADS=2 in the
  llama-server subprocess environment on Windows to prevent OpenMP
  from spin-waiting on all logical cores while the GPU decodes.
- Limit `--threads` to 2 on Windows when the model is fully
  GPU-offloaded (`-ngl -1`). Auto-detect otherwise.
- Pass `--cache-ram 0 --ctx-checkpoints 0 --no-cache-prompt
  --checkpoint-every-n-tokens -1` on Windows to disable prompt-cache
  snapshots that copy KV cache to system RAM over the WDDM/PCI-E bus.

Closes #5692.

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

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

* fix: use local import to avoid ruff F823 (sys used before assignment)

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

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

* review: address gemini review feedback

- Simplify _fully_gpu_offloaded init: default to False, only set True
  in the gpu_indices branch, drop redundant else.
- Log exceptions in _await_disconnect_then_close at debug level instead
  of silent pass, per review suggestion.

* Adjust review feedback for PR #5749

- _await_disconnect_then_close: set cancel_event before resp.aclose() so
  the streamer's RemoteProtocolError handler treats the watcher-driven
  close as cancellation, not an upstream error. Both call sites pass
  cancel_event through.
- Windows --cache-ram / --no-cache-prompt / --ctx-checkpoints block: gate
  on _fully_gpu_offloaded so CPU and partial-offload Windows runs keep
  prompt-cache reuse across turns.
- Windows OMP_WAIT_POLICY / OMP_NUM_THREADS env: same gate so CPU and
  partial-offload Windows runs keep default OpenMP parallelism.

* Shorten code comments touched by PR #5749

* Clean up local imports and rename underscore locals in PR #5749

- Drop the function-local `import sys as _sys` introduced as an F823
  workaround; remove the redundant in-function `import os`/`import sys`
  block so module-level imports resolve sys/os instead. F823 no longer
  triggers because no shadowing import remains inside load_model.
- Rename `_fully_gpu_offloaded` and `_t` to `fully_gpu_offloaded` and
  `threads_arg`. Underscore-prefixed names usually mean private/module-
  level; plain locals match Python style for in-function temporaries.

No behavior change. ruff clean, py_compile clean, 35 studio cancel-
infra tests + 13 launch-gating AST locks + 6 disconnect-watcher locks
+ 4 spoof live-import tests all pass.

* Fix Windows GGUF follow-ups for PR #5749

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

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

* Fix cache flag gating for PR #5749

* Fix Python 3.9 annotations for PR #5749

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

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

---------

Co-authored-by: Anmol Mishra <anmolx.work@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-15 10:32:10 +01:00
oobabooga
4176448fb8
Studio: enable stdio MCP servers on a loopback bind (#6295)
* Studio: enable stdio MCP servers on a loopback bind

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

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

* Studio: address codex review on stdio MCP loopback gate

* Studio: fix banner URL and preserve stdio MCP env opt-in on network binds

* Studio: scope loopback to exact aliases and honor force-disable on run_server reuse

* Studio: cover force-disable across a public re-bind and fix a stale test comment

* Studio: keep stdio MCP off on Colab loopback launches

* Studio: set tool policy before server startup

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
2026-06-15 03:02:32 +01:00
DoubleMathew
f372da407b
MLX Training updates (#5656)
* Expose MLX grad value clipping in Studio

* update test

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

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

* dataset ordering + wd

* fix mlx smoke step expectations

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

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

* cast norm activation output back to original input dtype

* address mlx studio review feedback

* Fix present-but-None seed override for PR #5656

studio/backend/core/training/worker.py
  `config.get("model_random_state", random_seed)` only fills the
  default when the key is absent. When a caller passes
  `config["model_random_state"] = None` explicitly (which happens
  any time a JSON payload sends an explicit `null`), the old code
  forwarded `None` to FastMLXModel and disabled deterministic init
  silently. Same for `lora_random_state`. Treat absent and explicit
  None the same way: fall back to random_seed.

studio/backend/tests/test_training_raw_support.py
  Update the source-string assertions to match the new lines.

* Guard optional MLXTrainingConfig fields and normalize random_seed for PR #5656

The MLX worker now passes `cast_norm_output_to_input_dtype` and
`dataset_order` only when the linked unsloth-zoo dataclass actually
declares them. Released zoo trees that predate the paired PR can still
construct `MLXTrainingConfig` without raising
`TypeError: unexpected keyword argument`. Once the dependency floor is
bumped to a release that contains both fields, the feature-detect
guards become no-ops.

`random_seed = config.get("random_seed", 3407)` was unguarded against
explicit `None` from raw / backend callers. The same value seeded the
trainer and was the fallback target for `model_random_state` /
`lora_random_state`. Normalize once at the top of the function and use
the normalized value everywhere so an explicit `None` cannot reach
FastMLXModel / get_peft_model / MLXTrainingConfig.

Existing seed source-pattern test updated to match the new normalize
helper. New test asserts the feature-detection guards exist and that
the unconditional kwargs do not include the gated fields.

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

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

* Normalize seed / cast / max_grad_value at TrainingBackend for PR #5656

Round-3 review consensus: the per-field guards that landed in the MLX
worker only protect the MLX path. The same `TrainingBackend.start_training`
config still reaches the CUDA/text trainer at `worker.py:2267`, the
embedding LoRA init at `worker.py:2450`, and embedding TrainingArguments
at `worker.py:2624` with raw `None` values, so an explicit
`random_seed=None` from a raw / backend caller still breaks non-MLX
training even after the previous fix.

Move the normalization into `TrainingBackend.start_training` itself,
where it runs once for every training mode:

- `_coerce_seed(value)`: explicit `None`, non-int, or absent all become
  3407. Every downstream worker now sees an int.
- `_coerce_optional_bool(value, default)`: explicit `None` falls back
  to `default` instead of `bool(None) == False`. Also normalizes the
  common raw-config / YAML string aliases ("true" / "false" / "0" /
  "1"). Used for `cast_norm_output_to_input_dtype`.
- `_coerce_optional_nonneg_float(name, value)`: rejects negative
  numerics from raw / backend callers, matching the Pydantic
  `ge=0` constraint the HTTP route already enforces. Used for
  `max_grad_value`.

worker.py MLX path: the existing `bool(config.get(key, True))` for
`cast_norm_output_to_input_dtype` was changed to also fall back on
explicit `None`, so direct worker callers (bypassing
`TrainingBackend.start_training`) are equally safe. `max_grad_value`
also raises on negative values inside the worker for the same reason.

TrainingStartRequest.random_seed default bumped from 42 to 3407 so
direct REST callers that omit the field receive the same default as
the Studio frontend and the MLX worker.

New regression test exercises the three new helpers across explicit
None, valid values, string aliases, and negative-value rejection.

* Tighten feature-detect test paren tracking for PR #5656

The block-extraction used , which stops at the
first inner closing paren (e.g. )
and would silently miss a future unconditional
/  added later in the same dict literal. Switched to
proper paren-depth tracking so the unconditional block is checked end-to-end.

* Shorten verbose comments in MLX Studio backend

* Handle MLX Studio EOS appending by mode

* Wire MLX leaf norm clipping through Studio

* Respect VLM layer filters for explicit LoRA targets

Rationale / guardrails for the local Studio/vision push:

When callers provide explicit VLM LoRA target_modules together with layer filters, FastVisionModel still needs to route the explicit targets through get_peft_regex. Otherwise the layer filters are ignored and adapters can be attached outside the requested language/vision scope.

Do not revert this to plain list(target_modules) for explicit module lists. The CUDA/Studio-facing contract is that explicit targets and layer filters compose: target_modules selects module names, while finetune_language_layers / finetune_vision_layers / finetune_attention_modules / finetune_mlp_modules constrain where those targets are allowed.

The regression test covers the language-only explicit q_proj case and source-checks that explicit targets are wrapped through get_peft_regex when filters are active.

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

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

* Refresh MLX smoke clip-config note for leaf_norm default

Trim the 11-line comment block to 5 lines and correct the stale claim
that MLXTrainingConfig defaults to max_grad_value=1.0. The new default
is max_grad_leaf_norm=1.0 (same memory profile as elementwise but
direction-preserving). The smoke still pins max_grad_value=1.0
explicitly to keep the 13-seed pass-rate fixture stable.

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

* Forward max_grad_leaf_norm through the training route and warn when layer filters constrain explicit target_modules for PR #5656

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han-Chen <info@unsloth.ai>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-06-14 04:58:50 -07:00
Daniel Han
a9a38da454
Studio: decide diffusion routing before the SWA resolver (#6299)
* Studio: decide diffusion routing before the SWA resolver

Loading a DiffusionGemma GGUF could fail with "llama-server failed to start. Check that the GGUF file is valid" while llama-diffusion-cli ran the same model fine.

_read_gguf_metadata set self._is_diffusion after calling _resolve_swa_pattern, both inside one try/except. The resolver reaches into transformers/HF, which can raise for an architecture transformers does not know (diffusion-gemma); the shared except then swallowed it and left _is_diffusion False, so the model was routed to plain llama-server instead of the diffusion runner. It only reproduced where the SWA pattern was not already cached/inline.

Set _is_diffusion right after the KV parse loop (before the resolver) so a resolver error can no longer drop the routing, and skip the resolver for diffusion models, which do not use Studio's SWA pattern.

* [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-13 04:22:50 -07:00
Daniel Han
368b19b237
Studio: fix training output dir escaping outputs root for models on another drive (#6293)
* Studio: derive training output dir from model basename for local-drive models

A LoRA/QLoRA run started from a model loaded by absolute path (common when
models live on a non-system drive, e.g. G:\modelsAI\...\gemma-4-12B-it) seeded
the default output dir with that full path. resolve_output_dir then raised
"path escapes root ... is not under the studio outputs folder", so training
could not start from a model stored off the system drive.

Add default_run_dir_name(): Hugging Face repo ids keep their namespace
(org/model becomes org_model), while local paths collapse to their final
component so an absolute source path can no longer leak into the output dir.
Use it at the three worker derivation sites and add a regression test.

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

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

* Studio: cap run dir name length and drop redundant resolve

Apply PR review feedback: length-cap the auto-generated output dir component so an unusually long model name stays under the filesystem name limit, and drop the redundant double resolve_output_dir at the embedding site so all three derivation sites match.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-13 04:06:17 -07:00
Daniel Han
cd270e2878
Studio: keep llama-server discovery from crashing on an access-denied candidate (#6268)
* Studio: keep llama-server discovery from crashing on an access-denied candidate

_find_llama_server_binary probed candidates with Path.is_file(), which raises
PermissionError (WinError 5) when a path exists but is momentarily inaccessible
(antivirus lock, an install replace in flight, an elevated-install ACL),
aborting model validation. Treat a denied-but-present path as the real binary
so discovery returns it; absent paths still skip.

* Retry a transiently locked binary instead of returning a denied path

Returning a still-denied path only moved the PermissionError to the next
is_file() (probe_server_capabilities). Retry briefly so a transient lock
clears and discovery returns an accessible path; on a persistent lock return
nothing rather than a path downstream cannot stat.

* Studio: do not fall back to another llama-server when a pinned one is locked

A denied LLAMA_SERVER_PATH made discovery skip the explicit pin and run a
lower-priority managed or PATH binary, so a load could silently use a stale or
incompatible server. Split the probe into a file/absent/denied status: when the
pinned path exists but stays access-denied, warn and stop rather than falling
back to a different executable.

* Studio: never downgrade past a denied pinned or managed llama-server

Extend the no-fallback rule beyond LLAMA_SERVER_PATH: a present-but-denied
UNSLOTH_LLAMA_CPP_PATH or managed ($STUDIO_HOME/llama.cpp, ~/.unsloth/llama.cpp)
binary now reports temporarily-unavailable instead of silently launching a
lower-priority legacy or PATH server. Shared _scan_pinned/_unavailable helpers;
legacy in-tree and PATH stay genuine fallbacks (a denied candidate there just
continues).

* Studio: let diffusion asset lookup use a locked llama-server path for its dir

DiffusionGemma does not run llama-server; _find_diffusion_assets only needs the
install dir to find the adjacent llama-diffusion-gemma-visual-server. The
no-fallback rule returning None on a transiently locked llama-server therefore
hid an available visual-server and raised 'runner not found'. Add an
include_denied option so diffusion lookup gets the locked path (its dir is all
it needs), while inference keeps the no-denied-path, no-downgrade behavior.

* Studio: report a locked llama-server as temporarily unavailable, not missing

When the pinned/managed binary stays access-denied through the retries, discovery
returns None and load_model raised 'binary not found', a terminal error that
points users at reinstalling rather than retrying a transient AV/install lock.
Reuse include_denied to detect the locked path and raise a distinct
temporarily-unavailable, retry message instead.

* Studio: GGUF preflight treats a locked llama-server as present

The pre-download preflight (and so /api/inference/validate) used the default
discovery, which returns None for a transiently access-denied binary, so it
raised 'binary not found' for a binary that merely needs the lock to clear. Use
include_denied so the existence check counts a locked binary as present; the
load itself still reports a still-locked binary as temporarily unavailable.
2026-06-12 11:20:07 -07:00
Lee Jackson
31439d9eed
Studio: extend llama.cpp first-token timeout (#5841)
* fix: extend llama.cpp first-token timeout

* fix: timeout label pluralization

* studio: distinguish llama stream timeout phases

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

* Fix/adjust timeout handling for PR #5841

* Fix lint failure for PR #5841

* Fix/adjust stream timeout handling for PR #5841

* Fix/adjust first token timeout for PR #5841

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

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

* Fix/adjust passthrough timeouts for PR #5841

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

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

* Fix/adjust preheader stream cancellation for PR #5841

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

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

* Fix/adjust timeout PR diff for PR #5841

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

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

* Fix/adjust Python 3.9 stream iteration for PR #5841

* Fix first body timeout for PR #5841

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

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

* Fix first token timeout deadlines for PR #5841

---------

Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
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-12 18:41:38 +02:00
Daniel Han
2a01ba8aad
DiffusionGemma: set UNSLOTH_IS_PRESENT for the shim subprocess (#6259)
A clean install could not run DiffusionGemma: the runner spawns
python -m unsloth_zoo.diffusion_studio.shim, and unsloth_zoo refuses to
import unless UNSLOTH_IS_PRESENT is set (normally by import unsloth). The
shim never imports unsloth, so the subprocess died with
'Please install Unsloth via pip install unsloth!' and the model load
failed with a 500. Set the flag in the runner child env, as unsloth does
on import.
2026-06-12 07:35:59 -07:00
hoobnn
f033213c0b
Studio: account for mmproj VRAM in GGUF fit budget (#5825) (#5849)
* Studio: account for mmproj VRAM in GGUF fit budget (#5825)

Vision GGUFs load the mmproj projector onto the GPU via --mmproj
alongside the weights, but the context auto-sizing / GPU-selection
budget sized off _get_gguf_size_bytes(model_path), which counts only
the weight file(s). The projector was never added, so the budget was
too optimistic: context got mis-estimated and tight vision loads
spilled to system RAM / OOM'd.

Resolve the launch projector once before GPU selection and fold its
size into the fit budget. The same resolved path feeds both the budget
and the --mmproj launch flag, so the two cannot disagree. The summary
log now reports the projector size separately, keeping "GGUF size"
accurate.

Adds _mmproj_vram_bytes() + unit tests (no GPU / network / subprocess).

* Studio: simplify mmproj summary-log concatenation (#5825)

Address review: the summary log mixed explicit `+` with implicit
f-string concatenation. Extract the optional projector fragment into
`mmproj_note` so the logger.info uses uniform implicit concatenation.
No behavioral change.

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

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

* Studio: trim mmproj VRAM comments

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-12 15:04:08 +01:00
Daniel Han
756265e3fd
DiffusionGemma: disable tools, enable artifacts canvas by default (#6255)
DiffusionGemma serves via the visual runner, which streams per-step
canvas frames so the answer resolves live in the bubble. The agentic
tool loop (generate_chat_completion_with_tools) does not forward those
frames, so whenever a tool pill (Search/Code) was on the live canvas
silently vanished while text still streamed. DiffusionGemma is not a
tool-calling target anyway, so report supports_tools=False for it: the
chat always takes the frame-forwarding path, and the Search/Code pills
disable themselves (a local model has no builtin web search either).

Also turn the artifacts canvas on by default for DiffusionGemma so a
full-HTML answer (e.g. a playable game) renders as an interactive
sandboxed card without the user flipping the global artifacts toggle.
2026-06-12 06:55:14 -07:00
Michael Han
99704ffe47
Studio: project sources backed by RAG (#6205)
* Studio: make project sources work with RAG and polish project UI

Projects had a disabled Sources tab with an Add sources placeholder.
This wires it up end to end on top of the RAG engine:

- Add a project scope to the RAG store, ingestion and retrieval
- New endpoints: POST/GET /api/rag/projects/{id}/documents
- search_knowledge_base resolves kb, project and thread scopes; an
  explicit KB stays exclusive, project and thread scopes combine
- Multi-scope search: FTS uses scope IN (...), vec0 KNN runs per
  scope and merges by cosine score
- Lazy ALTER TABLE adds documents.project_id on existing databases
- Deleting a project also removes its indexed sources
- Sources tab now uploads with progress chips and drag and drop
- Chats inside a project auto-enable retrieval over project sources
  when the project has indexed documents (cached probe, no Docs pill
  needed); external providers still never receive rag_scope

UI polish:
- Rounder project cards with folder icon chip and softer shadow
- Project header icon in a rounded chip
- Chats/Sources pills and Add sources button without borders

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

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

* Studio: match Add sources button shadow to the chat composer in light mode

* Studio: round project switcher hover pill and pad the folder icon

* Studio: remove border from project sources box

* Studio: grey hover on project cards and menu, move search into header, widen page spacing

* Studio: shorten sources copy, white header pills with composer shadow, fixed-width search, hub-size page headings

* Studio: align project landing blocks to the composer width

* Studio: restore muted background and flat look on projects header controls

* Studio: darker grey hover on project cards in light mode

* Studio: soften project card hover grey

* Studio: keep project card menu button visible while its menu is open

* Studio: drop focus outlines and rings on buttons and clickable icons, keep input focus styles

* Studio: address review feedback on project sources

- Remove uploaded files from disk when a project is deleted, confined
  to the uploads root
- 404 project uploads when the project does not exist, matching the KB
  endpoint
- Guard lexical search against an empty scope list
- Re-invalidate the project sources probe after uploads and removals
  settle so a chat sent mid-upload cannot cache a stale negative
- Keep keyboard focus rings: only mouse focus drops the Tailwind ring,
  the browser default outline stays removed

* Studio: add a green New badge to the project Sources tab

* Studio: unify New pills, fully round with soft emerald fill and no border

* Studio: a touch more vertical padding on New pills

* Fix project RAG source edge cases for PR #6205

* Fix duplicate RAG upload cleanup for PR #6205

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-06-12 15:42:51 +02:00
oobabooga
522f308be2
Studio: cache MCP tool discovery instead of re-probing every chat send (#5828)
* Studio: cache MCP tool discovery instead of re-probing every chat send

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

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

* Stop re-probing offline/down MCP servers every time

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

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

* Add tests for mid-probe delete and OAuth cool-off paths

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

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

* Don't cool-off a server edited or deleted mid-probe

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

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

* Guard MCP refresh cache writes

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
2026-06-12 14:18:06 +01:00
ashzak
aefe904d66
feat(studio): implement S3 dataset loading (completes #5951) (#6222)
* feat(studio): add S3 dataset configuration foundation (#4539)

Add foundational types and configuration for S3 bucket dataset loading:

- Add S3Config type to frontend training types
- Add S3Config Pydantic model to backend training models
- Add "s3" as a DatasetSource option
- Add s3Config state and setS3Config action to training config store
- Add i18n translations for S3 configuration (English and Chinese)

This provides the type definitions and UI text for S3 integration.
Full implementation requires boto3 dependency and data loading logic.

Refs: #4539

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

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

* Wire S3 config into training pipeline and prevent secrets persistence

- Pass s3_config from request into training_kwargs so it flows to training subprocess
- Add s3Config to NON_PERSISTED_STATE_KEYS to prevent AWS secrets from being
  saved to localStorage

Addresses code review feedback on PR #5951.

* Exclude S3 config from database persistence to protect secrets

Filter out s3_config (which contains secret_access_key) from the
config_json stored in training_runs table, preventing AWS credentials
from being persisted to disk.

Addresses P1 security feedback on PR #5951.

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

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

* Re-raise HTTPException in start_training and defer s3 DatasetSource widening for PR #5951

* Redact s3_config from W&B run config and accept camelCase S3 credential aliases for PR #5951

* feat(studio): implement S3 dataset loading end-to-end

Builds the actual S3 loader on top of the hardened #5951 foundation,
turning the 501-gated scaffold into a working dataset source.

Backend:
- Add core/training/s3_dataset.py: lists and downloads supported dataset
  files (parquet/json/jsonl/csv) from an S3 bucket to a temp dir, using
  IAM-role or access-key credentials. boto3 is imported lazily (optional dep).
- Wire s3_config into UnslothTrainer.load_and_format_dataset (downloads then
  reuses the existing local-file path) and thread it through worker.py.
- Replace the 501 "not implemented" gate with a boto3-availability guard so
  S3 works when boto3 is present and fails clearly when it is not.
- Add boto3 to studio.txt requirements.
- Add tests/test_s3_dataset.py (8 tests) covering download/filtering,
  collisions, missing-boto3, and S3Config camelCase/IAM validation.

Frontend:
- Widen DatasetSource to include "s3"; add s3_config to the training payload
  type and mapper; add an S3 validation branch and selectS3Source store action.
- Add s3-config-form.tsx (bucket/region/prefix/keys/IAM toggle) reusing the
  existing studio.dataset.s3.* i18n strings.
- Add a Hugging Face / Local / Amazon S3 source toggle in dataset-section;
  the S3 config card replaces the dataset combobox when S3 is selected.
- Fix DatasetPreviewDialog to accept the widened DatasetSource type.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

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

* Fix S3 dataset loader for PR #6222

* Fix S3 dataset edge cases for PR #6222

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

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

* Fix S3 IAM payload handling for PR #6222

* Block multimodal S3 datasets for PR #6222

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Ash <ash@MacBook-Pro.local>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-06-12 14:52:04 +02:00
Daniel Han
a70146df0f
Studio: bundle Gemma 4 chat templates (E2B/E4B + larger) and auto-apply to unsloth/gemma-4-*-GGUF (#6245)
* Studio: override chat template for unsloth/gemma-4-*-GGUF with bundled gemma-4.jinja

The chat templates baked into the shipped unsloth/gemma-4-*-GGUF quants predate
Google's gemma-4 chat-template PR #118 and lack the preserve_thinking flag, so
Studio cannot surface the "Preserve thinking" toggle for Gemma 4. Bundle the updated
template and override the embedded one at llama-server launch via --chat-template-file,
scoped to the gemma-4 GGUF family, so users do not need to re-download any quant.

- Add studio/backend/assets/chat_templates/gemma-4.jinja (PR #118 based;
  preserve_thinking defaults false, the one deliberate divergence from upstream).
- Add core/inference/chat_templates.py: gemma-4 GGUF matcher plus an
  effective-override resolver (explicit user template still wins).
- Wire the resolver into routes/inference.py ahead of the reload-dedup check and
  both load_model calls so the live backend and the incoming request compare against
  the same template text (no spurious reloads).
- Default preserve_thinking off in the launch-time chat_template_kwargs so direct
  API callers match the UI default.
- Ship the asset via package-data and add unit tests.

* Studio: ship E2B/E4B edge variant of the bundled Gemma 4 template

Google ships two distinct gemma-4 chat templates: E2B and E4B omit the empty
"<|channel>thought<channel|>" block on enable_thinking=false, while the
12b/26B-A4B/31B family emits it (confirmed against google/gemma-4-E2B-it,
-E4B-it, -12b-it, -26B-A4B-it, -31B-it; the two families differ only in that
one block). The single PR #118 based template followed the larger-model
behavior, which is wrong for the E2B/E4B GGUFs this feature most targets.

- Add studio/backend/assets/chat_templates/gemma-4-edge.jinja: identical to
  gemma-4.jinja minus the empty-thought-block, matching E2B/E4B behavior.
- Route unsloth/gemma-4-E2B-it-GGUF and -E4B-it-GGUF to the edge template;
  12b/26B-A4B/31B keep gemma-4.jinja.
- Extend tests for the edge matcher, per-family routing, and the empty-thought
  block difference (off for edge, on for standard).

* Studio: address review feedback on the gemma-4 template override

- Normalize owner-less shorthand model ids in the template matcher: a bare
  "gemma-4-E2B-it-GGUF" is canonicalized to "unsloth/" the same way
  ModelConfig.from_identifier does, so shorthand loads still get the override
  (and the preserve_thinking capability) instead of falling back to the
  embedded template.
- Scope the test's module stubs with unittest.mock.patch.dict instead of
  sys.modules.setdefault, and only stub deps that are missing, so the global
  module registry is not polluted for tests that run afterwards.
- Guard the Jinja render tests with pytest.importorskip("jinja2") so the suite
  stays runnable in minimal Studio environments where jinja2 is not present.
- Add tests for shorthand resolution.

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

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

* Studio: address 10-reviewer P1 findings on the gemma-4 template override

- /status no longer surfaces Studio's auto-applied bundled template as a
  user-authored chat_template_override. The frontend adopts that field as
  editable state and would otherwise re-send the gemma-4 template as an explicit
  override for a later, unrelated model. /status now reports None when the live
  override equals the model's auto-resolved bundled template.
- When a bundled family template is in effect, strip an inherited
  --chat-template-file from llama_extra_args too (not only when the raw request
  set chat_template_override). Otherwise a stale inherited template, appended
  last, shadows the bundled one while Studio reports the bundled template's
  capabilities.
- Write the temp chat-template file as UTF-8 explicitly, and keep the bundled
  templates ASCII (replaced em dashes), so non-UTF-8 Windows locales cannot raise
  UnicodeEncodeError or emit a mis-encoded template. Added an ASCII guard test.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-12 05:49:39 -07:00
Daniel Han
90cb9499e8
Studio: serve DiffusionGemma with live in-place denoising and honest stats (#6250)
* Studio: serve DiffusionGemma GGUFs with the on-device visual decoder

* Studio: render the DiffusionGemma denoising canvas live in chat with honest stats

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

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

* Studio: harden DiffusionGemma runner resolution (Windows .exe, build/bin lookup, clear stale audio flag, safe PYTHONPATH, Linux-only pdeathsig)

* [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-12 05:48:06 -07:00
Daniel Han
240c0c3500
Studio: fix WSL Strix Halo GPU on reinstall (ROCDXG drop-in + system HIP before bundle) (#6227)
* install.sh: persist ROCm-on-WSL drop-in even when rocminfo already works

_maybe_bootstrap_rocm_wsl calls _ensure_rocm_probe_env (which exports a
transient HSA_ENABLE_DXG_DETECTION + adds /opt/rocm/bin to PATH on the
installer process) right before the "rocminfo enumerates gfx1151 -> already
set up, return early" gate. On any reinstall over an existing /opt/rocm --
the common case, since the uninstaller keeps shared ROCm userspace but
removes /etc/profile.d/unsloth-rocm-wsl.sh -- that probe env makes rocminfo
succeed, so the gate returns 0 WITHOUT ever persisting the drop-in. The
transient env dies with the installer, so the next login shell (Studio,
llama-server) sees no GPU: torch cuda_avail=False, rocminfo finds nothing,
the llama.cpp ROCm prebuilt segfaults on a GPU it can't reach.

Factor the drop-in writer into _persist_rocm_wsl_dropin() and call it before
the early return so the persistent env is restored whenever librocdxg is
present. Idempotent (only writes when the drop-in is missing), gated on
librocdxg so it never fires on non-WSL/non-ROCDXG hosts, root-writes or
sudo-tees like before. The fast-path branch now reuses the same helper.

Reproduced on gfx1151 (Radeon 8060S) under dash (the curl|sh shell):
before the fix a reinstall left the drop-in absent and torch cuda_avail
False; after, the drop-in is persisted and a fresh login shell reports
cuda_avail True. Verified under both dash and bash, and idempotent on
re-run.

* Studio WSL: load system HIP before a prebuilt's bundled runtime (gfx1151)

The lemonade / published llama.cpp ROCm prebuilts bundle their own HIP
runtime (libamdhip64) built for bare-metal Linux. In WSL the GPU is reached
through the system ROCm's librocdxg bridge over /dev/dxg, which the bundled
runtime cannot drive -- it segfaults on the first GPU call. So:

  - install_llama_prebuilt.py: the prebuilt's llama-quantize/llama-server
    validation runs with the bundle dir first on LD_LIBRARY_PATH, segfaults
    (empty stderr), and the install silently falls back to a CPU source build
    (which on this host can't even build for GPU -- hipcc absent). The Strix
    Halo WSL user ends up on CPU despite a working GPU.
  - llama_cpp.py: even if a GPU prebuilt were kept, the serve-time launcher
    put the bundle dir first too, so it would crash at load.

Fix: on a ROCDXG WSL host (gated on /dev/dxg + "microsoft" /proc/version +
a librocdxg-providing /opt/rocm), prepend the system ROCm lib dir to
LD_LIBRARY_PATH so the WSL-capable libamdhip64 + librocdxg load first, while
the bundle still supplies libggml-hip / librocblas with the gfx1151 kernels.
Set HSA_ENABLE_DXG_DETECTION=1 alongside. Added _wsl_system_rocm_lib_dirs()
to both modules (kept identical so a prebuilt that passed install validation
runs the same way at serve time). Strict no-op on bare-metal Linux, NVIDIA,
macOS, and Windows.

Verified on gfx1151 (Radeon 8060S) in WSL (ROCm 7.2.1 + librocdxg, Adrenalin
ROCDXG): before, the lemonade gfx1151 prebuilt segfaulted and the install
fell back to a broken CPU build; after, install_llama_prebuilt validates and
keeps the GPU prebuilt (source=published, prebuilt_fallback_used=False), and
Studio serves Qwen3-1.7B-GGUF at 53 tok/s with the model resident in GPU
memory (llama-server device_info: ROCm0 = AMD Radeon 8060S).

* tests: cover the WSL ROCDXG drop-in + system-HIP-ordering fixes

- _wsl_system_rocm_lib_dirs: no-op without /dev/dxg, on bare-metal Linux,
  and on WSL without librocdxg; returns the system lib dir on a ROCDXG WSL
  host.
- binary_env: prepends the system ROCm lib dir ahead of the bundle and sets
  HSA_ENABLE_DXG_DETECTION on WSL; unchanged on bare-metal Linux.
- install.sh: _persist_rocm_wsl_dropin exists, is gated on librocdxg, and the
  rocminfo-already-works early return calls it before returning.
- llama_cpp.py: the serve-time launcher prepends the WSL rocm dirs before the
  bundle dir (mirrors binary_env).

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

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

* Tighten WSL ROCDXG fix comments (no logic change)

Condense the drop-in / system-HIP-ordering comments and docstrings added in
this PR. Verified comment-only via AST parse + py_compile + sh/bash -n, the
308-test rocm_support suite, and a dash functional re-run of the bootstrap
(drop-in still persisted, env still set).

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-12 04:46:39 -07:00
Tai An
b72cf0af24
fix(studio/responses): forward chat_template_kwargs enable_thinking to chat request (#6202)
* fix(studio/responses): forward chat_template_kwargs enable_thinking to chat request

The /v1/responses translation in _build_chat_request dropped
chat_template_kwargs (e.g. {"enable_thinking": true}) sent via the
Responses extra-body, so reasoning control was silently ignored.
Lift enable_thinking onto the typed ChatCompletionRequest field,
mirroring openai_chat_completions, so both the non-streaming and
streaming Responses pass-through paths honor it.

Fixes #6198

Signed-off-by: Tai An <antai12232931@outlook.com>

* Fix/adjust Responses reasoning for PR #6202

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

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

* Fix/adjust reasoning none for PR #6202

* Fix/adjust structured reasoning for PR #6202

* Fix/adjust responses reasoning review findings for PR #6202

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

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

* Fix/adjust responses reasoning follow-ups for PR #6202

* Fix/adjust think parsing gate for PR #6202

---------

Signed-off-by: Tai An <antai12232931@outlook.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-12 13:20:07 +02:00
Lee Jackson
7307fde839
Studio: Add custom provider option to Connections (#6112)
* feat: add custom connection

* Fix custom provider handling for PR #6112

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

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

* Fix custom provider connection test for PR #6112

---------

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-12 13:09:35 +02:00
oobabooga
72e67ae5a6
Studio: Add Tensor-Parallel llama.cpp support (#6040)
* Studio: Add Tensor-Parallel llama.cpp support

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

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

* Studio: harden Tensor-Parallel fallback and GPU selection

* Studio: reconcile split-mode extras and harden tensor-split planning

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

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

* Studio: reconcile split-mode extras in backend duplicate-load guard

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

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

* Studio: preserve inherited non-tensor split modes on reload

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

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

* Studio: honor cancellation in tensor fallback, preserve tensor mode on rollback, and don't raise an explicit small context

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

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

* Studio: reconcile split-mode in reload check and strip it on tensor downgrade

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

* Strip --tensor-split alongside --split-mode so inherited ratios don't override the tensor planner

An inherited or stale --tensor-split in llama_extra_args was appended after
Studio's computed --tensor-split and won last in llama.cpp, re-introducing the
asymmetric-GPU OOM tensor mode is meant to prevent. Group -ts/--tensor-split
into the split-mode shadow set so it is stripped on inherit and on the layer
fallback; parse_split_mode_override still keys on the mode value only.

* Drop quantized KV for the tensor attempt and report native max context

Tensor mode aborts on a quantized KV cache, so a user with q8_0/q4_1 etc. who
enabled Tensor Parallelism silently fell back to layer split. Clear the cache
type (and strip inherited/explicit --cache-type) for the tensor attempt only;
the layer fallback re-runs with tensor off and keeps the user's choice.

Also report max_available_ctx from the native context, not an explicit small
-c, so the context slider no longer warns too early in tensor mode.

* Reconcile inherited split-mode extras in the already-loaded check

When a same-model load omitted llama_extra_args, the tensor comparison resolved
the raw (None) request and treated an inherited --split-mode tensor server as a
mismatch, forcing a needless reload. Compare using the stored extras stripped
the same way the reload strips them.

* Pass tensor_parallel through compare-mode loads

The generalized compare path loaded each GGUF without tensor_parallel, so
compare ran layer split even with the toggle on and left the settings sheet
stale. Send the toggle and hydrate the loaded state from the response, matching
the main chat and recipe load paths.

* Add --tensor-parallel flag to unsloth studio run

The headless one-liner could only reach tensor mode by passing --split-mode
tensor as a raw llama.cpp extra. Add a first-class --tensor-parallel/
--no-tensor-parallel option that sets the tensor_parallel field on the
/api/inference/load payload, forwarded through the studio-venv re-exec like the
other polarity flags. Matches the web UI toggle and the API field.

* [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: danielhanchen <michaelhan2050@gmail.com>
2026-06-12 04:00:52 -07:00
Anmol Mishra
554c289538
fix: respect absolute export paths to prevent cross-drive copy failures (WinError 112) (#6088)
* fix: allow absolute save_directory in export paths to prevent cross-drive copy failures

The GGUF export pipeline (and all other export flows) forced every
save_directory through resolve_export_dir(), which always resolved
the path under exports_root() — typically ~/.unsloth/studio/exports/
on the system drive (C: on Windows).

When a user selected an output directory on a different drive (E:):
1. The absolute path was rejected at the Pydantic validator level.
2. Even if it got through, resolve_export_dir would re-resolve it
   under C:\Users\.unsloth\studio\exports\.
3. After GGUF conversion completed on E:, the relocation step would
   try to move/copy the finished files to C:, causing:
   - WinError 17 (cross-drive move failure when shutil.move falls
     through to a cross-filesystem copy)
   - WinError 112 (disk full on C:)

Fix both layers:
- _validate_save_directory: accept absolute paths (they represent an
  explicit user choice of output location).
- resolve_export_dir, resolve_output_dir, resolve_tensorboard_dir:
  return absolute paths as-is instead of forcing them under the
  default root. Keep the existing safety checks (null bytes, '..'
  segments) and fall through to resolve_under_root for relative paths.

Fixes: https://github.com/unslothai/unsloth/issues/6082

* refactor: centralize user path validation into _resolve_user_path helper

Addresses code review feedback: the null-byte, '..', and absolute-path
checks were duplicated across resolve_output_dir, resolve_export_dir,
and resolve_tensorboard_dir. Extract a single _resolve_user_path helper
that all three delegate to.

No behavioral change — pure consolidation.

* fix: address code review — contain destructive cleanup and scope absolute paths

Address all review feedback from gemini-code-assist:

1. P1: destructive subdirectory cleanup (export_gguf)
   The flattening loop in export_gguf previously rmtree'd every
   subdirectory under abs_save_dir. When targeting an existing user
   directory on a different drive (#6082), this could nuke unrelated
   subdirectories. Now snapshot existing subdirectories before the
   export and only clean up dirs created during this run.

2. P2: keep scan/read endpoints contained
   Only resolve_export_dir accepts absolute paths (export is a write
   path where user picks location). Reverted resolve_output_dir and
   resolve_tensorboard_dir to use resolve_under_root directly — these
   are used by scan/read/training endpoints that must stay contained
   under their respective roots.

3. Centralization feedback
   Removed the _resolve_user_path helper since it's no longer needed
   with the narrowed scope. resolve_export_dir has the absolute path
   logic inline with a clear docstring.

* fix: skip pre-existing subdirs in GGUF flatten loop and clean stale export intermediates

Two issues caught in code review (chatgpt-codex-connector):

1. The flattening loop moved ALL .gguf files from ALL subdirectories
   into abs_save_dir, including pre-existing unrelated user subdirs.
   Now skip pre-existing subdirs entirely unless they are known
   export-owned intermediates (model/, model_gguf/).

2. After a failed export, known export-owned subdirectories (model/,
   model_gguf/) were snapshotted as pre-existing on retry and never
   cleaned up. These are now always cleaned up regardless, since they
   are known intermediates created by the export pipeline.

* fix: separate write vs read export paths, guard same-dir rmtree

Three issues caught in code review (chatgpt-codex-connector):

1. P1: scan endpoint containment
   resolve_export_dir was changed to accept absolute paths, but it's
   also used by scan/read endpoints (routes/models.py) that must stay
   contained under exports_root(). Split into:
   - resolve_export_dir: contained, used by scans
   - resolve_export_write_dir: accepts absolute paths, used by export
     backend only

2. P1: same-directory rmtree
   When a non-PEFT checkpoint's gguf_dir resolves to the same path as
   abs_save_dir (user selected the checkpoint's gguf output as their
   export directory), shutil.rmtree(gguf_dir) would delete the user's
   chosen output directory. Now skip relocation when both paths resolve
   to the same location.

3. P1: pre-existing subdir flatten loop
   Reverted _EXPORT_OWNED_SUBDIRS logic — 'model/' and 'model_gguf/'
   are common directory names in shared model folders and don't prove
   export ownership. Now only clean up subdirs that didn't exist before
   the export started.

* fix: remove dead _EXPORT_OWNED_SUBDIRS and fix _export_details for absolute paths

Two fixes from review comments:

1. Remove unused _EXPORT_OWNED_SUBDIRS declaration (leftover from
   previous iteration that was intentionally removed).

2. _export_details now returns the full absolute path when the export
   target is outside exports_root(), instead of truncating to basename.
   Users who export to E:\ can now see the full destination path in
   the success dialog.

* fix: use unique tmp dir for GGUF intermediates to avoid overwriting user dirs

When exporting to an absolute destination that already contains a
model/ subdirectory (e.g. a shared models folder), the hard-coded
model_save_path would overwrite files in that unrelated directory.

Use _tmp_model_<uuid> as the intermediate path instead, so user
directories are never touched. The tmp dir is created as a new subdir
of abs_save_dir and cleaned up by the flatten loop after GGUF files
are relocated.

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

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

* Fix GGUF local export paths for PR #6088

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

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

* Address GGUF export follow-ups for PR #6088

* Clean GGUF temp dirs on export failure for PR #6088

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

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

* Fix/adjust export path tests for PR #6088

* Fix/adjust export path review findings for PR #6088

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

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

* Fix/adjust home export path handling for PR #6088

---------

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-12 12:52:57 +02:00
Daniel Han
25ccfebc0b
Studio: tune llama.cpp env for data-center GPUs (#6098)
* Studio: tune llama.cpp env for data-center GPUs

Detect datacenter/professional NVIDIA GPUs at llama-server launch and set
the llama.cpp env flags that help them, gated so consumer GeForce, AMD/ROCm,
CPU and macOS are never touched.

- GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F=1 for any DC GPU (FP32 cuBLAS
  accumulation). On a B200 this is ~0% throughput cost with identical
  perplexity (7.3230 wikitext-2-raw, baseline and on), where on GeForce the
  same flag costs real throughput, hence the gate.
- GGML_CUDA_P2P=1 and CUDA_SCALE_LAUNCH_QUEUES=4x for multi-GPU DC boxes.
  Benchmarked on 6x B200: +33-51% prompt processing on tensor (row) split and
  +8-16% on the default pipeline (layer) split, with no regression on the
  other split or on token generation.

Detection uses torch device names (A100/A30/H100/H200/H800/GH200/B200/GB200/
GB300/L40/L4/RTX PRO 6000/RTX 6000 Ada). A mixed box with one consumer GPU in
the selection is treated as non-DC. All writes are setdefault so a user value
always wins, and UNSLOTH_DISABLE_DC_TUNING=1 turns the whole thing off.

37 unit tests cover detection, multi-GPU gating, user-override precedence, the
disable flag and fail-open on error.

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

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

* Studio: fix data-center GPU detection false positives and physical-id mapping

Two issues in the data-center llama.cpp env tuning gate:

- _is_datacenter_gpu matched the marker allowlist as unbounded substrings, so
  workstation/laptop parts "NVIDIA RTX A1000" and "NVIDIA RTX A3000" matched
  "a100"/"a30" and were wrongly tuned as data-center GPUs (forcing FP32 cuBLAS
  accumulation and the multi-GPU env, which carry a real cost on those cards).
  Switch to a word-boundary regex.

- gpu_indices carries physical GPU ids (translated from torch ordinals by
  _get_gpu_free_memory via CUDA_VISIBLE_DEVICES), but they were passed straight
  into torch.cuda.get_device_properties, which expects mask-relative ordinals.
  On a masked host (e.g. CUDA_VISIBLE_DEVICES=4,5,6,7) a selection like [4,5]
  fell out of range and silently dropped the tuning, and on a mixed mask it could
  probe the wrong GPU class. Build a physical-id to device-name map mirroring
  _get_gpu_free_memory, then look up the selection by physical id.

Add regression tests for the A1000/A3000 false positives and for masked-host
physical-id selection (reordered and mixed-class masks included).

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

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

* Studio: tighten data-center GPU tuning comments

Comment-only pass over the DC tuning block and its tests: shorten verbose
docstrings/comments, drop ones that restate the code, collapse multi-line
blocks. Keep the load-bearing rationale (physical-id vs ordinal mapping, the
word-boundary reason, the B200 benchmark numbers). No code change: verified
with comment_tools.py check --strip-docstrings (code unchanged, comments only).

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-12 02:39:01 -07:00
Irakli
95a2627bf6
Fix step count mismatch when sequence packing is enabled (#5967)
* Fix step count mismatch when sequence packing is enabled

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

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

* Emit a single step-0 progress event and guard applyStatus totalSteps

Merge the two consecutive _update_progress calls before train() so the
step-0 gate in _on_progress fires once instead of twice, avoiding a
duplicate startup event and a null-metric step-0 row in training_metrics.

Apply the same positive-number guard to applyStatus that applyProgress
uses, so a stale or startup status poll can no longer overwrite the
packed step count with 0 or replace it with a stale total.

* Log debug message when train_dataset length is unavailable

The TypeError fallback for length-less datasets (e.g. streaming
IterableDataset) was silent, leaving no trace that the step estimate
came from the raw dataset rather than the packed one.

* [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: Etherll <61019402+Etherll@users.noreply.github.com>
2026-06-12 12:30:53 +03:00
oobabooga
7f2986a413
Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls (#5869)
* Studio: Add inline confirmation (Allow/Always allow/Deny) for tool calls

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

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

* Fix race in tool-call confirmation gate

* Studio: gate built-in tool calls and harden the confirmation handshake

The Allow / Always allow / Deny controls only lived in the fallback tool
card, but the built-in tools (web search, python, terminal, code
execution, image generation) render with their own components and so
never showed the buttons. Those calls paused after tool_start with no way
to approve them, hanging until the 1 hour timeout. Only MCP tools, which
use the fallback renderer, actually worked.

Render the controls for every tool card by wrapping each registered tool
component (and the fallback) in thread.tsx with a shared
ToolConfirmationControls, so the gate applies uniformly.

Also make the handshake robust:
- The gate keys on a per-call approval_id minted by the backend and
  echoed in tool_start, instead of session_id alone, so a stale or
  concurrent confirmation can no longer resolve the wrong call.
- The approval slot is registered before tool_start is yielded, closing
  the race where a fast click or an auto "Always allow" could reach the
  backend before the waiter existed.
- The frontend resolves with the same session id the request was sent
  with (plus the approval_id), fixing the new-thread mismatch where the
  confirmation targeted a different session than the blocked stream.
- The confirm endpoint returns {resolved}; the UI keeps the buttons and
  shows a retry hint until the backend confirms a match, instead of
  hiding them on a failed or mistargeted post.
- The gate runs after the disabled-tool and duplicate-call checks, so a
  call that will not execute is not put up for approval. A denied call is
  still excluded from duplicate detection, so re-issuing and approving it
  works.
- "Always allow" is scoped per session to match the backend gate.

Add backend tests for the approval registry, the SSE no-deadlock
handshake, and the loop integration (allow, deny, disabled, duplicate,
re-issue after deny).

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

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

* Move "Confirm tool calls" to the Tools section

* Studio: Keep tool group open while a tool call awaits confirmation

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

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

* Fix tool confirmation session scope for PR #5869

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

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

* Fix confirmation follow-ups for PR #5869

* Apply pre-commit formatting for PR #5869

* Fix confirmation cleanup for PR #5869

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

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

* Harden confirmation lookups for PR #5869

* Studio: make the tool-call confirmation decision immutable

resolve_tool_decision accepted a second confirmation for the same approval_id
and overwrote slot["decision"] in the window before the waiter reads it and
pops the slot, so a duplicate or out-of-order POST could flip an Allow to Deny
(and returned a misleading resolved:true). Reject once the slot's event is
already set so the first decision wins. Adds a regression test.

* Fix/adjust tool confirmations for PR #5869

* [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: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
2026-06-12 10:55:26 +02:00
James Dawdy
f22e890ab8
fix(studio): inherit llama_extra_args and honor --no-mmproj (#5902)
* fix(studio): inherit llama_extra_args and honor --no-mmproj

Reloading the same GGUF from the UI without gguf_variant no longer drops
CLI pass-through args like --no-mmproj. Skip mmproj download and launch
when --no-mmproj is present in llama_extra_args.

Co-authored-by: Cursor <cursoragent@cursor.com>

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

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

* fix(studio): tighten GGUF llama_extra_args variant inheritance guard

Reject inherited CLI args when the request changes gguf_variant or when
omitted variant resolves differently from the stored extra_args source.

Co-authored-by: Cursor <cursoragent@cursor.com>

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

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

* Treat --no-mmproj-auto and --mmproj-auto with last-wins parsing for PR #5902

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
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-12 00:27:04 -07:00
alkinun
672d8f0581
Expose runtime context length for hub models (#6154)
* expose runtime context length for hub models

* runtime context helper review

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

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

---------

Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-11 22:13:53 +03:00
Viktor Ferenczi
120daf9d8b
fix(studio/rocm): don't stack ROCR_VISIBLE_DEVICES on HIP_VISIBLE_DEVICES (#6176)
When pinning GPUs for the llama-server child, the ROCm path set both
HIP_VISIBLE_DEVICES and ROCR_VISIBLE_DEVICES to the same physical
indices. These masks filter at different layers and stack:
ROCR_VISIBLE_DEVICES reduces the visible set at the HSA/ROCr layer and
re-indexes from 0, then HIP_VISIBLE_DEVICES indexes into that reduced
set. _select_gpus ranks by free VRAM and picks the most-free card, so a
single non-zero pin (e.g. "1") becomes out of range at the HIP layer,
HIP enumerates 0 devices, and the model silently runs on CPU
("ggml_cuda_init: failed to initialize ROCm: no ROCm-capable device is
detected").

Set only HIP_VISIBLE_DEVICES (which narrows correctly on its own) and
clear any inherited ROCR mask so it can't double up.

Verified on a 2x Radeon AI PRO R9700 (gfx1201) host, ROCm 7.1.1: the
same selected=[1] load that fell back to CPU (~7.7 tok/s) now runs on
the GPU (~78 tok/s).

Fixes #6175

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-11 16:21:48 +01:00
Nilay
004577c9cd
studio: show MCP "Import config" on the add-server form (#6030)
* studio: import MCP servers from a config file

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

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

* import config' on the add-server form

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

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

* fix: defensively handle MCP config imports

* fix: address MCP import review follow-ups

* fix: preserve apostrophes in Windows MCP commands

* fix: preserve apostrophe-wrapped Windows MCP args

* fix: align Windows MCP parsing with list2cmdline

* fix: preserve explicit MCP remote transport intent

* fix: trim MCP remote URLs before transport checks

---------

Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
2026-06-11 16:17:22 +01:00
Daniel Han
3733e0b274
fix(studio): surface live step with null loss through the SSE progress stream (#6206)
* fix(studio): surface live step with null loss through the SSE progress stream

The metric histories skip non-finite steps, so during a NaN stretch the
SSE live loop and final complete event replayed the last finite
step/loss pair. Follow the live progress step when it is ahead of the
history tail and report its loss honestly (null until recovery).

Completes the NaN honesty fix for the SSE consumer flagged in review.

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

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

* Apply live-step handling to inactive streams and clear the UI loss on null for PR #6206

Fresh /progress connections after a finished run took the inactive branch
which still replayed the last finite step and loss pair; apply the same
live-step correction there. On the frontend, applyProgress kept the stale
currentLoss when a payload advanced the step with a null loss; clear it so
the display shows -- until the loss recovers. Widen the runtime state type
to number | null, which the view layer already handles.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-11 07:50:13 -07:00
Daniel Han
bc85ecd145
Studio: report the real llama-server context window and add an opt-in overflow policy for OpenAI-compatible serving (#6164)
* Studio: report the real llama-server context window and add an opt-in overflow policy for OpenAI-compatible serving

A community report showed OpenCode failing tool calls every few minutes
against Studio's OpenAI-compatible API while the same GGUF was stable on
LM Studio. Root cause: Studio advertises the requested context length, but
llama-server can allocate less (memory-fit step on small GPUs, --parallel
slot split), so clients budget against a window that does not exist. Their
generations truncate mid tool call at the real wall (finish_reason=length
with cut JSON arguments) and eventually the prompt itself exceeds the real
window, returning a 400 that agentic clients treat as non-retryable.

Changes:
- After llama-server health, read default_generation_settings.n_ctx from
  /props and adopt it whenever it is below Studio's computed context, with
  a warning. The load response, status route, UI value, and the passthrough
  max_tokens ceiling all become honest automatically.
- Expose context_length and max_context_length on /v1/models so clients can
  budget against the enforced window.
- Accept empty role=tool content (commands with no output are routine in
  agentic loops; OpenAI and llama-server both accept it) instead of a 400.
- Add context_overflow=truncate_middle (per request, or server-wide via
  UNSLOTH_CONTEXT_OVERFLOW=truncate_middle): on exceed_context_size_error
  the passthrough drops whole middle turn-groups (system prompt, first turn,
  and recent turns kept; tool calls stay paired with their results), clips
  oversized contents middle-out when group-dropping is not enough, clamps
  max_tokens to the generation headroom, and retries. Default stays 'error'
  with code=context_length_exceeded so clients running their own compaction
  keep full control.

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

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

* Studio: allocate the requested context for real (kv-unified, fit-ctx floor)

Two launch-flag gaps caused the advertised vs allocated divergence at the
source:
- llama-server enables --kv-unified only when the slot count is auto; Studio
  always passes --parallel N, which silently splits -c into per-slot windows
  of -c/N. Pass --kv-unified when N > 1 so a single request can use the full
  advertised window (same total KV memory, shared pool).
- with --fit on the fit step may set ctx as low as 4096; pass
  --fit-ctx <requested> for explicit requests so fit offloads or fails into
  the existing --fit off retry instead of silently shrinking the window.

Both flags are gated on --help capability probing so older builds keep the
current behavior, where the /props readback remains the backstop. Verified
live: -c 98304 --parallel 4 now serves per-slot n_ctx 98304 (was 24576),
48k-token requests pass through the passthrough, and the readback warning no
longer fires.

* [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-11 07:49:55 -07:00
BardiaKoopah
0ac051177a
fix(mlx): forward resume_from_checkpoint to MLXTrainer.train() (#6173)
Studio's frontend exposes a Resume action and submits requests with
resume_from_checkpoint set to a previous run's output_dir. The CUDA
training paths in worker.py read this field from config and pass it to
trainer.train() (see lines 2729-2787 and 3108-3229). The MLX path
_run_mlx_training did neither: it never read config['resume_from_checkpoint']
and called trainer.train() with no args. The MLX trainer also did not
accept the kwarg, so even threading it through would have been a no-op.

With this PR + the unsloth-zoo companion PR adding the trainer-side
support (saves optimizer_state + trainer_state, accepts and applies
resume_from_checkpoint in MLXTrainer.train()), MLX Resume now works
end-to-end. Verified on M2 16GB with Qwen3-0.6B + unsloth/LaTeX_OCR:
loss at every post-resume step matches a fresh run bit for bit
(2.168627977371216 == 2.168627977371216 at step 6, etc).

Two lines: read the field near the other config.get() extractions in
_run_mlx_training, pass it as a kwarg at the trainer.train() call site.

Companion PR: unslothai/unsloth-zoo#751
2026-06-11 07:29:05 -07:00
BardiaKoopah
14f679da8a
fix(studio): surface NaN loss honestly instead of laundering to last finite value (#6016)
When training produced a NaN or Inf loss event, the handler filtered the
value to None but never updated progress.loss — clients kept seeing the
last finite value as if everything were fine.

Now: on non-finite loss, clear progress.loss to None and log a one-shot
warning. Training continues (no phase=error, no _should_stop), matching
the expected behavior for a non-fatal numerical event.

Test: tests/test_training_nan_loss_handling.py with 6 cases covering
finite, NaN, +/-Inf, idempotency of the one-shot warning, and recovery
when a finite step follows a non-finite one.
2026-06-11 07:25:24 -07:00
Daniel Han
73973435ad
Studio: require an installed ROCm DLL before forcing BNB_ROCM_VERSION; drop shadowing shutil imports in save.py (#6194)
* Require a found ROCm DLL before forcing BNB_ROCM_VERSION in Studio paths

main.py previously set BNB_ROCM_VERSION=72 whenever HIP_PATH or ROCM_PATH
was set, and the training worker fell back to a blind 72 when DLL
detection found nothing. On a Windows machine with the AMD HIP SDK
installed but CUDA or CPU torch, that forces a ROCm backend onto a
non-ROCm bitsandbytes wheel, which raises at import. Both paths now only
write the override when a libbitsandbytes_rocm DLL actually exists (or a
seeded value is already present), matching the strict gates in
unsloth/import_fixes.py.

Also removes four redundant local import shutil statements in
unsloth/save.py that shadow the module-level import, the same pattern
that caused the UnboundLocalError fixed in #6149.

* Worker: gate the BNB override on a found ROCm DLL, preserving seeded marker

Review follow-ups: track _found_rocm_bnb in the worker like main.py so a
ROCm DLL with an unparsable name still gets the seeded or 72 fallback,
and skip the env write entirely when no DLL exists so a seeded value
keeps its sitecustomize marker and stays redetectable by later import
fixes.
2026-06-11 06:52:38 -07:00
Daniel Han
a5d6e6928d
Studio: surface the llama.cpp update affordance when MTP is disabled (#6192)
* Studio: surface the llama.cpp update affordance when MTP is disabled

When a model asks for MTP (auto on an MTP model, or forced mtp / mtp+ngram)
but it gets disabled, the load already degrades gracefully and serves without
speculative decoding. Until now the UI gave no hint why, or that an update
would fix it.

Record why MTP was dropped on the backend (spec_fallback_reason): the probe
found no mtp token (binary_no_mtp), the spawn aborted with an outdated-arch /
context-build error such as a prebuilt that predates the Gemma drafter
(binary_outdated), or the current build could not run it, e.g. a CUDA kernel
limit (runtime_error). Expose it in the inference status. In the chat
Speculative Decoding section, show a short note and, for the two update-fixable
reasons, an inline Update llama.cpp button that reuses the existing update flow.
A runtime_error gets the note without an update push, since a newer build may
not fix it.

Backend tests cover the reason being set / cleared. Frontend typechecks.

* Address review: tighten the update hint to genuinely outdated binaries

Reserve binary_outdated (which surfaces the Update llama.cpp affordance) for an
unknown-architecture abort, which proves the prebuilt predates the model;
classify the generic memory/context build failures as runtime_error, where an
update may not help. Frontend: only append the "Update llama.cpp to enable it"
sentence when an update is actually available, so the text never points at an
action the UI is not offering.
2026-06-11 06:10:17 -07:00
Daniel Han
c3604d01f7
Studio: enable MTP for sub-3B Gemma separate-drafter GGUFs (#6191)
* Studio: enable MTP for sub-3B Gemma separate-drafter GGUFs

The sub-3B auto-drop to ngram-mod was tuned for an embedded draft head
(Qwen), whose per-token cost regresses below 3B. Gemma ships the head as a
separate root mtp-*.gguf drafter, a tiny standalone model that is cheap
enough to win below 3B: B200 Q4_K_XL bench, draft-mtp n=2 vs spec-off,
gemma-4-E2B (2B) = 1.21x (accept ~0.65) while ngram-mod is 1.00x.

Exempt a separate drafter from the sub-3B gate everywhere the threshold is
applied: the resolver (_mtp_too_small), the auto-fit VRAM reserve, the
drafter auto-download decision, and the reload-skip mirror via a
has_separate_drafter flag on _auto_mode_drops_mtp. Embedded sub-3B heads
(Qwen) still drop to ngram-mod. A drafter the binary cannot build (older
prebuilt, or a CUDA kernel limit) still aborts the spawn and the load
retries once without speculative decoding.

Adds the full Qwen3.5 + Gemma-4 (regular and QAT) auto/off/forced resolver
matrix, plus explicit sub-3B exemption tests.

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

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

* Always compare the separate drafter in the reload-skip mirror

The sub-3B wrapper around the drafter compare could skip it when the drafter
was deleted out from under a running sub-3B server (detected None, stored set),
leaving a stale launch. The resolved-path compare is cheap and already handles
every case, so drop the _auto_mode_drops_mtp guard (and its now-unused imports)
and always compare when the mode can use a drafter and the user does not own
--spec-type. Addresses review feedback on #6191.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-11 05:27:46 -07:00
Lee Jackson
49024462f3
Studio: ignore unsupported env proxy during Studio startup (#6102)
* fix: ignore unsupported env proxy during Studio startup

* fix: handle missing socksio env proxy at startup

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

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

* Match printf logging style and inline the proxy predicate for PR #6102

* [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: Daniel Han <danielhanchen@gmail.com>
2026-06-11 05:13:27 -07:00