New base_precision config for the DiT trainers: nf4 (unchanged default) |
bf16 | int8 | fp8 | auto, advertised per family + per machine through
/api/train/diffusion/info (precision_modes, recommended_precision,
supports_compile) so the UI can gate the selector.
- bf16: dense transformer + regional torch.compile (auto-armed). The
measured speed mode: 2.3x nf4 on FLUX (1.81 -> 4.12 steps/s), 2.6x on
Z-Image (2.5 -> 6.38 steps/s) on B200, at dense-weight VRAM
(FLUX 24.7 GB / Z-Image 13.6 GB peak vs 10.4 / 4.7 for nf4).
- int8: torchao weight-only int8 on the frozen base, quantized AFTER
add_adapter (quantizing first trips peft 0.18's TorchaoLoraLinear,
which is incompatible with the torchao 0.16 config API). Runs eager:
inductor rejects the int8 subclass training graph (aliased subclass
outputs), so compile is force-disabled for it.
- fp8: torchao convert_to_float8_training on the frozen linears
(filter skips lora_ modules, proj_out, non-divisible-by-16 dims,
pad_inner_dim), applied after add_adapter, compile auto-armed.
Works and round-trips, but measured SLOWER than compiled bf16 at
LoRA-training shapes (FLUX 3.15 vs 4.12 steps/s; Z-Image similar),
so it is an explicit opt-in and auto never picks it.
- auto: free VRAM (measured before load) + dense-size table -> bf16
when it fits with headroom, int8 in the middle band, else nf4.
Prequant bnb repos always resolve to nf4; dense modes on them are
rejected at validation with a pointer to the family's dense base.
Two crashes found and fixed along the way:
- The cuDNN SDPA backend's training graph fails on the FLUX attention
shapes (torch 2.10 + cu130, B200): mha_graph.execute errors, then the
context degrades into illegal memory accesses. The perf-flag guard now
pins flash/mem-efficient SDPA for the run (mathematically equivalent,
snapshot/restored). nf4 escaped it by routing attention differently.
- Regional compile now uses dynamic=True (the inference layer's proven
default): dynamic=False specialisation fused a gemm_and_bias epilogue
that failed with CUBLAS_STATUS_EXECUTION_FAILED on the FLUX training
graph; dynamic=True is also faster (Z-Image 3.84 -> 6.38 steps/s).
Verified: 98 backend tests green (new test_diffusion_base_precision.py:
validation, auto policy table, fp8 filter, compile gating, /info fields);
per-mode 40-step runs on FLUX + Z-Image with loss means inside the nf4
envelope and adapter round-trip generation through the normal LoRA path
for bf16-, fp8-, and int8-trained adapters.
Perf core for the diffusion trainers, defaults preserving the training math:
- Phased model loading: the pipeline now loads without its transformer
(conditioning only), captions are encoded and the text encoders freed,
the VAE latent cache is built and the VAE freed, and only then does the
transformer load. The multi-GB denoiser never shares VRAM with the
encoders, cutting measured peak VRAM on B200: FLUX 17.1 -> 10.4 GB,
Qwen-Image 19.1 -> 12.8 GB, Z-Image 7.3 -> 4.7 GB.
- Latent cache (cache_latents, default on): per-image crop/flip variants
(cache_variants, default 4 vs the single frozen variant of the diffusers
--cache_latents) store the VAE posterior's affine parameters, so every
step still draws a fresh VAE sample; a cached center-crop Z-Image run
matches the uncached one at the bf16 nondeterminism floor.
- True batching: train_batch_size now actually batches the transformer
forward (it was silently 1). nf4 dequant dominates the step cost, so
batch 4 lands near batch-1 step time: 4.0x samples/s on Qwen-Image,
3.1x on FLUX, 2.1x on Z-Image, with multi-seed loss envelopes
overlapping batch-1.
- LR scheduler support in the DiT loop (lr_scheduler / lr_warmup_steps
were accepted but ignored); progress events now report the real
per-step LR.
- TF32 + high fp32 matmul precision under enable_tf32 (default on),
snapshot/restored around the run. cudnn.benchmark is scoped to a
caller opt-in only: autotuning the fp32 VAE convs doubled peak VRAM
on the DiT families for zero steady-state gain.
- Vectorized sigma gathering (drops a per-step Python search loop),
cached FLUX img_ids/guidance, fused torch AdamW fallback, steady-state
samples_per_second (excludes the first-step warmup).
- Regional torch.compile plumbing (compile_transformer off/on/auto with
eager fallback): auto stays off over a bitsandbytes base where compile
is a net loss (27 s warmup, slightly slower steady on Z-Image); it
arms automatically for the dense/quantized speed modes that follow.
- Stop parity with the LLM trainer: /api/train/diffusion/stop accepts an
optional {save} body and the service forwards save=False as a
no-save cancel; a new preparing event surfaces cache-build progress.
- SDXL trainer gets the same latent cache, perf flags, and fused
fallback; its batching, LR schedule, and min-SNR stay as they were.
Verified: 83 backend tests green; per-family 30-40 step runs with
adapter round-trip generation through the normal LoRA path (FLUX,
Qwen-Image, Z-Image all pass).
The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single
compute dtype during the forward, exactly like the diffusers dreambooth
scripts run under accelerator.autocast. Without it the 4-bit backward on
FLUX.1-dev fails with an illegal-address CUBLAS error partway into the
first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL
path (its own trainer) is untouched.
SDXL re-encoded every caption with both CLIP text encoders on every step (pure
waste, since captions are constant) and kept the encoders resident. Precompute
each unique caption's embeddings once, then free the text encoders before the
loop: numerically identical (embeddings are deterministic and this consumes no
torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB
lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32
fallback, halving optimizer state with no meaningful LoRA quality cost. Env
toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard
A/B the paths.
Extends diffusion LoRA training beyond SDXL to the three popular DiT families
via a single shared flow-matching loop parameterised by small per-family specs
(loading, prompt/latent encoding, transformer forward, save). Verified against
diffusers 0.38.0:
- FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly
nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the
user's HF token).
- Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std,
img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16
base).
- Z-Image: list I/O with the reversed timestep convention and a negated
prediction, bf16 only.
The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now
route these families to the DiT trainer; the SDXL blocklist guard is replaced by
a positive family resolution that also rejects GGUF repos (inference-only) and
still-unsupported families. Per-family defaults + labels + VRAM notes are exposed
via family_train_infos for the Train UI.
Memory: caption embeddings are precomputed once and the text encoders freed
before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit)
and 8-bit AdamW are on by default.
The training service kept only the latest loss, so a live loss chart could show a
single point. Fold each progress event into bounded (step, loss, lr) history arrays
(capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM,
and record the family / base model / catalog path on completion. The status endpoint
returns these as a nested metric_history object the UI can chart directly, and the
start request accepts an optional model_family override.
Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:
- New core/training/diffusion_train_common.py holds the config + validation, dataset
discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
from the base model (or an explicit model_family) via the diffusion family registry,
and a known-but-not-yet-trainable family is refused with a clear message. Unknown
custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
and Qwen-Image repos (safetensors-only, no remote code).
The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.
Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.
Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.
Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.
lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
official bases, or a local path) before from_pretrained, so an untrusted remote repo
is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
cancel during download is honoured; a stop may carry save=False to cancel without
leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.
Adds regression tests for the config coercion/validation and aliasing.
start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.
Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).
core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.
routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.
Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.
The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.
core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
.txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
AdamW + get_scheduler + grad accumulation + grad clipping, then export via
save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
so the training worker can spawn it; plus a CLI entry point.
Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.
Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).
Extracted and narrowed from unslothai/unsloth#6543 by @TheJagStudio.
This keeps the startup/banner and text file encoding hardening separate from the already-merged Python code-exec UTF-8 fix in #6548.
Co-authored-by: Jagrat Patel <81472856+TheJagStudio@users.noreply.github.com>
* (feat) Add project names to studio training runs to avoid models being overwritten when doing similar training runs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update studio/frontend/src/features/export/export-page.tsx
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
* Update studio/frontend/src/features/export/export-page.tsx
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
* Update studio/frontend/src/features/export/export-page.tsx
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
* better project name sanitization, removed duplicated project name normalization
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* implement checkpoint scanning utilities and tests for base model inference
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Guard project_name against null and use leading important modifiers
* Fix/adjust training project names for PR #6512
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix/adjust training project names for PR #6512
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address project-name review feedback
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Show project names in training recents
* Keep GGUF export directories source-specific
---------
Co-authored-by: NZ-Linix <nz-linix@outlook.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: NZ-Linix <linus.ordowski@outlook.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
* Studio: keep the training event pump alive so progress can't silently freeze
The parent-side event pump is the only writer of the in-memory progress state
that SSE /progress, /status, /metrics and the DB history all read. It ran in a
single unsupervised daemon thread with no guard around event handling, so one
malformed event or a transient queue/DB error would terminate it permanently.
The worker subprocess keeps training regardless (mp.Queue puts never block on an
unbounded queue), so a run kept burning GPU for hours while every progress
surface froze on the last step the pump saw.
- Guard each pump iteration: a bad event or queue-read error is logged and
skipped instead of ending the loop. _read_queue now reads any error as
"no event", not just Empty/EOFError/OSError/ValueError.
- Add a _pump_running flag and an _ensure_pump_alive watchdog wired into
is_training_active, so a pump that dies while the worker is alive is restarted
on the next status poll and the UI catches up from the still-open queue.
- Start respawned and restarted pumps under the lock so the watchdog can never
spawn a duplicate during the brief start window.
Adds tests/test_training_pump_resilience.py covering both guarantees.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio training pump: address review (drain guard, start race, read backoff, respawn flag)
Follow-up to the event-pump resilience change, closing four edge cases a
review surfaced in the same pump/queue surface:
- _drain_queue now tolerates any error during the worker-exit drain and
finalizes with whatever it drained, instead of skipping finalization and
leaving the run wedged "active" with a dead worker.
- start_training clears a stale _pump_running flag during reset and assigns
the subprocess handles plus starts the pump under the lock, so a concurrent
status/SSE poll can't spawn a duplicate pump during setup.
- _read_queue goes back to the narrow EOFError/OSError/ValueError catch;
truly unexpected errors are left to _pump_loop's guarded read, which logs
and backs off so a persistently raising queue can't spin a hot loop.
- The xet respawn-failure path clears _pump_running so a later run can't
inherit a stale flag.
Adds regression tests for all four.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: revive a crashed pump after worker exit + stop test module pollution
Two review follow-ups on the training event pump:
- _ensure_pump_alive refused to restart once the worker had exited
(not self._proc.is_alive()), so a pump that crashed just before the worker
finished never drained the terminal complete/error events still sitting in
the queue. progress.is_training stayed True and is_training_active() returned
True forever, leaving the run stuck "running" behind a dead pump. A True
_pump_running flag with a dead thread is an unambiguous crash regardless of
worker state, so restart there too: the fresh pump drains the backlog and
finalizes. Updated the watchdog test to assert the revive-and-finalize.
- The resilience test imports core.training.training while heavy module-level
deps are stubbed, then restores the stubs -- but the cached training module
kept the stubs bound in its globals, so a later test in the same session
could exercise the fakes (e.g. prepare_gpu_selection) instead of the real
code. Evict the training module (and its package) after import when this file
created it, so subsequent tests re-import it cleanly.
* Studio: finalize training run when queue reads keep failing on a dead worker
reviewer.py follow-up. _read_queue only swallows EOFError/OSError/ValueError;
an unexpected error escapes to the pump's outer guard, which logged, slept and
`continue`d. If those reads keep raising after the worker has already exited
(e.g. a broken queue pipe), the loop never reaches the dead-worker finalize
block, so the pump spins on with _pump_running True and progress.is_training
stuck True -- the run looks like it is still training forever. On a read failure
now fall through to finalize when the worker is gone, only backing off and
retrying while it is still alive. Mirrors the data-recipe pump fix; added a
regression test.
* Tighten training pump resilience comments and docstrings
Condense the verbose explanatory comments and docstrings on the training event
pump and its tests to shorter, clearer forms. Comment/whitespace only; verified
no code changed via AST diff. No behaviour change.
* Studio: create the training DB run before starting the event pump
start_training started the event pump before the eager _ensure_db_run_created()
call, so for a worker that completes or fails immediately the pump could race the
main thread into creating and finalizing the same run row (duplicate INSERT, or a
finalize skipped while _db_run_created was still false). Create the run first; the
pump then only ever finalizes. Adds a regression test asserting the pump observes
an already-created run.
* [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>
The post-filter safety net for 'Train on completions' fires when
train_on_responses_only() masks every token in too many rows. Its trigger is
a row-drop ratio, not a token-length check, but the message hardcoded
"max_seq_length is too short, try increasing (e.g. 8192)" -- advice that
fires identically at any max_seq_length and can recommend a value below the
user's current setting (telling someone already at 16384 to use 8192).
The dominant real cause is that the model's response template is not found in
the formatted samples: the dataset is already formatted, or its structure
doesn't match the model's chat template, so every token gets masked and the
rows are dropped. Reword the error (and the comment above it) to lead with
that cause and the actionable fix (turn off 'Train on completions'), and
mention max_seq_length only as a secondary possibility without a hardcoded
recommendation.
* Studio: lazy-import matplotlib so the server starts when the wheel is blocked
matplotlib.pyplot was imported at the top of core/training/training.py, on the
server boot path. When matplotlib's native extension fails to load (e.g. an
unsigned wheel blocked by Windows Smart App Control), that import crashed the
whole Studio server at startup instead of just disabling loss plots.
Move it into a lazy _load_pyplot() helper called from _create_loss_plot, using
the headless Agg backend, and return None when matplotlib is unavailable so
plotting degrades gracefully. The plot return was already Optional, so callers
need no changes. Keep the type-only import under TYPE_CHECKING and quote the
annotations.
Fixes#6588
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: pin matplotlib==3.11.0
Pin matplotlib to the current latest so a new unsigned release does not
reintroduce the Smart App Control block on Windows. Belt-and-suspenders on
top of the lazy import. Pinned in both studio.txt and extras.txt.
* Pin matplotlib to 3.10.9 so Studio still installs on Python 3.10
matplotlib 3.11.0 requires Python >=3.11, so the pin had no installable wheel on
Python 3.10 (still supported) and pip install failed there. 3.10.9 is the latest
3.10.x (requires-python >=3.10) and covers Python 3.10 through 3.13. Also tighten
the lazy-import docstrings.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
* Resolve the transformers tier by probing AutoConfig instead of guessing
When the only signal is a 5.x tokenizer class, get_transformers_tier guessed the
lowest 5.x sidecar (530). That misroutes models whose built-in config parser needs
a higher tier: dense NemotronH ships a 5.x tokenizer but its '-' (MLP) layer only
transformers 5.10 can parse, so 5.3/5.5 raise KeyError '-'. The config.json
transformers_version field records the saving version, not the minimum to load, so
it cannot drive routing either.
Replace the weak tokenizer->530 guesses (local and remote) with a probe: parse
config.json with the built-in parser (trust_remote_code=False) in each sidecar,
escalating 530->550->510, and pick the first that succeeds. This generalizes to any
architecture without hardcoded lists. Strong signals stay fast paths (no subprocess);
the probe runs only when the tier is otherwise ambiguous and is cached by (model,
commit sha). It never executes repo code, never downloads weights, never raises, and
falls back to the legacy 530 guess on a transient/auth/offline failure or when no
sidecar is available. UNSLOTH_DISABLE_TIER_PROBE restores the old behavior.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: tier probe fallbacks and cross-platform robustness
Codex:
- Never escalate to 510 on uncertainty. When every sidecar was probed and none
parsed with the built-in parser, the model is a remote-code / custom model_type
that loads via its own code; keep the legacy 530 route instead of jumping to
510 (which would change the behavior of models that worked on the 5.3 stack).
- Only cache the 530 fallback when the result is conclusive (every tier actually
probed). If a sidecar was missing/uninstallable the environment is incomplete,
so return 530 uncached and retry on the next call.
- Do not pin the tier cache under an unknown revision: _resolve_commit_sha no
longer memoizes a None sha (a transient Hub failure is retried), and _probe_tier
only caches a tier when the commit sha is known.
Gemini:
- Wrap Path.exists() in the sha resolver in try/except OSError (a remote repo id
can raise WinError 123 on Windows).
- Probe script writes the error to sys.stderr.buffer as UTF-8 bytes so a non-ASCII
message cannot itself raise UnicodeEncodeError under cp1252.
- subprocess.run decodes stderr with errors="replace" to avoid UnicodeDecodeError
on non-UTF-8 consoles.
Tests: 72 passed (added partial-sidecar uncached, sha-unresolved not cached,
all-failed stays 530 + cached, sha resolver retries None / handles OSError).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review round 2: authenticate tier checks, stop memoizing local sigs
Codex:
- Thread hf_token through _check_config_needs_510/550 and
_check_tokenizer_config_needs_v5 (and the underlying raw fetches). Previously a
gated/private model whose only 5.x signal is tokenizer_config.json never reached
the authenticated probe: the unauthenticated raw fetch failed and cached False,
so the model fell through to the default 4.x tier. The per-check caches are now
keyed by (model, token) so an unauthenticated miss cannot poison a later authed
read, mirroring _load_config_json.
- _resolve_commit_sha no longer memoizes a local directory signature. A local
signature is mutable (size/mtime of config/tokenizer), so a reused/overwritten
checkpoint path would otherwise keep selecting the previous tier; it is now
recomputed every call. Only the immutable remote commit sha is memoized.
Tests: 75 passed (added token-cache isolation + auth header, local signature not
memoized, token threaded into all checks/probe).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review round 3: reach activation with the token, drop SHA tier cache
Codex round 3:
- Thread hf_token into the activation path that actually selects a sidecar. The
token-aware tier checks added last round were unreachable:
activate_transformers_for_subprocess called get_transformers_tier without a
token, and the inference/training/export workers passed only the model name even
though they hold a request-scoped hf_token. activate_transformers_for_subprocess
now takes hf_token and the three workers forward config["hf_token"], so a
gated/private model whose only 5.x signal is an authenticated config/tokenizer is
routed to the right sidecar instead of falling to default 4.x.
- Stop importing huggingface_hub during tier detection. _probe_tier no longer
resolves a commit sha, so it never pulls huggingface_hub into the worker before
the sidecar venv is prepended to sys.path (activation only prepends, never
purges), which would otherwise pin the default-env hub over the sidecar's
pinned huggingface_hub==1.8.0.
- The tier cache is now keyed by model_name for the process lifetime (a model's
required tier is a property of its architecture; cleared on restart). This drops
the mutable-SHA memo that masked remote revision changes and the mutable
local-signature memo, removing _resolve_commit_sha / _local_dir_signature /
_probe_sha_cache entirely.
- Do not cache a probe success that depended on a skipped lower tier: if a lower
sidecar was unavailable, the lowest valid tier may change once it installs, so
the result is returned uncached and re-probed next call.
Tests: 73 passed (probe imports no hub; success uncached when a lower tier is
skipped; activation forwards the token).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments to be more succinct
* Re-probe overwritten local checkpoints and authenticate the probe child
The AutoConfig tier probe cached its result under the bare model_name, so a
local checkpoint overwritten in place (same path, new config.json) kept serving
the stale sidecar. Fold a cheap config.json signature (size + mtime) into the
cache key for local paths; remote ids stay name-keyed so no huggingface_hub
import lands before the sidecar is activated.
The probe relies on the implicit HF_TOKEN env, so an inherited
HF_HUB_DISABLE_IMPLICIT_TOKEN=1 left it unauthenticated and a gated repo 401ed
into the 530 fail-safe. Clear that flag in the child env when a token is set.
* Keep tier probes off the log-only path and probe new 5.x archs default-first
- get_transformers_tier gains probe=True/False. needs_transformers_5 (a coarse
4-vs-5 boolean used only for a spawn log and a vision-check branch) now passes
probe=False, so a parent/log-only caller never spawns sidecar probes. The real
activation path keeps probe=True and resolves the exact tier in the worker.
- A config.json saved by transformers 5.x but matched by no fast path is now probed
default-first: _probe_tier gains include_default + floor, prepending the ambient
4.57.x tier to the escalation. A model that still parses on the default is left on
it (no mis-route onto a sidecar); only a config the default parser cannot read
escalates to the lowest 5.x tier that parses. The transformers_version field is a
cheap 'worth probing' hint only, read from the already-fetched config (no extra
network); ordinary 4.x configs never probe.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Separate probe cache by mode and keep version-field 5.x visible to needs_transformers_5
- _probe_tier cache was keyed only by config.json signature, so a default-first probe
that returned 'default' could be handed back to a later tokenizer/known-5.x caller
(floor=530), leaving a model with a 5.x-only tokenizer on transformers 4.x. Key the
cache by probe mode (floor + include_default); the legacy 530 mode keeps the bare key.
- The version-field 5.x detection is a cheap config read, not a probe, so run it even
when probe=False: a standard-tokenizer model whose only signal is transformers_version
>= 5 now classifies as 5.x via needs_transformers_5 (returns '530' without spawning a
probe), so the vision-routing fallback uses the 5.x subprocess instead of failing the
default parser and marking it non-vision. The real activation path still probes
default-first and may resolve 'default'.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Don't treat local checkpoints as Hub ids, and fix stale activation test double
- _load_config_json / _check_tokenizer_config_needs_v5: a local checkpoint dir whose
config.json / tokenizer_config.json is not yet present was being fetched from the Hub
as if the path were a repo id, and the 404 miss was cached. A later call after the
file is written (in-progress checkpoint) then served the stale miss, so a
TokenizersBackend checkpoint fell through to the default tier. Skip the Hub fetch for
local dirs and do not cache the miss, so the file is read once it appears.
- test_activate_transformers_version_or_warn_*: the worker now threads hf_token into
_activate_transformers_version (model_name, hf_token); update the one-arg test doubles
to the real two-arg signature so the silent-success path stays silent.
* Tighten comments in the AutoConfig probe and tier-selection paths
* Address review: canonical probe cache key and reuse _token_cache_key
- _probe_cache_key resolves config.json to its absolute realpath before
keying, so a relative path or a changed cwd can't collide with or miss a
prior probe result. Remote ids still fall back to the name (stat raises,
caught).
- _cached_config_json reuses _token_cache_key instead of re-hashing the
token inline, keeping the (model, token) key derivation in one place.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Add HF dataset streaming mode to Studio
* Added default value for datasetStreaming in training-config-store.ts
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Handle None max_steps for streaming validation
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: fast-fail streaming validation and guard incompatible modes
Reject dataset_streaming at the API boundary when hf_dataset is empty,
the dataset is vision/audio, or max_steps is not set. Probe eval split
with get_dataset_split_names before the streaming load so typos fail
immediately instead of mid-training. Guard column_names=None after map
on iterables. Hide the UI toggle for non-text configurations and clear
the stale flag when config becomes incompatible.
* studio: add streaming dataset tests, iterable helper, and streaming template/format support (WIP)
Work-in-progress on top of feat/studio-dataset-streaming-mode (PR #4946):
- new test_training_streaming.py and iterable.py dataset helper
- streaming support in chat_templates.py and format_conversion.py
- additional streaming guards in trainer.py / models / routes
- frontend streaming wiring in params-section and training-config-store
Committed to preserve uncommitted work before merging latest main.
* studio: fix review-team findings for streaming + main merge
BLOCKER: streaming + raw-text/CPT crashed on len(IterableDataset). Guard it in the
start route (reject format_type=="raw" or training_type=="Continued Pretraining")
and in isStreamingSupported (datasetFormat !== "raw").
Also:
- models/training.py: validate hf_dataset/subset/split (charset+length, block ..//);
cap dataset slice indices (le=1e9); note validator ordering
- chat_templates.py: guard _apply_custom_mapping .map() for streaming
- trainer.py: warn when packing+streaming
- training-config-store.ts: persist-migration bump to v11 (standalone datasetStreaming
backfill); add isVisionModel to NON_PERSISTED; toast on silent streamingCompatiblePatch
mutations in the 4 indirect setters
- tests: route rejections (max_steps, raw/cpt), slice cap, unsafe hf_dataset
* studio: enable raw-text/CPT dataset streaming + streaming UX polish
- raw_text: keep the lazy filter but skip len()-based row counting for
IterableDatasets so raw-text / CPT can stream; guard the eval-size log
- routes/trainer: drop the raw/CPT streaming block; add a defensive
not-streaming guard on the eval auto-split (train_test_split)
- dataset-section: streaming toggle is visible-but-disabled and lists the
exact unmet requirement(s) in its tooltip; block embedding models
- training-start-overlay: show "streaming (no full download)" instead of a
stuck download bar for streaming runs
- trim the streaming test suite to the high-value cases
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: address streaming review (MLX/embedding guards, sliced eval split, rehydrate timing)
- routes: reject dataset_streaming for embedding training and on Apple Silicon
(MLX); both loaders materialize the full dataset instead of streaming
- trainer: validate the base eval split name so streaming eval accepts HF slice
syntax such as "validation[:1000]"
- training-config-store: defer the onRehydrateStorage setState to a microtask so
it doesn't hit the store's TDZ during synchronous hydration
- test: streaming start rejects embedding models
* studio: harden HF dataset streaming (column_names, split slicing, empty/eval bounds, gating)
Address a deeper streaming review:
- raw_text: resolve_column_names() guards IterableDataset.column_names=None
(from_generator / unresolved features) so raw-text and CPT streaming no longer
raise TypeError before training
- models/routes: reject HF slice syntax in train_split/eval_split when streaming
(load_dataset(streaming=True) raises "Bad split"); reject mixed sources
(local/S3) and embedding/MLX streaming at the API, not just in the UI
- trainer: an empty post-slice/filter stream fails preflight with a clear message;
streaming eval is capped (STREAMING_EVAL_MAX_SAMPLES) so each eval terminates;
the manual-slice shortcut falls back to a regular load when train_split is sliced
- format_conversion: streaming conversions preflight the first mapped row so
format errors surface before training, not mid-iteration
- frontend: block streaming on Apple Silicon; clear datasetStreaming when a
dataset is detected as image/audio at start
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: fix CI for streaming PR (lint blocker + no-torch sandbox + preflight test)
- trainer.py: drop unused `IterableDataset` import (hoist safety-net blocker).
- test_training_streaming.py: only select real classes (isinstance type) when
locating the trainer class, so a MagicMock-stubbed global is never passed to
object.__new__ (fixes TypeError on the Python 3.10-3.13 jobs).
- no-torch import sandboxes (test_e2e_no_torch_sandbox.py,
test_studio_import_no_torch.py): teach the chat_templates/format_conversion
exec stubs and the full-import-chain copy list about the new `.iterable`
module so the AFTER/runtime cases import without torch again.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
* Studio: persistent per-user trust_remote_code approval cache
The consent gate pins each approval to a content fingerprint (sha256 over every
repo .py), but nothing was persisted, so the dialog reappeared on every fresh
load of the same unchanged repo. This adds an on-disk, per-user approval cache
that lets the gate skip the dialog when the same user reloads the same code,
while keeping the safety guarantees intact.
Two-tier validation, both must hold or the user is re-prompted:
- Commit SHA (cheap, one HfApi.model_info().sha, no download): a match means a
byte-identical tree to the approved revision, so the scan/download is skipped.
- Content fingerprint (authoritative): used whenever the SHA is unavailable
(local path / offline) and always recomputed on a SHA miss. A new or edited
.py changes both the SHA and the fingerprint, so it is caught in every mode.
Safety:
- Keyed per subject; one user's approval never auto-runs code for another.
- CRITICAL is never stored or honored (guarded on both write and read), so a
hand-edited store cannot smuggle in an auto-approval.
- The malware (HF unsafe-file) gate stays unconditional.
- Fail-safe: a corrupt store, an unresolvable SHA, or any error degrades to
"ask again", never to "auto-approve". UNSLOTH_TRC_APPROVAL_CACHE_DISABLE=1
turns the cache off entirely.
New module utils/security/remote_code_approvals.py holds the store
(studio_root()/security/remote_code_approvals.json, atomic write, 0600, RLock)
plus the SHA resolvers. Recording happens at the single gate chokepoint when the
caller supplies the matching fingerprint, so subject is just threaded through
inference/training/export (orchestrators, routes, workers). The scan endpoint
returns already_approved so the frontend can skip the dialog on a cache hit.
Tests: new tests/test_trc_approval_cache.py covers cache miss, SHA-match skip,
SHA-moved re-scan, new-file re-consent, CRITICAL never cached (write + forged
read), disable flag, subject isolation, combined adapter+base key, corrupt
store, and no-subject bypass. Full security suite: 101 passed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: make the approval cache skip only the prompt, never the scan
Codex found that the SHA "no-scan" fast path could run untrusted code without
re-consent. Removed it; the gate now always re-scans and the cache only seeds the
authoritative fingerprint check, so it can skip the dialog but never the scan.
- CRITICAL is hard-blocked on every load (the scan always runs), so a hand-edited
store that downgrades a CRITICAL repo's severity can no longer auto-run it
(P2: do not trust editable severity for SHA approvals).
- The fingerprint covers external auto_map repos, so changed third-party code
always re-prompts even when the primary commit SHA is unchanged; there is no
longer a SHA path that bypasses the fingerprint (P1: external auto_map repos).
- resolve_commit_sha is resolved fresh on every call (no memoization), so a repo
whose default branch moves after approval re-prompts instead of reusing a stale
cached SHA (P1: revalidate mutable Hub SHAs). The SHA is now only a conservative
secondary gate: a fresh resolvable SHA must match the approved revision, else the
seed is withheld; a None (local/offline) falls back to the fingerprint.
- Approvals record the scanner ruleset version (SCAN_RULES_VERSION); the gate
ignores approvals from an older ruleset so reclassified bytes are re-scanned and
re-shown instead of silently auto-approved (P2: invalidate on scan-policy change).
Tests: test_trc_approval_cache.py rewritten around the prompt-skip semantics
(unchanged repo still scans; SHA move / changed code / scanner-version bump /
disable flag all re-prompt; forged downgraded severity still blocks CRITICAL).
105 passed with test_consent_gate.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments to be more succinct
* Keep run-owner subject out of persisted config; serialize approval writes
Threading subject (the run owner's username / API-key id) into the training
config meant _sanitize_db_config persisted it into config_json, which
training-history GET returns to any authenticated user, leaking who started a run
in multi-user installs. Filter subject alongside the token fields; the worker
still receives it from the live config.
The approval store's RLock only guards one process, but approvals are recorded
from separate inference/export/training subprocesses, so concurrent writers could
clobber each other on os.replace and drop an approval (re-prompt). Hold a
best-effort cross-process file lock around the read-modify-write.
* Fail safe on a malformed approval store
A store with the right version but a non-dict shape (e.g. a hand-edited
"subjects": []) passed _load()'s check, then lookup chained .get() on a list and
raised, breaking every remote-code load until the file was removed. Validate that
subjects is a dict in _load(), and tolerate a non-dict per-subject entry in
lookup/record/forget, so a corrupt store fails safe (re-prompt) instead.
* Keep subject out of the MLX W&B run config
_run_mlx_training uploads the whole training config to W&B minus a sensitive set
that only listed hf_token/wandb_token/s3_config, so the authenticated subject
(username / API-key id) was sent to W&B as run config even though DB history
already strips it. Add subject to the W&B-sensitive filter, mirroring
training._sanitize_db_config.
* Tighten the W&B subject-filter comment
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Studio: self-heal unsloth namespace-package shadows in all subprocess workers
A directory named `unsloth` (or `unsloth_zoo`) without an __init__.py on
PYTHONPATH/sys.path, a stray source checkout or a polluted PYTHONPATH, makes
`import unsloth` resolve to an empty namespace package, so a worker's
`from unsloth import FastLanguageModel` dies with a cryptic
"cannot import name ... (unknown location)".
The LLM training path already recovered from this via `_ensure_real_packages`
in trainer.py (PR #6269), but the inference, export, and embedding-training
subprocesses imported Unsloth directly with no guard. Extract that helper into
a shared, dependency-free core/import_guards.py and call it before the Unsloth
import in every subprocess: it drops the offending sys.path entries, imports
the real packages (unsloth before unsloth_zoo so the pre-zoo GPU fixes run),
then restores sys.path. trainer.py now imports the shared helper instead of its
local copy.
Covers both unsloth and unsloth_zoo and both namespace origin forms (None and
"namespace"). The existing PR #6269 test now exercises the shared helper.
* Studio: distinguish a failed model load from no model in the attach gates
A failed load never sets the checkpoint, so the image and audio attach gates
fell through to "Load a model before adding images/audio", which reads as if
the user simply forgot to pick a model rather than that the load errored. Add a
dedicated lastModelLoadError to the chat runtime store, set only when an actual
load attempt fails (not on refresh, list, status, or unload errors, which keep
using modelsError) and cleared when the next load starts. The image gate (all
three call sites) and the audio gate now use it to report a failed load and
point at the server logs, while still blocking in exactly the same cases.
* Tighten namespace-shadow guard and load-error comments
* Studio: free chat model VRAM at training start only when the GPU is tight
The training start route unconditionally tore down the transformers/MLX
inference subprocess before training, and never stopped the llama.cpp GGUF
server at all, so a loaded GGUF chat model kept holding VRAM for the whole
run. Conversely the HF model was always unloaded even when there was plenty
of room to keep it.
Make the unload VRAM aware and cover every inference backend:
- Add routes/training_vram.py with summarize_resident_chat(),
can_keep_chat_during_training() and free_chat_models_for_training(). The
keep/unload decision reuses the same estimator and live per device free
VRAM reader the training GPU selection already uses (auto_select_gpu_ids,
estimate_required_model_memory_gb, get_visible_gpu_utilization), so the
probe agrees with the placement computed later in start_training.
- When a chat model is resident and training fits alongside it with a
conservative margin (required_gb * 1.15 + 4 GB), keep it loaded so the
user can train and chat at the same time; on a multi GPU box training
lands on a different GPU and both coexist. Otherwise unload the HF/MLX
orchestrator and the llama.cpp GGUF server before training starts.
- The export subprocess shutdown stays unconditional and now runs first so
its freed VRAM is reflected in the decision.
Default deny: non CUDA backends, unestimable models, or any probe error
fall back to the previous always unload behavior.
Adds tests/test_training_vram_coexistence.py and updates two existing route
tests in test_gpu_selection.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: per-GPU floor for explicit GPU lists + don't unload chat on invalid gpu_ids
Address review feedback on the chat coexistence probe:
- Explicit gpu_ids mode now enforces a per-GPU floor in addition to the
aggregate free-VRAM check, mirroring auto_select_gpu_ids' min_per_gpu_N.
Without it, an uneven split such as free [45, 10] for a 40 GB job passed
the aggregate threshold and kept chat loaded even though the 10 GB GPU
could not hold its training shard, risking an OOM.
- Invalid explicit gpu_ids (ids outside the visible set, or a UUID/MIG
mask) make resolve_requested_gpu_ids raise. That request is rejected with
a 400 before training starts, so leave the resident chat model untouched
instead of unloading it.
- Tighten the target_modules / gpu_ids type hints to List[str] / List[int].
Adds tests for the per-GPU floor (uneven split unloads, even split keeps)
and for invalid gpu_ids keeping the chat model loaded.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: only free chat VRAM once training will start; handle in-flight and CPU-only chat
Address the second review pass on the chat-coexistence path:
- Run the chat/export VRAM teardown as a before_spawn hook inside
TrainingBackend.start_training, fired only after the start guards pass.
Previously the route freed chat VRAM before calling start_training, so a
refused start (e.g. a lingering pump thread) would tear down the resident
chat model even though no training job began.
- Treat an in-flight HF chat load (loading_models set, no active model yet)
as not safely sizeable: free it rather than risk both OOMing as the load
keeps allocating after training starts.
- Do not count or tear down a GGUF llama-server confirmed to run entirely on
CPU (_gpu_offload_active is False): it holds no VRAM, so killing it cannot
help training fit.
Adds tests for the before_spawn hook (runs on start, skipped when a
subprocess is alive or a pump thread will not die, survives a hook error),
the in-flight load flag, and the CPU-only GGUF exclusion in both the resident
summary and the unload path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: treat any in-flight chat load (HF swap / mid-start GGUF) as unsafe to keep
Tighten the in-flight detection in summarize_resident_chat so the keep check
never sizes a load that is still allocating:
- Flag loading on ANY non-empty loading_models, not only when active_model_name
is empty. load_model adds the new model to loading_models before clearing the
old active_model_name, so a replacement load during a swap was previously
sized as a normal resident and could OOM as the new model finishes loading.
- Flag a GGUF server that is active but not yet healthy (is_loaded False) as
in-flight: it is still mmaping/offloading layers, so its final VRAM footprint
is unknown.
Consolidates the signal into a single resident["loading"] flag; the route frees
the chat model whenever it is set. Adds tests for the replacement HF load and
the mid-start GGUF cases.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: tighten comments in chat/training VRAM coexistence (comments only)
* Studio: run before_spawn VRAM hook only after GPU-selection validation
Reviewers found the before_spawn hook fired before prepare_gpu_selection
validated gpu_ids (and before config build), so a refused start (invalid
gpu_ids -> 400, or a bad grad-clip value) could still tear down chat/export
VRAM. Move the hook to immediately before proc.start(), once all synchronous
validation and process construction have passed. This also fixes the route's
in-flight-chat loading branch, since that teardown runs inside the same hook.
Add test_hook_skipped_when_gpu_selection_rejects.
* Studio: recompute GPU auto-selection after the before_spawn VRAM hook
Codex P2: with before_spawn moved after prepare_gpu_selection, placement was
frozen against the pre-teardown VRAM state while the hook freed export/chat
afterward. Auto-selection could pin training onto a GPU the hook then cleared
(or onto a kept chat model). Split validation from placement: explicit gpu_ids
are still validated before the hook (raise -> 400, no teardown; explicit
placement is VRAM-independent), but VRAM-dependent auto-selection now runs
after the hook so it sees the freed memory.
Add test_auto_placement_runs_after_hook and test_explicit_placement_validated_before_hook.
* Studio: allow chatting during training (lift sidebar gate + VRAM-aware load guard) (#6335)
* Studio: allow chatting during training (lift sidebar gate + VRAM-aware load guard)
The sidebar disabled New Chat, project, and home navigation while a training
run was active, so users could not chat during training even though the backend
serves inference fine alongside a run. This removes that gate and adds a backend
guard so the one genuinely risky operation, loading a new local chat model
mid-training, is refused with a clear 409 when it would not fit beside the run.
Frontend (app-sidebar.tsx): drop the chatDisabled = isTrainingRunning gate and
its consumers. Navigation triggers no model load on its own, so chat stays
usable during training.
Backend (routes/training_vram.py, routes/inference.py): add
can_load_chat_during_training plus a load/validate guard that sizes the same
effective load the loader performs (LoRA 4-bit to 16-bit resolved first, HF auto
placement via auto_select_gpu_ids, explicit multi-GPU per-GPU floor, GGUF sized
from on-disk shards and companions or the selected remote variant). It is a
no-op when training is inactive, never blocks external providers or
already-resident models, and default-denies only on a CUDA sizing failure so a
load can never OOM the run. Validate refuses early with the real settings so the
frontend does not unload the resident chat model for a load that would be
rejected.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: address review feedback for chat-during-training load guard
- Run the load/validate VRAM guard via asyncio.to_thread so the sync
nvidia-smi + HF metadata work never blocks the event loop.
- Size the GGUF KV cache at the requested context (_estimate_gguf_kv_gb)
and add it to the local GGUF estimate so large-context picks are not
under-counted.
- Keep the requested quantization when adapter_config.json is malformed
(not a JSON object) instead of raising in _effective_load_in_4bit.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: size the training load guard at the launcher's effective GGUF context
The GGUF KV-cache estimate used max_seq_length only, but the llama.cpp
launcher honors a user --ctx-size/-c in llama_extra_args. A load such as
max_seq_length=4096 with --ctx-size 131072 was sized against a 4k cache
while the server allocates 131k, so the guard could approve a long-context
GGUF load that then OOMs training. Size the guard's KV at the larger of
max_seq_length and the parsed --ctx-size (reusing the launcher's own
parse_ctx_override), keeping the conservative f16 cache so the estimate is
never smaller than what the server allocates.
The chat model picker also validated with the raw max_seq_length while
/load sizes with resolveLoadMaxSeqLength, so validate could pass, unload
the current model, then have /load reject the native-context load. Validate
now uses the same effective context; the load path is unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: size the GGUF training guard at the server parallel-slot count
The KV-cache estimate assumed a single slot, but llama-server allocates the
cache across --parallel slots (app.state.llama_parallel_slots). On a Studio
launched with --parallel N>1 the guard under-sized the cache N-fold and could
approve a GGUF chat load that then OOMs training. Thread the same slot count
the loader uses into the guard's KV estimate; default 1 leaves single-slot
setups unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments for chat-during-training guard
* Studio: keep chat generation alive across navigation; Train spinner + Return to Chat
Hoist the base chat runtime above the routed outlet so navigating to Train (or any tab) no longer aborts an in-flight generation; only an explicit Stop cancels. Add a Train sidebar spinner and swap New Chat to Return to Chat while a run is active, with a lightweight completion watch so the spinner clears from any tab. Also respawn a chat llama-server killed mid-session and guard unreadable HF cache dirs that 500'd the hub model list.
* Studio: show Return to Chat on the Train tab whenever a chat is live
Previously the top sidebar item only swapped to Return to Chat while training was running; on the Train tab with an idle/just-finished run it stayed New Chat, which started a fresh thread and cancelled an in-flight generation. Show Return to Chat (and navigate back, preserving the run) whenever a generation is running or its thread is still active, or training is in progress.
* Studio: keep a running chat alive when starting a New Chat
Starting a New Chat (or switching threads) while a generation was in flight
remounted the single-chat runtime provider, which detached the in-flight run
and cut the previous chat off (it showed up frozen / empty when reopened).
Key the single-chat view by project instead of by thread or new-chat nonce so
the provider stays mounted and assistant-ui switches to a fresh thread in place.
The previous generation keeps streaming in the background and autosaves on
completion, and returning to that thread reattaches the live run instead of
reloading a half-saved one.
Also:
- "Return to Chat" now lands on the thread that is still generating rather than
the empty new chat that became active after New Chat.
- Skip the explicit /inference/cancel POST when an abort comes from a runtime
detach (navigation / background switch) rather than an explicit Stop, so a
backgrounded generation is never cancelled behind the scenes.
* Studio: make model export non-blocking and inline
The Export tab opened a full-screen modal that trapped focus, could not be
closed or cancelled while running, and showed no progress. It also stopped
training and unloaded the chat model before loading, so export could not run
alongside them.
Export now mirrors the training runtime pattern:
- Inline panel embedded where the Export Model button was, with no modal or
backdrop, so the rest of the UI stays usable during an export.
- Global export runtime store plus an app-root lifecycle hook, so a run keeps
going and streaming across navigation and is reflected on the Export nav item
from any tab.
- The worker log stream now stays connected across the load to export phase
boundary instead of stranding on "Waiting for worker output".
- Progress bar driven by phase and quant index (quant N of M for GGUF), with
elapsed time and a working Cancel.
- load-checkpoint no longer stops training or unloads inference; export loads in
its own subprocess in parallel and surfaces out-of-memory as a clear error.
- Add POST /api/export/cancel and is_export_active on /api/export/status.
* Studio: show Return to Chat on the Export tab too
Extend the New Chat to Return to Chat swap to the Export route so leaving a
running chat for Export offers a way back to the live generation, matching the
Train tab.
* Studio: smooth out Export animations and polish the panel
- Drop the height-based reveal animations (source switch, run panel, quant
picker, hub fields) that caused flashing and reflow; use instant swaps and
quick opacity fades instead.
- Method and quant cards now transition colors only, with no transition-all or
hover lift, so selecting a method or quant is crisp instead of jumpy.
- Auto-scroll the export panel into view when it opens and add a scroll-to-bottom
button when its output is below the fold, like Chat.
- Show Return to Chat on the Export tab while an export is running, matching how
training drives it on the Train tab.
- Surface the current phase or stage in the live output before the first worker
line arrives so the panel never looks stuck while progress is advancing.
* Studio: show Return to Chat on every non-chat tab
Generalize the Return to Chat swap from just Train/Export to any non-chat route
(Recipes, Projects, Hub, ...) so a running or active chat is always one click
away, instead of showing New Chat there.
* Studio: stream export logs over the Cloudflare tunnel; drop janky export animations
Exporting over a --secure Cloudflare quick tunnel showed "connecting..." with no
logs while the progress bar advanced. Cloudflare buffers text/event-stream and
only flushes when the stream closes, so the SSE log stream never reached the
browser during the run (direct localhost is unaffected, which is why this only
showed up over the tunnel).
Add a tunnel-safe JSON poll fallback (GET /api/export/logs?since=) that the
runtime lifecycle hook polls while a run is active. Short JSON responses are not
buffered by the proxy, so logs show up in near real time over the tunnel. It
shares the orchestrator's monotonic seq cursor with the SSE stream and the store
de-dupes by seq, so the two transports run together (SSE on localhost, poll over
the tunnel) without double-printing. A successful poll marks the panel
"streaming" instead of leaving it stuck on "connecting...".
Also remove the framer-motion AnimatePresence reveals from the export config and
run panel (quant picker, hub fields, the inline run panel, and the live log
section). The expand/slide animations flashed and felt clunky; the sections now
render in place.
* Studio: recover export over the Cloudflare tunnel when the blocking POST times out (524)
A model export over a --secure Cloudflare quick tunnel showed "Request failed
(524)" even though the export succeeded on the backend (the GGUF was written).
Cloudflare returns 524 when a single request takes longer than ~100s to respond,
and a GGUF conversion routinely runs for minutes, so the blocking per-method
export POST is cut off while the backend keeps going.
Confirm completion via short status polls instead of relying on the long POST
response (the same approach that fixed log streaming):
- The orchestrator records each finished op's outcome (status / output_path /
error) with a monotonic seq, exposed on GET /api/export/status.
- parseJson now preserves the HTTP status; a 524/520/522/523/502/503 or a
status-less network drop is classified as a recoverable transport error.
- runExport wraps each phase (load, every export method, each GGUF quant): on a
recoverable failure it keeps the run alive (logs keep streaming, the panel
shows "reconnecting...") and polls status until the still-running op finishes,
then settles from the recorded result, recovering the output path for the
success banner. A real 4xx still fails immediately; localhost still uses the
fast POST response. applyBackendStatus also settles a reloaded run from the
last-op record.
Verified over the tunnel: a 3m14s gemma-4-E4B-it GGUF export now ends on the
success banner with the output path instead of 524.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: keep the export method + logs visible after navigating away mid-export
While an export was running, navigating to another tab and back to Export
remounted the page and reset the local form state (exportMethod, quant levels),
so the method card showed unselected and the run panel's log area was hidden
until the card was re-clicked. The run itself lives in the global store and was
unaffected.
Seed exportMethod / quantLevels from the active run's summary via lazy useState
initializers on (re)mount, and gate the panel's log area on the live run
(isExporting / logLines / the run's method) rather than only the local form
selection. The card stays selected and the logs/progress stay visible across
navigation; nothing changes when no run is active.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Studio: address export/training review findings
- Export: guard Start against an empty GGUF quant selection so an inline-panel
run with no quant can't settle as success with no file produced.
- Export: thread the source HF token into the background load so gated/private
HF source exports (and gated bases) authenticate, matching the consent path.
- Export: only settle a recovered (non-owned) run as a finished export when the
last backend op was an export, not a standalone load_checkpoint.
- Training: free the export subprocess whenever an export is active, not only
once a checkpoint is loaded, so an in-flight export load can't race training
for VRAM (current_checkpoint is unset during the load phase).
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Studio: fail fast on an invalid first training batch
Training a base vision-language model (e.g. Qwen/Qwen2-VL-7B or
unsloth/Qwen2-VL-7B) on a conversational image dataset crashed on the first
step with 'Expected ... Long, Int; but got torch.cuda.FloatTensor (embedding)'.
Root cause: the base model's chat template is a flat, media-only template that
renders to an empty string for role-based messages, so UnslothVisionDataCollator
hands the processor empty text, the processor returns empty input_ids, torch
defaults the empty tensor to float32, and the embedding lookup rejects it.
Add a preflight that runs one real batch through the trainer's own tokenization
and collation right before train(), and stops the run with an actionable message
when input_ids is empty or non-integer (pointing to the instruction-tuned variant
for the base-model case). Faithful across text, vision and audio-VLM paths, and
never blocks a run whose first batch is valid.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments in the training preflight
* Stub unsloth/trl in preflight test so backend CI collection passes
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Reap Studio child processes when the parent dies abnormally
Standalone `unsloth studio` launches orphaned cloudflared and llama-server when
the parent exited without running the cooperative shutdown path (terminal-window
close, Task Manager End Task, SIGKILL): the children reparented to init and kept
running, leaving an authenticated Cloudflare tunnel up for days.
Add utils/process_lifetime.py: a parent-owned Windows Job Object
(JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, children auto-inherit) plus Linux
PR_SET_PDEATHSIG, behind a best-effort helper that mirrors the desktop app's
windows_job.rs. initialize_parent_lifetime() runs at the top of run_server;
long-lived spawns (cloudflared, llama-server, RAG embedder, llama.cpp updater)
get the PDEATHSIG preexec, multiprocessing workers are adopted into the job, and
_graceful_shutdown plus atexit gain a terminate_all() backstop sweep. The
cooperative shutdown path is otherwise unchanged.
Verified on Linux: killing the parent now reaps cloudflared and llama-server
within ~2s instead of orphaning them.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* test: add real Windows kill-on-job-close integration test
Spawn a parent that installs the job and a child that inherits it, terminate
the parent, and assert the child is reaped. Skipped off Windows. Also make the
liveness probe Windows-safe (os.kill(pid, 0) terminates on Windows).
* Fix Win64 handle truncation in the Job Object calls
Set explicit argtypes so the 64-bit job/process handles are not marshaled as
c_int (which truncated them on Win64, failing AssignProcessToJobObject). Assert
install success in the Windows integration test.
* Bind multiprocessing workers to parent death; harden the sweep
Review follow-ups:
- Multiprocessing workers (inference/export/training/data-recipe/Xet) cannot be
given a preexec_fn by the parent, so adopt_pid alone left them orphanable on a
Linux SIGKILL. They now bind themselves with PR_SET_PDEATHSIG at startup via
bind_current_process_to_parent_lifetime(), wired into the shared
run_without_native_path_secret entrypoint and the Xet child entry.
- Wire the previously-missed data-recipe worker through adopt_pid.
- terminate_all now honors its timeout: SIGTERM, wait, then SIGKILL the
survivors, so cooperative children can exit cleanly.
- Track adopted pids with a /proc starttime identity and add forget_pid, so the
shutdown sweep never signals a recycled pid.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: Michael Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Harden model fetching: consent gate for trust_remote_code
Add a load-path consent gate that scans a model's auto_map repository code
before it executes and blocks CRITICAL/HIGH findings unless the user pins
approval of that exact code version. Capability detection stays code-free,
reading raw config.json instead of AutoConfig.
- Scan config.json and tokenizer_config.json auto_map, nested local helpers,
and external owner/name--module repos; fail closed on partial downloads.
- Gate inference, training, and export workers, including the MLX path and a
LoRA's base model, and report requires_trust_remote_code from the raw config
so chat and auto-load surface the dialog.
- Verify trusted-org auto-enable against the Hub with the request token and key
the verdict cache by token; reject local-path and spoofed names.
- Add a consent dialog showing the flagged file, line, and surrounding code.
- Thread hf_token through the scan and load paths for gated repos.
* Address review: token handling, tokenizer/LoRA scan coverage, rollback
- Send the HF token for remote-code scans in the POST body, not the URL, so it
never lands in a log or browser history.
- Collect tokenizer_config.json auto_map files directly instead of relying only
on the repo file listing.
- Resolve a LoRA's base model for the validate flag and the scan endpoint so the
dialog scans the code the workers actually gate.
- Pass the request token to the training YAML trusted-org auto-enable.
- Resend a previously approved fingerprint when rolling back to a custom-code
model after a failed switch.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Consent UX: drop legacy chat toggle, fix decline copy, purge declined downloads
The per-model consent dialog is now the single approval path for custom
(auto_map) code in chat, so three leftovers from before it existed are removed:
- Remove the "Enable custom code" switch from Chat Settings and stop persisting
trust_remote_code, so a previously saved blanket-on cannot linger and load a
model without going through per-version review. The flag stays as an internal
YAML/preset default (e.g. first-party auto-enable); the load path still gates
every custom-code load on a fingerprint only the dialog produces.
- Reword the decline message and the auto-load toast to describe approving the
model's code from the dialog, not a missing settings toggle.
- On decline, purge the repo the scan downloaded so untrusted code is not left
on disk. A new /api/models/discard-remote-code endpoint deletes only a
metadata-only cache entry the scan created; it refuses local paths, loaded
models, and any repo with weight files cached, so a model the user already had
or pre-downloaded is always left untouched. The frontend only calls it when
the scan reported created_by_scan.
Adds discard-endpoint tests (delete metadata-only, refuse on weights/gguf,
refuse local, no-op when not cached) and a created_by_scan payload assertion.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Export: remove the user-facing trust remote code toggle
The Export page kept a "Trust remote code" switch (default on) next to the HF
token field. Like chat, custom (auto_map) code should be approved per model
through the load-time review dialog, not a persistent blanket switch, so the
toggle is removed. The export load path already routes through the same consent
dialog: an HF source now starts with trust_remote_code off and only enables it
when the user approves the scanned code in the dialog (a local checkpoint the
user exported stays trusted by default). With the dialog unreachable and no
approval, an HF source loads with trust_remote_code off, which fails closed
rather than running unreviewed code.
* Block loads of repos with unsafe files using Hugging Face's security scan
The trust_remote_code consent gate covers one load-time RCE vector (a repo's
auto_map Python). It does not cover the other: a malicious pickle inside a weight
file (pytorch_model.bin, *.pkl, *.dat) deserializes during from_pretrained even
with trust_remote_code False, so a repo with a normal config plus a poisoned
pickle slips past the existing gate.
Add a metadata-only malware gate that uses Hugging Face's own scan (picklescan +
ClamAV), read via model_info(securityStatus=True).security_repo_status. It never
downloads, opens, or unpickles the flagged files; it only reads the Hub's verdict
and surfaces the flagged file names. New evaluate_file_security runs
unconditionally (independent of trust_remote_code) in every load path (inference,
training SFT/MLX, export), blocking the load when a file is flagged
unsafe/suspicious/malicious. The /remote-code-scan preflight and the validate
endpoint also report the result so the consent dialog opens as a hard block (no
override) listing the flagged files, even for a repo with no custom code.
Policy: hard block with no user override; fail open when the scan is unavailable
(offline/unscanned) so legitimate loads are not broken; no first-party exemption
(a poisoned pickle in a compromised trusted repo still blocks); local paths and
GGUF are skipped (no Hub scan, non-pickle format). Blocking does not gate on
scansDone, since that is often false for clean repos and a file already flagged
unsafe is unsafe regardless.
Adds test_file_security.py covering the block/allow/fail-open/skip matrix.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: scan list-form tokenizer auto_map, gate unsafe files on all load paths
Fixes from a 10-reviewer pass on the model-fetching hardening:
- The remote-code scanner skipped tokenizer auto_map encoded as a [slow, fast]
list (transformers' standard tokenizer shape, e.g.
{"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}). External
tokenizer code in that form was never fetched, scanned, or fingerprinted, so an
AutoTokenizer(trust_remote_code=True) load could run it. _auto_map_refs now
flattens string, list, and nested values. Adds a regression test.
- Compare-mode chat loads and background auto-load only gated on
requires_trust_remote_code, so a repo flagged unsafe by the Hub scan but with no
custom code skipped the hard-block dialog. Both now also gate on
requires_security_review, matching the main chat path.
- The /remote-code-scan and /validate routes collapsed a LoRA adapter to its base
before the malware scan, so unsafe files in the adapter repo itself were missed
in the pre-load review (the workers already scan both). Both routes now run the
file-security scan over the adapter and the base.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Require approval for all HIGH remote code, fail closed when unscannable
Tighten the load-time security gates based on review:
Consent gate
- HIGH-severity auto_map code now requires explicit, per-version approval for
every repo, including first-party unsloth/nvidia. The org is no longer a
blanket bypass: a compromised first-party repo with HIGH code still warrants
review. CRITICAL stays a hard block; clean code still loads after the consent
prompt.
- Fail closed when auto_map code is present but cannot be fully fetched or
listed to scan (gated, offline, transient, or a repo-listing failure that
could hide an imported helper). We cannot fingerprint code we cannot see, so
this is a non-approvable block, retryable once the repo is reachable.
- Scan auto_map from every config that can carry one (model, tokenizer, image
and feature processor, processor, video processor), not just config.json and
tokenizer_config.json, so a custom-processor model is not missed. The file
list is the single source of truth in remote_code_scan and is pinned to the
transformers filename constants by a guard test.
- Distinguish a genuine 404 (config truly absent) from a transient error: only
the latter forces a scan, so a repo with no config is correctly a no-op.
Malware gate
- Scan a remote repo even when its name ends in .gguf; only local paths skip the
Hub scan, so a repo cannot dodge the scan by naming itself "*.gguf".
- Correct the docstring: a file already flagged unsafe blocks regardless of
scansDone; the only fail-open path is an unavailable scan.
Coverage
- Resolve a remote LoRA adapter's base model (not just local directories) so the
base, where the code and weights actually execute, is scanned in validate,
the scan route, and the training and export workers.
- Gate the embedding training path (FastSentenceTransformer) with the malware
and consent checks, matching the other load paths.
Tests updated and added for each change.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scope malware gate to the load-path vector; stop false-blocking first-party models
Follow-up hardening from a second review pass + a broad live model matrix
(unsloth/* , nvidia/* , third-party, and the eicar malware repo).
Malware / unsafe-file gate
- Scope the block to the actual RCE vector: a root-level file in a code-executing
format. from_pretrained deserializes weight files at the repo ROOT, so a flag is
only a load-path pickle vector there. Two exclusions, because neither is loaded:
inert formats (safetensors is tensor-only, gguf is non-pickle, configs/text/
images) and files in subdirectories. This keeps eicar blocked (its *.pkl/*.dat/
eicar_test_file sit at the repo root) while no longer false-blocking legitimate
first-party repos: nvidia/Nemotron-H-8B-Base-8K ships root safetensors plus NeMo
pickle checkpoints under nemo/ that the loader never touches, and the Hub flags
both; the gate previously hard-blocked it.
- Unknown / future non-"safe" levels now fail closed (block) instead of being
silently allowed, so Hub schema drift cannot introduce a bypass; in-progress
("pending"/"scanning"/"error") levels stay non-blocking to avoid false blocks.
Consent gate
- Ignore a STALE own-repo auto_map target that is absent from the repo listing (an
older config pointing at a file the repo no longer ships) instead of failing the
whole repo closed as unscannable. The present .py are still fully scanned, which
is the stronger coverage, and a file that is not there cannot execute. This
unblocks first-party models like unsloth/PaddleOCR-VL (its tokenizer_config.json
names processing_ppocrvl.py while the repo ships processing_paddleocr_vl.py). A
referenced .py that IS present but cannot be fetched, and a repo-listing failure,
still fail closed.
Remote LoRA base resolution
- Distinguish a genuine 404 (not a LoRA / repo absent -> None) from a transient
error: the transient case is retried once, then logged as a WARNING (a missed
base is scanned by neither gate) rather than silently skipped.
Discard endpoint
- Treat .onnx and .ckpt as weights so a repo whose only heavy artifact is one of
those is never eligible for the declined-download purge.
Tests added for each: load-path scoping (safetensors/subdir/Nemotron-H shapes,
unknown-level fail-closed, pending non-block), stale own-repo auto_map ref, remote
LoRA transient retry, and the empty-config-list (all-404 -> []) semantics.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Make LoRA-base transient-warning test robust to logging backend
Assert on the logger object directly instead of capsys, so the test does not
depend on whether the real structlog logger or the module-stub logger is active
(which varies with test collection order).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Allow a repo with auto_map but no executable code (e.g. GGUF) instead of blocking
A config can declare an auto_map yet the repo ship NO executable .py -- most
commonly a GGUF repo whose config.json carries an auto_map copied from the original
model (e.g. unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF references
modeling_decilm.py, which the GGUF-only repo does not contain). A GGUF model loads
through llama.cpp, which never executes auto_map, and transformers cannot run a file
that is not present, so there is nothing to scan and trust_remote_code is a no-op.
The fail-closed change treated this empty result the same as "code is present but we
could not fetch it" and hard-blocked the load. Distinguish the two: repo_remote_code_files
now RAISES RemoteCodeUnscannable when code is present but cannot be fully fetched or
listed (offline / gated / transient / a present .py that 404s / a listing failure),
and returns an empty dict only when the listing succeeded and the repo genuinely ships
no executable .py. The consent gate blocks on the exception (fail closed) and allows the
empty case as a no-op. Real unscannable code still hard-blocks; eicar and CRITICAL/HIGH
custom code are unaffected.
Verified against all 37 unsloth/*Nemotron* models (two GGUF repos were false-blocked,
now load) and the existing matrix (eicar still blocks; DeepSeek-OCR / NVLM-D-72B still
prompt approvable consent). Tests updated to expect the raise for unscannable cases and
added for the no-executable-code no-op.
* Ignore vestigial auto_map in GGUF repos (llama.cpp never runs it)
A GGUF repo's config.json is often copied verbatim from the original
transformers model, auto_map and all, but a GGUF load goes through
llama.cpp which never executes auto_map, so the config is inert. Treat
a direct .gguf reference, and a repo that ships .gguf weights with no
.safetensors, as having no remote code so the consent flow is never
triggered. A mixed repo with both .gguf and .safetensors is still gated,
since the safetensors variant would load through transformers where
auto_map does run. The check sits behind the existing auto_map-present
gate so normal models pay no extra repo listing.
* Add scanner-result copy to the remote-code consent dialog
Make the consent dialog state the scan outcome in plain language for
every model. When the static scan finds nothing, reassure the user with
'Our automatic scanner did not flag any worrying files, but please
double check.' (shown only for the clean, approvable case). When the
scan flags custom code or unsafe files, label the list with 'Our
automatic scanner flagged issues including:'. The Hugging Face
attribution for unsafe files stays in the dialog description.
* Close GGUF-suffix consent bypass for repo ids ending in .gguf
The .gguf short-circuit in _config_has_auto_map skipped the scan for any
model name ending in .gguf, including a bare two-segment repo id like
'evil/model.gguf'. Such a repo can still ship safetensors plus auto_map
Python that transformers would execute, so skipping the scan was an
asymmetric bypass (file_security already scans those repos). Restrict the
short-circuit to genuine direct GGUF file references via
_is_direct_gguf_file_ref: a local .gguf path, or a remote repo_id plus
filename (three or more segments). A two-segment repo id named *.gguf now
falls through to the config scan and _is_gguf_repo file inspection, so it
only skips consent when it actually ships .gguf weights and no safetensors.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Align consent dialog body with the title and fix narrow-width overflow
The scan results (the 'Our automatic scanner...' label, finding/unsafe
cards, and the clean-scan reassurance) sat at the dialog's left padding
while the title and description were indented past the status icon, so
the body did not line up under the description. Move the title,
description and results into one column to the right of the icon so they
share a left edge, and let that column fill its width so the description
no longer wraps early.
Also stop a wide code snippet from pushing the dialog off-screen on
narrow viewports: AlertDialogHeader is a grid with place-items-center,
which sized the content row to its content; give the row w-full so it
fills the track, and add min-w-0 down the results chain so the snippet
scrolls inside its card instead of widening the dialog. Verified aligned
and contained from mobile portrait through ultrawide.
* Treat a repo as GGUF-only only when it ships no transformers weights
_is_gguf_repo excluded only .safetensors, so a repo with a .gguf and a
pytorch_model.bin (or .pt/.pth/.h5/.msgpack/.onnx/.ckpt) and no
safetensors was treated as GGUF-only and skipped the consent scan, even
though transformers can load that weight set and execute the repo's
auto_map code. Require the absence of ANY transformers-loadable weight
before treating the repo as a llama.cpp-only GGUF load. A genuine
GGUF-only repo (only .gguf) is still inert; a mixed repo with any pickle
or safetensors weight is gated. Adds a regression test across all the
non-safetensors weight formats.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Block flagged subdir weight shards referenced by a root index
The malware gate treated every subdirectory file as non-loadable, but
from_pretrained deserializes a subdir shard a root index references
(pytorch_model.bin.index.json -> shards/...-00001-of-00002.bin). Read the
root weight indexes and block a flagged subdir pickle the weight_map
points at; a flagged subdir pickle no index lists (NeMo nemo/*.distcp)
stays non-blocking, and an inconclusive index lookup fails closed.
* Pass hf_token to the export checkpoint load
ExportBackend.load_checkpoint scanned with hf_token in the worker but
loaded the weights unauthenticated, so a gated/private checkpoint passed
preflight then 401'd at from_pretrained. Add hf_token to load_checkpoint
and forward token to every from_pretrained branch; the worker passes the
command's hf_token.
* Scope created_by_scan to every HF cache the discard searches
created_by_scan used get_cache_path (active HF_HUB_CACHE only) while
/discard-remote-code deletes across active, legacy, and default caches. A
repo the user already had in a legacy/default cache was marked
scan-created and deleted on decline. Check all three caches for the repo
dir before declaring the scan created it.
* Scan the full .py closure of external auto_map repos
An auto_map cross-repo ref (owner/name--module.Class) only had its entry
file downloaded, but transformers also fetches that file's relative
imports from the same repo, so a dangerous helper.py was left outside the
scanned fingerprint. List each external repo's .py and scan the whole set
(plus the referenced entry files); fail closed if the repo cannot be
listed or fetched.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fail closed when a weight index cannot be fully read
_indexed_shard_paths treated a partial result as definitive: if one weight
index read cleanly but another failed transiently, it returned the shard
paths it did see. A flagged subdirectory pickle listed only by the index we
could not read would then be classed as "not a load input" and skipped,
re-opening the very fail-open this guard was added to close.
Return None whenever any index read is inconclusive, even if another read
cleanly, so the caller blocks the already-flagged subdir pickle. A repo that
ships no index files raises EntryNotFoundError for each (never inconclusive)
and still returns an empty set.
* Match cached repos case-insensitively in the created_by_scan guard
_repo_in_any_hf_cache resolved casing only against the active cache and then
probed every cache with an exact directory name. A case-variant already
present in a legacy or default cache (models--Unsloth--Foo for a scan of
unsloth/foo) was missed, so the repo was marked created_by_scan and deleted
on decline -- but discard_remote_code_download deletes case-insensitively,
so that delete would hit the user's pre-existing cache entry. Detect
case-insensitively too, mirroring the deletion path.
* Skip remote-code and security review for selected GGUF variants
validate_model ran the trust_remote_code and Hugging Face security-scan
preflight against the repo even when the selected artifact is a .gguf. A
GGUF loads through llama.cpp, which never executes the repo's auto_map
Python and never deserializes root pickle weights, so repo-level Transformers
artifacts (a config.json with auto_map, or an unsafe pytorch_model.bin next
to the .gguf in a mixed repo) are inert for that load. Gating the GGUF on
them is a false positive. Run both preflights only for non-GGUF loads.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scope the malware gate to actual load roots and serialized files
Two fixes to evaluate_file_security so it neither misses a load-path pickle nor
false-blocks an inert file:
- Honor subdirectory load roots. Spark-TTS / BiCodec call from_pretrained on the
snapshot's LLM subdirectory, so a flagged pickle directly under it is a
root-level load artifact there. A new load_subdirs parameter (set from the
model's audio type via security_load_subdirs) reclassifies those files relative
to the load root and looks for weight indexes under it, so a flagged shard in
that subdir is no longer skipped as "not root-level".
- Exempt source files. A root .py is never deserialized by from_pretrained;
executable repo code runs only through auto_map, which the remote-code consent
gate scans. Flagging a Python helper here would false-block a repo that merely
ships a build or train script.
* Scan a LoRA adapter and base as one consent unit, and gate MEDIUM code
A LoRA load runs both the adapter's and the base's repo code. The consent gate
scanned them separately and pinned one fingerprint per repo, so an adapter that
shipped its own auto_map code was either never shown in the dialog (which only
saw the base) or impossible to approve with the base's fingerprint.
evaluate_remote_code_consent_for_targets now scans all of a load's repos as a
single combined unit and pins ONE fingerprint over the union of their code, so
approving the load approves every repo's code together. evaluate_remote_code_consent
becomes a thin single-target wrapper, and an unscannable target fails the whole
load closed.
Also gate MEDIUM findings: like HIGH they now block pending pinned approval, so a
direct API caller cannot run flagged code by setting trust_remote_code=True
without consenting. Only a clean scan loads without a fingerprint.
* Preflight a LoRA load's adapter and base as one combined consent scan
scan_model_remote_code rewrote a LoRA adapter to its base and scanned only the
base for remote code, so the dialog never surfaced an adapter's own auto_map
code. Scan the adapter and base together through
preflight_remote_code_consent_for_targets, which pins one combined fingerprint
the worker gate accepts. The malware preflight is also scoped to each target's
load subdirectories.
* Apply combined consent and subdir-aware malware scan in load workers
Each load worker (inference, export, training) evaluated remote-code consent
once per target with a single shared fingerprint, so a LoRA adapter that ships
its own auto_map code could not be approved by the base's fingerprint. They now
scan the adapter and base together via evaluate_remote_code_consent_for_targets,
which pins one combined fingerprint over the union of their code. The malware
scan in each worker is also scoped to the model's load subdirectories so a
flagged pickle under a from_pretrained load subdir is not missed.
* Report a consistent trust_remote_code requirement after a model loads
validate_model reports requires_trust_remote_code from the YAML default OR the
raw auto_map, but the load, already-loaded, and status responses reported only
the YAML default. A custom-code model approved and loaded via auto_map was then
reported as not requiring trust_remote_code, so the frontend stored false and a
later retry or rollback sent trust_remote_code=false and failed.
A shared resolver reports the same requirement for a loaded model (a value
stored at load time, else the trust_remote_code the load used, else the YAML
default, else the raw auto_map check), and the load response persists it so the
status and already-loaded paths stay consistent. The selected-GGUF security
review is also scoped to the model's load subdirectories.
* Run the consent gate on training resume and for YAML-only trust_remote_code
Three frontend gaps left a model loading without the trust_remote_code it needs:
- The shared consent helper returned early when the scan found no auto_map and no
unsafe files, dropping a requirement that comes from a model's Studio YAML
default (e.g. GLM-4.7-Flash). It now grants the caller's requirement with an
empty pin instead of sending trust_remote_code=false.
- Resume-from-history called startTraining directly with no consent gate, so a
resumed run whose model needs custom code (or an old run with no approved
fingerprint) hit the worker block with no dialog. It now runs the same gate as
a fresh start.
- HF export passed requiresTrustRemoteCode=false for every HF source, so a
YAML-only model could not flip the flag before export. It now signals the
requirement for HF sources.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Cover both LoRA repos in validate, report GGUF as inert, purge all declined repos
Three follow-on gaps from the combined adapter+base consent work:
- validate_model resolved requires_trust_remote_code from the base alone, so a
LoRA adapter that ships its OWN auto_map code (with a plain base) was reported
as not needing trust_remote_code and the consent dialog never opened. It now
checks the [adapter, base] target set, matching the scan route and the workers
(which already gate both) and the security review already running over both.
- The already-loaded, loaded, and status responses for a selected GGUF reported
requires_trust_remote_code from the model's YAML default. A GGUF loads through
llama.cpp, which never executes the repo's auto_map Python, so the requirement
is inert for that load. They now report False, matching validate_model (which
already skips both gates for GGUF) so a status refresh cannot flip the flag
back on.
- The remote-code scan downloads both the adapter's and the base's config, but
created_by_scan tracked only the primary, so a base the scan was first to pull
into the cache was left on disk when the user declined. The scan now reports
scan_created_repos (every repo it newly cached) and the decline cleanup purges
each; created_by_scan stays for older clients. The frontend falls back to the
primary flag when the list is absent.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scan the repo the load fetches, purge external code on decline, harden consent pins
Six follow-on hardening fixes from a fresh review pass over the gate:
- The malware gate scanned the literal "Spark-TTS-0.5B/LLM" alias, but the trainer
downloads it as unsloth/Spark-TTS-0.5B and loads LLM/, so the alias 404'd and
failed open, missing a flagged LLM/ pickle. evaluate_file_security now resolves
the alias to the repo the loader fetches and scans LLM/ as a load root.
- security_load_subdirs relied only on tokenizer detection, which fails on an
unresolved alias or offline; it now also honors the Studio YAML audio_type
default, so a BiCodec LLM/ load root is not missed.
- The remote-code scan downloads external auto_map repos (owner/name--module.Class),
but the decline cleanup tracked only the model/adapter/base, leaving the external
untrusted code cached. The scan now enumerates external auto_map repos and reports
the ones it created in scan_created_repos, so a decline purges them too.
- External auto_map refs failed the whole load closed on a stale or mis-derived
dotted ref (sub.mod.py vs the real sub/mod.py) even though the actual file was
present and scanned. They now drop such refs when the repo listing is real, exactly
like the own-repo path; an empty/incomplete listing still fetches and fails closed.
- The combined consent fingerprint keyed code by the raw target string, so the scan
endpoint's canonicalized casing and a worker's raw user input produced different
pins for identical code, rejecting a valid approval. Hub repo ids are now folded to
lowercase in the key (local paths stay case-sensitive), so the pin tracks the code.
- Export threaded hf_token into the weight load but not into detect_audio_type /
is_vision_model, so a gated multimodal base 404'd in detection and fell through to
the text loader. Both probes now use the same token.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Thread the token through check-vision and guard the gate's parallel sites
The /check-vision endpoint classified a model without the hf_token, so a gated or
private vision model 404'd in the probe and was reported as a plain text model --
the same dropped-token shape as the export probes, at a sibling site. It now passes
the token like the neighboring /check-embedding endpoint.
Add deterministic consistency guards (tests/test_security_gate_consistency.py) that
enumerate the gate's parallel sites mechanically instead of relying on a review to
spot a missed sibling: every is_vision_model / is_embedding_model / detect_audio_type
caller under routes/ and core/ must thread the token, every GGUF response must report
trust_remote_code via the resolver or False (never the raw YAML default), and every
load worker that runs the malware or consent gate must resolve the LoRA base. A new
site that drops the token or mis-reports the requirement now fails CI directly.
* Narrow the LLM alias rewrite and make audio detection token-aware
Three fixes from the confirmatory review, one a regression from the previous round:
- _load_scan_target rewrote EVERY remote repo ending in "/LLM" to unsloth/<parent>,
so a real third-party repo named "<owner>/LLM" was scanned as unsloth/<owner>
while the loader still fetched the real repo -- a fail-open hole introduced when
the Spark-TTS alias handling was added. It now rewrites only a registry-known
bicodec alias; every other "/LLM" repo is scanned as itself.
- detect_audio_type cached results under the bare model name, so an unauthenticated
probe of a gated/private repo cached None and poisoned a later authenticated call
with the token. The cache is now keyed by (normalized_name, token_fingerprint),
matching the vision cache.
- The training fallback /check-vision call dropped the hf_token, misclassifying a
gated/private VLM when the config endpoint failed. It now passes the token, like
the getModelConfig call it falls back from; checkEmbeddingModel takes the token too.
Extend the consistency guards: every capability cache must be keyed by a tuple
including the token, so a cache re-declared as Dict[str, ...] fails CI.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Document the broad .py scan as deliberate and enforce it with a test
The remote-code scanner scans every .py in a repo once an auto_map exists, not
just the auto_map entry's static import closure. This is intentional: the entry
module can reach a sibling via an absolute import, importlib, or exec, none of
which a static relative-import closure follows, so closure-only scanning would be
a real bypass of a load-time RCE gate. The broad scan never under-scans; the cost
is that an unrelated benign script can over-block, which is the safe failure
direction (HIGH stays approvable; only CRITICAL hard-blocks).
Spell this out at both the local and remote scan sites so the choice reads as
deliberate, and add a test asserting an unrelated, never-imported .py is still
scanned -- so a future narrowing to the static closure fails CI.
* Purge a declined remote LoRA adapter the scan downloaded
scan_model_remote_code probed the created-by-scan state AFTER resolving the base,
but get_base_model_from_lora_identifier downloads a remote adapter's own
adapter_config.json, so the adapter looked already-cached and was dropped from
scan_created_repos. On decline the adapter -- including the auto_map .py the
preflight fetched -- was left on disk, defeating the "untrusted code is not left
on disk" guarantee for the adapter itself.
Snapshot the primary's cache state BEFORE base resolution and use it when marking
the adapter scan-created; on any probe error treat it as pre-existing so a decline
never deletes it. The base and external repos are unaffected (their configs are not
downloaded before their own probe). Add a test that models the mid-scan download
side effect, which the prior static-stub tests did not.
* Clear remote-code approval when the training model changes
Switching the training model from an approved custom-code model to a clean one
kept the previous model's trust_remote_code=true and approved fingerprint in the
store: setSelectedModel reset visionImageSize on a true switch but not the
remote-code approval. The clean model then trained with trust_remote_code=true,
which bypasses the compiler and disables fused cross-entropy.
Reset trustRemoteCode and approvedRemoteCodeFingerprint on a true model switch.
The new model's own YAML default is re-applied by loadAndApplyModelDefaults, and a
custom-code model still re-opens the consent dialog before training starts, so the
only change is that a clean model no longer inherits a stale approval.
* Trim verbose comments across the model-fetching hardening changes
Condense the explanatory comments and docstrings introduced across the
trust_remote_code consent gate, the malware/unsafe-file gate, the remote-code
scanner, the load workers, the model routes, and the security frontend into
fewer, tighter lines while preserving every security rationale (fail-open vs
fail-closed direction, the deliberate broad-scan anti-bypass note, the
empty-vs-unscannable distinction, stale-ref handling, and the alias-rewrite
spoof guard).
Comments and docstrings only. No code, logic, identifiers, or test behaviour
changed; verified comment-only via the AST/TypeScript checker (40/40), with the
backend test suite and frontend tsc green.
* Do not cache transient audio-detection failures
detect_audio_type cached _detect_audio_from_tokenizer's result
unconditionally, so a transient read failure (network error or 5xx,
returned as None) poisoned the cache and the later successful probe never
ran. Mirror the vision cache: _detect_audio_from_tokenizer now returns
(audio_type, definitive) and the caller caches only definitive results.
A read that succeeds with no audio tokens, or clean 404s for every
tokenizer path, stays a cacheable None; only a genuine transient failure
(connection error, timeout, 5xx, malformed body) skips the cache so the
next call retries.
---------
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* 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>
* 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>
pre-commit-ci[bot]