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

Studio: offer the latest transformers release for brand-new architectures (#7056)

Browse source 
* Studio: offer the latest transformers release for brand-new architectures

When a model's config.json model_type is absent from every installed
transformers overlay (base 4.57.x and the .venv_t5_530/550/510 sidecars),
Studio now checks, unauthenticated and cached, whether the newest
transformers ships it:

- utils/transformers_latest.py fetches the latest release version from
  https://pypi.org/pypi/transformers/json and the CONFIG_MAPPING_NAMES
  sources for that tag and for main from raw.githubusercontent.com
  (never api.github.com), parsing them with the same AST extractor the
  static router uses (no code execution, no trust_remote_code). Results
  are cached in memory and in a JSON snapshot under studio_root()/cache
  with a one day ttl; fetches are bounded to 5s with one retry and a
  failure backoff, and offline mode or the new kill switch
  UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS=1 short-circuits to None.

- POST /api/inference/validate gains requires_transformers_upgrade plus
  a transformers_upgrade payload (model_type, pypi_version,
  supported_in_pypi, supported_in_main) so the frontend can raise the
  install consent dialog before /load, mirroring the existing
  remote-code consent flow. The check fires only when the model_type is
  unknown to all installed overlays and the hardcoded tier tables.

- POST /api/inference/install-latest-transformers provisions a new
  persistent .venv_t5_latest sidecar after user consent, pinned to the
  exact PyPI version (re-verified server-side) with the same
  --target/--no-deps recipe as the fixed sidecars. A JSON pin marker
  inside the dir records the installed package set, so restarts
  revalidate it and routing resolves the new highest-ranked tier
  automatically. A dependency preflight (compat_plan) compares the
  release's requires_dist against the running env: unsatisfied
  tokenizers/safetensors floors are shadow-installed as exact pins into
  the sidecar, anything else unsatisfied blocks the install with a
  clear message.

Routing for every already-supported model_type is unchanged: the
hardcoded lists and the 530/550/510 static resolver run first, the new
tier only participates once its venv exists, and the probe order gains
the latest sidecar only when provisioned. Verified against live PyPI
and GitHub (transformers 5.13.0: 674 model_types, 26 absent from all
installed overlays, e.g. cosmos3_omni; 4 dev-only on main) and with a
real sidecar install plus restart persistence. 64 new tests; the
existing 200-test transformers_version suite passes unchanged.

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

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

* Latest-transformers check: fetch outside the lock, serialize installs

Release the module lock during the network refresh so a slow fetch cannot
stall other threads in the ASGI pool; concurrent callers during a fetch get
None (the graceful fallthrough) via an in-flight flag instead of stacking
fetches. Serialize install_latest_transformers with an in-progress flag so
concurrent consents cannot race the sidecar delete and recreate; the loser
gets a structured already-in-progress refusal.

* Latest-transformers check: LoRA bases, pin-gated mapping, live reverify

Run the upgrade check over the [adapter, base] target set so a LoRA whose
base model is a brand-new architecture surfaces the prompt (the worker
activates transformers for the base, not the adapter).

Gate the latest overlay's mapping lookup on a valid pin marker, matching
activation and the probe order, so a partial or manual .venv_t5_latest dir
cannot be routed to and then refused at activation.

Re-verify the requested version against a live PyPI snapshot at install
time, falling back to the cached one on fetch failure, so a release
published inside the cache TTL is not silently missed.

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

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

* Latest-transformers check: nested config types and latest-tier vision probe

Collect every model_type in the config (top level plus each nested
sub-config) and signal on the first one missing from all installed
overlays, so a supported wrapper carrying a brand-new backbone still
surfaces the upgrade prompt; wrappers instantiate sub-configs through
CONFIG_MAPPING and would fail on the nested type.

Route the vision capability subprocess through the pinned latest sidecar
when the model resolves to the latest tier, so latest-only VLMs are not
misclassified as text-only; every other tier keeps the 5.5 sidecar used
today.

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

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

* Latest tier: nested routing, vision probe after raw miss, safe upgrades

Route by every model_type in the config: a nested sub-config type can raise
the tier (wrappers instantiate sub-configs through CONFIG_MAPPING), so a
supported wrapper with a latest-only backbone routes to latest once
installed instead of staying on default. An unknown nested type never
vetoes; the primary type keeps its previous semantics. The collector is
shared with the upgrade checker.

Vision detection: when the raw heuristics say False for a model that routes
to the latest tier, run the AutoConfig subprocess under the pinned latest
sidecar instead of trusting heuristics built from older transformers.

Provisioning: stage-and-swap. Build the new sidecar in .venv_t5_latest.staging
and swap it in only when the install and pin marker are complete, so a failed
upgrade never destroys a previously working sidecar; restore the old dir if
the final swap fails.

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

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

* Latest-transformers checker, vision subprocess, and cache fixes

Require the latest release to support every missing model_type (the
primary included) before prompting; a nested-only match cannot make the
model loadable, so no install is offered for it.

The vision-check subprocess now unions the active sidecar's own
registry mappings into the inlined parent-process detection sets, so
architectures only the sidecar knows classify correctly.

A successful sidecar install clears the tier probe cache, the latest
tier's model_type mapping, and the vision-detection cache so the new
venv takes effect without a restart. Tests for all three.

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

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

* Aggregate upgrade support flags and keep install off /v1

The upgrade signal now reports supported_in_pypi only when the latest
release covers every missing model_type; a mix with a main-only nested
type surfaces as dev-only so no PyPI install is offered that would
still fail at load. The consented install endpoint moves to
studio_router so it is not reachable through the OpenAI-compatible /v1
mount. Tests for both.

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

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

* Honor the latest-transformers kill switch in routing

With UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS set after the sidecar was
provisioned, the latest tier still joined mapping and probe routing
because only the pin was checked. Both admission points now also check
the kill switch, so operators can roll back a problematic sidecar
without deleting files.

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

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

* Repair the latest sidecar through stage-and-swap

The lazy repair path installed into the live .venv_t5_latest, which
_ensure_venv_dir wipes first, so a failed repair deleted the pinned
sidecar and its marker. Both the consented install and the repair now
share one stage-and-swap helper: the incomplete-but-pinned dir survives
any failure and a later attempt can still repair it.

* Tighten comments

* Remove the staging dir when a latest-sidecar install fails

A pip failure inside _ensure_venv_dir returns False without raising, so
the except cleanup never ran and the partial .venv_t5_latest.staging
leaked until a later attempt. Also note on the validate response fields
that frontend consumption ships in the follow-up PR.

* Add the transformers-upgrade consent dialog to the frontend

When /validate reports requires_transformers_upgrade, every explicit load
path (chat runtime and the compare composer) now pauses on a consent
dialog modeled on the remote-code one: it names the model_type and the
latest PyPI transformers version, and on Accept calls
/api/inference/install-latest-transformers itself, shows an installing
state, and resumes the original load automatically on success. Errors
surface in the dialog with a retry; Cancel aborts the load like the
trust dialog's deny path. Architectures shipped only on transformers
main get a dev-only notice with no install button. Background auto-load
skips upgrade-requiring candidates instead of prompting, mirroring the
trust_remote_code rule. The dialog mounts once in the root layout and
runs before the security dialogs, since no load can proceed without the
runtime.

* Route a non-installable new architecture to the custom-code consent as a last resort

When the upgrade dialog has no installable PyPI release (the architecture
is only on transformers main, which Studio never installs), the dialog now
says so explicitly, and when the model also declares custom (auto_map)
code it offers Continue with custom code: resolving the paused load into
the existing trust_remote_code consent gate instead of hard-aborting.
Models with no custom code keep the Cancel-only notice. The backend
returns no upgrade signal at all for architectures unknown to both PyPI
and main, so those still route straight to the unchanged security gate.

* Force a 16-bit load for models on the latest-transformers sidecar

Live validation with Zyphra/ZAYA1-8B (model_type zaya, shipped by
transformers 5.13.1 but unknown to every installed tier) surfaced a
generation crash when the consented sidecar load kept the default bnb
4-bit quantization: transformers' grouped-MoE kernels feed the packed
uint8 expert weights straight into torch._grouped_mm, and generation
dies (plain 16-bit works). New latest_tier_active_for() mirrors the
sidecar activation's tier resolution and never raises; the inference
worker flips load_in_4bit off when it reports true, and the load route
applies the same flip so the pre-load VRAM guard and the worker command
agree. Fixed tiers are untouched. With the guard, ZAYA1-8B loads and
generates correctly in Studio chat.

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

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

* Offer the custom-code fallback when a latest-sidecar install fails

* Fail remote mapping fetches wholesale and mirror the 16-bit flip in validate

A transient fetch or parse failure of one auto-mapping file no longer caches
a partial latest-release map for the TTL (a real 404 on pre-5.10 tags is
still tolerated), and validate_model now applies the same latest-sidecar
16-bit sizing flip as /load before the training guard so the two agree.

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

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

* Tighten comments in the latest-transformers changes

* Resolve remote LoRA bases, fold nested tiers, and guard the sidecar swap

latest_tier_active_for now resolves a remote adapter's base model the same
way worker pre-activation does (and returns early without a sidecar pin), a
hardcoded fast-path tier is raised when a nested sub-config's model_type
needs a higher sidecar, and the install route refuses to swap .venv_t5_latest
while training runs on it and unloads a latest-tier chat model first.

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

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

* Gate the sidecar install on worker liveness and size installable upgrades 16-bit

The install route now refuses while any training or export runs (tier
re-resolution without the load token is unreliable for gated repos), holds
the inference lifecycle gate across the unload and the swap so no load can
interleave, and passes the model name to unload_model. validate_model runs
the upgrade check before the training guard and sizes an installable
upgrade as 16-bit, matching what /load and the worker will force after the
consented install.

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

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

* Close the sidecar install races and honor the kill switch over cached mappings

Training starts and mutating export routes now refuse while a transformers
install is in progress (shared is_install_in_progress flag), the chat unload
and idle export-worker teardown moved into a before_swap hook that runs only
once the staged install succeeded, and _config_model_types checks the kill
switch before returning a cached latest mapping.

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

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

* Reserve the sidecar swap before the gate wait and abort it on failed teardown

The install-in-progress flag moved into a shared sidecar swap reservation in
transformers_version, taken by the install route before awaiting the
inference lifecycle gate (so training and export starts see it for the whole
window) and by the lazy .venv_t5_latest repair path. The before_swap hook
now raises when the chat unload or export teardown reports failure, leaving
the previous sidecar untouched.

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

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

* Back the sidecar swap reservation with a cross-process lock file

The lazy repair runs inside worker subprocesses, where a module-level flag
is invisible to the parent's route checks. The reservation now also creates
a lock file next to .venv_t5_latest (O_EXCL, owner-only removal, stale after
two hours for crashed owners), so is_install_in_progress sees a repair from
any Studio process.

* Hand the swap reservation to the installer thread and harden pre-swap teardown

A cancelled install request no longer releases the reservation while the
installer thread is still staging (the thread owns and releases it, shielded
from cancellation). The route refuses while another inference request is
generating, export teardown runs before the chat unload and is judged by
worker liveness rather than the cleanup return value, and a live inference
worker with no active model (failed load residue) is shut down before the
swap.

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

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

* Keep the lifecycle gate with the installer and recheck the swap at spawn time

The gate moved into the shielded install task so a cancelled POST cannot
release the guard /load honors while the installer still runs, cached latest
probe results are ignored while the kill switch is set, and the training and
export subprocess spawns recheck the sidecar swap reservation right before
spawning (the route-level guards are one-shot and validation can outlast an
install's start).

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

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

* Close the spawn-registration windows against the sidecar install

Training marks the spawn in progress before its reservation recheck and
is_training_active honors the flag, so the install route sees a start that
has passed proc.start() but not yet recorded _proc. Export load-checkpoint
rechecks the reservation after setting _export_active and before tearing
down the old worker, so losing the race keeps the loaded checkpoint instead
of surfacing a 500.

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

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

* Refine the install-window interleavings around worker teardown

The inference busy count is rechecked under the lifecycle gate (streams
start by taking that gate, so nothing slips past a held gate), the training
handshake moved ahead of the VRAM-freeing before_spawn hook so a lost race
leaves chat/export intact, the export spawn-time check is op-aware (inside
an active op the install is the side that aborts), and the Xet-stall respawn
waits out a transient reservation instead of stranding the run.

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

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

* Track the install's server-side unload and guard export ops against the swap

The upgrade dialog store records when its install actually ran (the server
unloads the active chat model before swapping), and the load flow then marks
the previous model as unloaded so a later cancelled gate still triggers
rollback; the custom-code fallback leaves the flag unset. _run_export gained
the same reservation handshake as load_checkpoint so an install cannot block
behind an hours-long export op instead of returning 409.

* Tighten comments in the install-guard and upgrade-consent changes

* Surface install-race refusals cleanly and roll back after a failed swap unload

/load refuses while the sidecar swap is reserved so a load cannot succeed
and immediately be unloaded by the pre-swap teardown, worker starts that
lose the install race raise a typed SidecarSwapInProgress mapped to 409
instead of a 500, the install response reports model_unloaded even on a
structured failure so the client can restore its state, and the compare
flow tracks the server-side unload like the primary load path and clears a
stale checkpoint on abort.

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

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

* Type the export install races, scope the lock release, and keep the unload signal

Export load-checkpoint and export ops raise SidecarSwapInProgress (mapped to
409 in every export route) instead of a 400-shaped failure, the export spawn
check distinguishes repair reservations (always refused) from install ones
(op-aware), the swap lock release only unlinks a lock this process wrote so
a stale-superseded owner cannot drop the new owner's live lock, and the
frontend unload signal survives a superseding consent via read-and-clear
consumption instead of a reset.

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

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

* Finalize a stalled run when the respawn loses the install race and latch the unload signal

The Xet-stall respawn timeout now finalizes the run as a failure instead of
raising into the pump's broad finalization catch (which stranded it in a
training state with no worker), and a successful install retry ORs the
model_unloaded signal with the latched value so a failed-after-unload first
attempt still triggers rollback.

* Recheck the swap under the load gate and latch the unload before resolver checks

/load rechecks the sidecar reservation after acquiring the lifecycle gate
(an install can reserve while the load queues on it), and the dialog store
latches model_unloaded as soon as the install response arrives, before any
resolver-identity guard, so a superseded consent's unload still reaches
whichever load consumes the signal next.

* Report cleared-state unload failures, guard queued installs, and fold name tiers

A failed chat unload that still cleared the orchestrator's model state now
reports model_unloaded so the client rolls back, the installer aborts with
a 409 when a model load completed while it waited on the lifecycle gate,
and the fixed-tier name fast path consults the config mapping when a latest
sidecar is pinned so an accepted upgrade routes to the sidecar it installed
(no I/O added to the unpinned path).

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

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

* Report cleared-state unload failures and harden the spawn handshake flag

The failed-unload branch in before_swap now detects that the orchestrator
cleared its model state and reports model_unloaded before aborting (the
earlier commit claimed this fix but a scripting error dropped the edit),
the installer's queued-load check compares a load generation counter so a
same-model reload is caught, and both training spawn sites wrap everything
after the handshake in a guard that resets _spawn_in_progress on any
exception so a failed start cannot wedge is_training_active.

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

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

* Bump the load generation when the load is published, not at load start

A start-time bump is already visible when the installer snapshots mid-load,
so a same-model reload completing after the snapshot looked unchanged and
could be unloaded by the swap. The counter now increments alongside the
active_model_name publish.

* Self-heal a broken pinned sidecar, guard lazy repairs, and refresh stale retries

A valid pin whose transformers source dir vanished now triggers the repair
from the routing path (with a five minute backoff after failures) instead of
silently routing latest-only models to older tiers, the lazy repair refuses
while parent-visible chat/training/export workers are active since it has no
teardown of its own, and a version-mismatch install failure carries the
superseding release so the dialog's Retry re-requests a version that can
succeed.

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

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

* Flip latest-tier loads to 16-bit outside chat and protect export state

Training and export workers now apply the same latest-sidecar 16-bit flip
as the chat worker so a brand-new grouped-MoE architecture cannot reach bnb
4-bit through those paths, the latest-tier vision override returns None on
an inconclusive probe so a transient failure is not cached as not-vision,
and the install route refuses while an idle export checkpoint is loaded
rather than discard it with no rollback signal on a failed swap.

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

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

* Address parallel-review findings on the sidecar guards and install checks

The training route sizes latest-tier jobs 16-bit before GPU selection, the
inference subprocess spawn rechecks the swap reservation like training and
export (covering the OpenAI auto-switch path) with the typed error mapped
to a retryable 409, compat_plan blocks the install when dependency metadata
cannot be fetched instead of proceeding unverified, snapshot model-type
lists must contain only strings, and pin-marker package specs are validated
against the sidecar's own package set before ever reaching pip.

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

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

* Parent-only repairs, live-owner locks, remote-base activation, pre-teardown recheck

Lazy sidecar repairs now refuse inside worker children (whose empty backend
singletons cannot see live siblings) and run only in the parent where the
active-worker guard is real, swap-lock staleness requires the owner pid to
be dead so a slow live install is never superseded, both activation entry
points resolve a remote adapter's base model like the inference worker and
latest_tier_active_for already do, and load_model rechecks the reservation
before tearing down the old worker so losing the race keeps the current
model loaded.

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

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

* Check workers under the repair reservation and keep state on refused swaps

The lazy repair now reserves first and checks workers under the reservation
(worker starts set their active markers before rechecking, so every
interleaving aborts one side), with export ops and in-flight inference loads
counted as active. The inference pre-teardown and spawn guards refuse only
repair reservations since an install shares the load's lifecycle gate and
aborts via its queued-load snapshot, a SidecarSwapInProgress raised before
teardown no longer clears the live model mirrors, and an export spawn abort
after teardown clears current_checkpoint so the page cannot claim a loaded
checkpoint with no worker.

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

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

* Repair a present-but-incomplete latest sidecar from routing

The routing self-heal only fired when the pinned sidecar's transformers/
dir was missing. A sidecar that kept transformers/ but lost another pinned
package still routed models to the latest tier, and workers refuse
parent-only repairs, so every load failed until a manual reinstall. Routing
now validates the full pin (via _venv_dir_is_valid) and repairs any
incomplete sidecar under the same swap reservation and 5-minute backoff.

* Treat an unrepaired latest sidecar as unavailable in routing

When the pinned sidecar is incomplete and the lazy repair fails (offline,
pip failure, workers active) or is inside the backoff window, routing
returned the source dir anyway, sending models to a tier whose worker
activation is known to fail. Return None instead so models an older tier
supports keep loading there until a repair succeeds, matching the behavior
when the sidecar dir is missing entirely.

* Harden sidecar swap and repair against crash, survivor, and 16-bit paths

Reclaim a swap lock as soon as its recorded owner PID is dead instead of
waiting out the two-hour cutoff, so a crash mid-install no longer wedges
/load, training, export, and repair for hours. A lock whose PID cannot be
read yet still uses the long cutoff so the create-before-write window is
never mistaken for dead.

Probe process liveness with OpenProcess on Windows: os.kill(pid, 0) there
is CTRL_C_EVENT (a real Ctrl+C via GenerateConsoleCtrlEvent), not a
harmless check, and psutil is not always present.

Return whether _shutdown_subprocess actually killed the worker and keep the
live handle when it survives terminate/kill (an uninterruptible CUDA
syscall can outlive SIGKILL). The pre-swap liveness guard now trusts that
result, so the destructive .venv_t5_latest rename cannot proceed while a
live worker still holds sidecar modules.

Recover a sidecar stranded at .old when a swap's activation rename and its
rollback both fail: reading the pin restores it when no swap holds the
reservation, so latest-tier models are not permanently broken.

Resolve the latest tier in the parent for export loads and for explicitly
16-bit training runs, not only 4-bit ones: tier resolution self-heals an
incomplete sidecar, and repairs are parent-only, so those paths could not
recover before. Sidecar integrity and quantization are independent.

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

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

* Revert the parent-side latest-tier repair probe on training and export loads

The probe ran before the route freed VRAM, so a resident chat or export worker
made _workers_active_for_repair() refuse the parent-only repair; the route then
tore that worker down and spawned a child that also cannot repair, so an
incomplete sidecar still failed to load. Repairing correctly requires running the
repair between the worker teardown and the child spawn, decoupled from VRAM
sizing, which is a larger change tracked separately. Restore the prior behavior
so these paths match the reviewed form and do not partially attempt a repair that
cannot complete while workers are resident.

* Honor failed worker shutdowns on load and revalidate the cached latest mapping

The fresh-load paths spawned a new worker straight after _shutdown_subprocess
without checking its result, so a worker that outlived terminate/kill (a wedged
CUDA syscall) had its handle overwritten by the replacement while it still held
GPU memory, and is_worker_alive/the pre-swap guard could no longer see it. Both
the inference load and the export checkpoint load now abort when the old worker
did not exit, so the load can be retried once it does.

_config_model_types returned a cached latest mapping without re-checking the
sidecar, so a sidecar deleted or broken in-process after its first parse was
never re-validated: routing kept sending latest-only models to the stale latest
tier while activation failed. The cached latest mapping is now dropped and
re-resolved (self-healing) when the sidecar is no longer intact.

* Drop cached latest mapping when the pin is gone; keep 4-bit for custom-code fallback

_latest_sidecar_intact now returns False when the pin marker itself is gone, not
just when a pinned package is missing. Otherwise a cached latest mapping outlived
a deleted pin: _config_model_types kept returning it, so routing sent latest-only
models to a tier whose worker activation then failed (no pinned version) until
restart. It now drops the cache and re-resolves to no latest tier. The
_overlay_transformers_dir caller already gates on a present pin, so it is
unaffected.

validate_model forced 16-bit sizing whenever a PyPI upgrade was merely offered,
even for a model that can fall back to its own auto_map code. /load loads such a
model 4-bit without the install, and the install route refuses while training is
active, so 16-bit sizing here returned a VRAM 409 for the only viable 4-bit path.
The offered-upgrade flip is now gated on the absence of a custom-code fallback;
an already-active latest sidecar still always sizes 16-bit.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-15 05:25:26 -07:00 • committed by GitHub
commit 815f242970
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
29 changed files with 4637 additions and 204 deletions
Download patch file Download diff file Expand all files Collapse all files

90
studio/backend/core/export/orchestrator.py
View file

@ -132,6 +132,11 @@ class ExportOrchestrator:
"""True while an export / load / cleanup command is running."""
return self._export_active
def is_worker_alive(self) -> bool:
"""True while the persistent export subprocess is running (op or idle)."""
proc = self._proc
return proc is not None and proc.is_alive()
def was_cancelled(self) -> bool:
"""True if the in-flight (or most recent) run was cancelled by the user."""
return self._cancel_requested
@ -204,6 +209,23 @@ class ExportOrchestrator:
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new export subprocess."""
# Last-resort recheck for spawns outside an active op. Inside an op, _export_active is set and
# load_checkpoint already rechecked, so a reservation here is an install about to observe
# is_export_active() and abort; raising would kill this export for an install that never proceeds.
from utils.transformers_version import sidecar_swap_in_progress
from utils.transformers_version import sidecar_swap_kind
_swap_kind = sidecar_swap_kind()
# Inside an active op an INSTALL reservation is about to abort on the
# is_export_active check, but a lazy REPAIR has no such check and can be
# rebuilding the sidecar right now, so it must always refuse the spawn.
if _swap_kind == "repair" or (_swap_kind is not None and not self._export_active):
from utils.transformers_version import SidecarSwapInProgress
raise SidecarSwapInProgress(
"A transformers installation is replacing the latest sidecar; "
"retry when it completes."
)
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
@ -231,11 +253,17 @@ class ExportOrchestrator:
adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep)
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
"""Gracefully shut down the export subprocess."""
def _shutdown_subprocess(self, timeout: float = 10.0) -> bool:
"""Gracefully shut down the export subprocess.
Returns True only once the worker is confirmed dead. If it survives
terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives
SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the
pre-swap liveness guard can still observe the survivor instead of a cleared
handle and refuse the destructive sidecar swap."""
if self._proc is None or not self._proc.is_alive():
self._proc = None
return
return True
self._drain_queue()
@ -265,10 +293,20 @@ class ExportOrchestrator:
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
# Survived SIGKILL (uninterruptible syscall): keep the handle so callers
# and the pre-swap guard see a live worker rather than a nulled one.
logger.error(
"Export subprocess still alive after terminate/kill; "
"preserving its handle for the pre-swap liveness check"
)
return False
self._proc = None
self._cmd_queue = None
self._resp_queue = None
logger.info("Export subprocess shut down")
return True
def _cleanup(self):
"""atexit handler."""
@ -409,14 +447,44 @@ class ExportOrchestrator:
self._export_active = True
op_success, op_message = False, ""
try:
# Handshake with the sidecar install route: _export_active is set above, so either this
# recheck refuses BEFORE tearing down the old worker (keeping the loaded checkpoint), or
# the install sees is_export_active() and 409s. The spawn-time recheck stays as a last resort.
from utils.transformers_version import sidecar_swap_in_progress
if sidecar_swap_in_progress():
from utils.transformers_version import SidecarSwapInProgress
op_message = (
"A transformers installation is replacing the latest "
"sidecar; retry when it completes."
)
raise SidecarSwapInProgress(op_message)
# Always kill any existing subprocess and spawn fresh.
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
if self._shutdown_subprocess() is False:
# Survivor still holds GPU memory (a wedged CUDA syscall outliving
# SIGKILL); its handle is kept so is_worker_alive() and the pre-swap
# guard still see it. Do not spawn a second worker over it -- fail so
# the load can retry once it exits.
op_message = (
"The current export worker did not exit and still holds GPU "
"memory; not starting a new checkpoint load over it. Retry shortly."
)
return False, op_message
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
try:
self._spawn_subprocess(sub_config)
except Exception:
# The old worker is already gone; a stale current_checkpoint
# would make the Export page claim a loaded checkpoint that
# the next op then fails on with "no subprocess running".
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
raise
try:
resp = self._wait_response("loaded")
@ -560,6 +628,18 @@ class ExportOrchestrator:
self._export_active = True
op_success, op_message, op_output_path = False, "", None
try:
# Handshake with the sidecar install route (see load_checkpoint): _export_active is set
# above, so this recheck refuses before the command is sent, or the install sees the active
# op and 409s. Without it, an install would block in cleanup_memory behind a long export op.
from utils.transformers_version import sidecar_swap_in_progress
if sidecar_swap_in_progress():
from utils.transformers_version import SidecarSwapInProgress
op_message = (
"A transformers installation is replacing the latest "
"sidecar; retry when it completes."
)
raise SidecarSwapInProgress(op_message)
cmd = {"type": "export", "export_type": export_type, **params}
try:
self._send_cmd(cmd)

11
studio/backend/core/export/worker.py
View file

@ -236,6 +236,17 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
checkpoint_path = cmd["checkpoint_path"]
max_seq_length = cmd.get("max_seq_length", 2048)
load_in_4bit = cmd.get("load_in_4bit", True)
# Latest-sidecar checkpoints load 16-bit here too: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (same flip as the chat worker).
if load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(checkpoint_path, cmd.get("hf_token")):
load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"export load (4-bit is disabled for brand-new architectures)",
checkpoint_path,
)
trust_remote_code = cmd.get("trust_remote_code", False)
# Auto-enable trust_remote_code for NemotronH/Nano models.

87
studio/backend/core/inference/orchestrator.py
View file

@ -174,6 +174,21 @@ class InferenceOrchestrator:
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new inference subprocess."""
# Same recheck as the training/export spawns, REPAIR reservations only: a
# repair swaps without holding the lifecycle gate this load's caller owns,
# while an install cannot swap until this gate is released (and then its
# queued-load snapshot aborts it), so tolerating installs here lets the
# load win instead of failing both sides. Also covers the OpenAI
# auto-switch path, which enters _load_model_impl without route guards.
from utils.transformers_version import (
SidecarSwapInProgress,
sidecar_swap_kind,
)
if sidecar_swap_kind() == "repair":
raise SidecarSwapInProgress(
"A transformers repair is replacing the latest sidecar; retry when it completes."
)
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
@ -210,12 +225,24 @@ class InferenceOrchestrator:
if self._cancel_event is not None:
self._cancel_event.set()
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
"""Gracefully shut down the inference subprocess."""
def is_worker_alive(self) -> bool:
"""True while the inference subprocess is running, even with no model
active (a failed load can leave a live worker holding sidecar modules)."""
proc = self._proc
return proc is not None and proc.is_alive()
def _shutdown_subprocess(self, timeout: float = 10.0) -> bool:
"""Gracefully shut down the inference subprocess.
Returns True only once the worker is confirmed dead. If it survives
terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives
SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the
pre-swap liveness guard can still observe the survivor instead of a cleared
handle and refuse the destructive sidecar swap."""
self._stop_dispatcher() # before killing subprocess
if self._proc is None or not self._proc.is_alive():
self._proc = None
return
return True
# 1. Cancel any ongoing generation first (instant via mp.Event)
self._cancel_generation()
@ -252,12 +279,22 @@ class InferenceOrchestrator:
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
# Survived SIGKILL (uninterruptible syscall): keep the handle so callers
# and the pre-swap guard see a live worker rather than a nulled one.
logger.error(
"Inference subprocess still alive after terminate/kill; "
"preserving its handle for the pre-swap liveness check"
)
return False
self._proc = None
self._cmd_queue = None
self._resp_queue = None
self._cancel_event = None
self._drain_event = None
logger.info("Inference subprocess shut down")
return True
def _cleanup(self):
"""atexit handler."""
@ -882,6 +919,13 @@ class InferenceOrchestrator:
# Public API — same interface as InferenceBackend
# ------------------------------------------------------------------
# Monotonic count of PUBLISHED loads; lets the install route detect a load
# (including a same-model reload) that completed while it waited on the gate.
# Bumped when the load result is published, not at load start: a start-time
# bump is already visible when the installer snapshots mid-load, so the
# completed reload would look unchanged and get unloaded by the swap.
load_generation: int = 0
def load_model(
self,
config, # ModelConfig
@ -935,13 +979,36 @@ class InferenceOrchestrator:
sub_config["resolved_gpu_ids"] = resolved_gpu_ids
sub_config["gpu_selection"] = gpu_selection
# Recheck the sidecar reservation BEFORE tearing the old worker down,
# for REPAIRS only: an install holds this same lifecycle gate, so it
# cannot swap while this load runs, and its queued-load snapshot
# aborts it after this load publishes -- the load wins cleanly.
# Raising here (repair) keeps the current model loaded.
from utils.transformers_version import (
SidecarSwapInProgress,
sidecar_swap_kind,
)
if sidecar_swap_kind() == "repair":
raise SidecarSwapInProgress(
"A transformers repair is replacing the latest sidecar; "
"retry when it completes."
)
# Always kill the existing subprocess and spawn fresh: reusing one
# after unsloth patches torch internals breaks getsource on reload.
if self._ensure_subprocess_alive():
self._cancel_generation()
time.sleep(0.3)
self._shutdown_subprocess()
if self._shutdown_subprocess() is False:
# The worker survived terminate/kill (e.g. a wedged CUDA syscall that
# outlives SIGKILL). Its handle is kept, so is_worker_alive() and the
# pre-swap guard still see it; do not spawn a second worker over one
# still holding GPU memory. Fail so the load can retry once it exits.
raise RuntimeError(
"The current inference worker did not exit and still holds GPU "
"memory; not starting a new model over it. Retry shortly."
)
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
@ -1030,6 +1097,7 @@ class InferenceOrchestrator:
return False
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
self.load_generation += 1
# A load always spawns a fresh subprocess holding only this model, so
# mirror that. A lingering stale name would pass unload_model's "not in
# self.models" guard, and the worker's absent-name fallback would unload
@ -1061,8 +1129,15 @@ class InferenceOrchestrator:
self.models.clear()
raise Exception(error)
except Exception:
except Exception as exc:
self.loading_models.discard(model_name)
from utils.transformers_version import SidecarSwapInProgress
if isinstance(exc, SidecarSwapInProgress) and self._ensure_subprocess_alive():
# Raised before the old worker was torn down: the previous model
# is still live, so keep the mirrors (clearing them would let the
# installer treat the worker as inactive and kill it unreported).
raise
self.active_model_name = None
self.models.clear()
raise

12
studio/backend/core/inference/worker.py
View file

@ -291,6 +291,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
hf_token = _clean_token(config.get("hf_token"))
load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True))
# Latest-transformers sidecar models load 16-bit: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (e.g. grouped-MoE torch._grouped_mm).
if load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(config["model_name"], hf_token):
load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"load (4-bit is disabled for brand-new architectures)",
config["model_name"],
)
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token):
trust_remote_code = True

312
studio/backend/core/training/training.py
View file

@ -881,93 +881,125 @@ class TrainingBackend:
else:
defer_auto_selection = True
# Synchronous validation passed -> free VRAM (export + chat) now, before
# auto-selection and the spawn, so placement sees the freed memory.
if before_spawn is not None:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
# Handshake with the sidecar install route: mark the spawn in progress BEFORE rechecking
# the reservation, so either this recheck aborts, or the install's is_training_active()
# sees this flag (or the recorded proc) and refuses.
from utils.transformers_version import sidecar_swap_in_progress
if defer_auto_selection:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(None, **gpu_selection_kwargs)
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
from .worker import run_training_process
self._spawn_in_progress = True
if sidecar_swap_in_progress():
self._spawn_in_progress = False
from utils.transformers_version import SidecarSwapInProgress
raise SidecarSwapInProgress(
"A transformers installation is replacing the latest sidecar; "
"retry when it completes."
)
# Any exception between the handshake above and the flag reset below would
# otherwise leave _spawn_in_progress latched, wedging is_training_active
# (and the install route) until restart.
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
# Synchronous validation passed -> free VRAM (export + chat) now, before
# auto-selection and the spawn, so placement sees the freed memory. Runs AFTER the handshake
# so a lost race to an install can't tear down chat/export for a training run that never spawns.
if before_spawn is not None:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
if defer_auto_selection:
try:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
None, **gpu_selection_kwargs
)
except Exception:
# Flag is already set; a failed GPU selection must not leave is_training_active stuck True.
self._spawn_in_progress = False
raise
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
from .worker import run_training_process
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
self._spawn_in_progress = False
return False
logger.info("Training subprocess started (pid=%s)", proc.pid)
# Reset state (old pump thread dead, proc.start() succeeded).
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._complete_seen.clear()
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
)
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
self.grad_norm_history.clear()
self.grad_norm_step_history.clear()
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
self._db_create_in_progress = False # a stale watchdog create can't block this run
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Start each job Xet-first; keep config so a stall can respawn over HTTP.
self._last_full_config = config
self._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
self._spawn_in_progress = False
return True
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
return False
logger.info("Training subprocess started (pid=%s)", proc.pid)
# Reset state (old pump thread dead, proc.start() succeeded).
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._complete_seen.clear()
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
)
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
self.grad_norm_history.clear()
self.grad_norm_step_history.clear()
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
self._db_create_in_progress = False # a stale watchdog create can't block this run
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Start each job Xet-first; keep config so a stall can respawn over HTTP.
self._last_full_config = config
self._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
return True
self._spawn_in_progress = False
raise
def stop_training(self, save: bool = True) -> bool:
"""Send stop signal to the training subprocess."""
@ -1266,50 +1298,84 @@ class TrainingBackend:
from .worker import run_training_process
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
new_proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
new_proc.start()
from utils.process_lifetime import adopt_pid
# This run is active, so an install request 409s rather than proceeds: a reservation seen here
# is transient (an aborting install or short lazy repair). Wait it out instead of stranding the
# stalled run; only a wedged reservation fails the respawn.
from utils.transformers_version import sidecar_swap_in_progress
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "error",
error_message = "Failed to recover stalled model download",
self._spawn_in_progress = True
_swap_wait_deadline = time.time() + 120
while sidecar_swap_in_progress() and time.time() < _swap_wait_deadline:
time.sleep(1)
if sidecar_swap_in_progress():
# Raising here would land in the pump's broad finalization catch and
# strand the run in a training state with no worker: finalize it as a
# failure explicitly instead.
self._spawn_in_progress = False
msg = (
"A transformers installation is replacing the latest sidecar; "
"cannot respawn the training worker."
)
logger.error(msg)
with self._lock:
self._progress.is_training = False
self._progress.error = msg
self._ensure_db_run_created()
self._finalize_run_in_db(status = "error", error_message = msg)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = new_proc
self._pump_thread = new_pump
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
# Reset the handshake flag on any unexpected failure past this point, so a
# crashed respawn cannot wedge is_training_active until restart.
try:
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
new_proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
new_proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
self._spawn_in_progress = False
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "error",
error_message = "Failed to recover stalled model download",
)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = new_proc
self._spawn_in_progress = False
self._pump_thread = new_pump
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
except Exception:
self._spawn_in_progress = False
raise
def _ensure_pump_alive(self) -> bool:
"""Restart the event pump if it crashed, even after the worker exited.
@ -1342,6 +1408,10 @@ class TrainingBackend:
def is_training_active(self) -> bool:
"""Check if training is currently active."""
# A spawn past its sidecar-swap recheck counts as active even before _proc is recorded,
# so an install cannot slip in mid-spawn.
if getattr(self, "_spawn_in_progress", False):
return True
# Self-heal a crashed pump first: a dead pump must never leave the worker
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()

15
studio/backend/core/training/worker.py
View file

@ -3019,11 +3019,24 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
),
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
)
# Latest-sidecar models load 16-bit here too: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (same flip as the chat worker).
_train_load_in_4bit = config["load_in_4bit"]
if _train_load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(model_name, hf_token):
_train_load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"training load (4-bit is disabled for brand-new architectures)",
model_name,
)
try:
success = trainer.load_model(
model_name = model_name,
max_seq_length = config["max_seq_length"],
load_in_4bit = config["load_in_4bit"],
load_in_4bit = _train_load_in_4bit,
full_finetuning = not use_lora,
hf_token = hf_token,
is_dataset_image = config.get("is_dataset_image", False),

63
studio/backend/models/inference.py
View file

@ -140,6 +140,27 @@ class ValidateModelRequest(BaseModel):
)
class TransformersUpgradeInfo(BaseModel):
"""A model architecture no installed transformers ships, but a newer release does."""
model_type: str = Field(
..., description = "config.json model_type unknown to every installed transformers"
)
pypi_version: Optional[str] = Field(
None, description = "Latest transformers release on PyPI at check time"
)
supported_in_pypi: bool = Field(
False,
description = "True if the latest PyPI release ships this model_type; Studio can "
"install it into a persistent sidecar after user consent.",
)
supported_in_main: bool = Field(
False,
description = "True if transformers GitHub main ships this model_type (dev-only; "
"not installable through Studio yet).",
)
class ValidateModelResponse(BaseModel):
"""Result of model validation.
@ -167,6 +188,48 @@ class ValidateModelResponse(BaseModel):
description = "Native training context length, read from the GGUF header when the file "
"is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
)
# Additive fields; the consuming consent dialog ships in a follow-up frontend PR.
requires_transformers_upgrade: bool = Field(
False,
description = "True when the model's architecture is unknown to every installed "
"transformers but a newer transformers ships it; the UI should offer the "
"install-latest-transformers consent dialog (or the dev-only notice).",
)
transformers_upgrade: Optional[TransformersUpgradeInfo] = Field(
None,
description = "Details for the transformers-upgrade dialog; set only when "
"requires_transformers_upgrade is true.",
)
class InstallLatestTransformersRequest(BaseModel):
"""Consented request to install the latest transformers release into a sidecar."""
version: str = Field(
...,
min_length = 1,
max_length = 64,
description = "Exact transformers version to install; must match the current "
"latest PyPI release reported by /validate.",
)
class InstallLatestTransformersResponse(BaseModel):
"""Result of the consented latest-transformers sidecar install."""
success: bool = Field(..., description = "Whether the sidecar was provisioned")
version: str = Field(..., description = "The requested transformers version")
message: str = Field(..., description = "Human-readable result")
model_unloaded: bool = Field(
False,
description = "Whether the active chat model was unloaded before the swap "
"(reported even on failure, so the client can restore its state)",
)
latest_version: Optional[str] = Field(
None,
description = "On a version-mismatch failure: the release that superseded "
"the requested one, so the client can retry with it",
)
class GenerateRequest(BaseModel):

35
studio/backend/routes/export.py
View file

@ -51,7 +51,17 @@ def _ensure_export_supported() -> None:
Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints
(scan/status/logs) are intentionally NOT gated so the Export page can still render the reason.
Also refuses (409) while a latest-transformers install is swapping .venv_t5_latest: an
export worker spawned mid-swap could activate a half-replaced sidecar.
"""
from utils.transformers_latest import is_install_in_progress
if is_install_in_progress():
raise HTTPException(
status_code = 409,
detail = "A transformers installation is in progress. Retry when it completes.",
)
from utils.hardware import export_capability
cap = export_capability()
@ -97,6 +107,11 @@ async def load_checkpoint(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error loading checkpoint: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -308,6 +323,11 @@ async def export_merged_model(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting merged model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -347,6 +367,11 @@ async def export_base_model(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting base model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -388,6 +413,11 @@ async def export_gguf(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting GGUF model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -428,6 +458,11 @@ async def export_lora_adapter(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True)
raise HTTPException(
status_code = 500,

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

@ -1693,6 +1693,9 @@ from models.inference import (
CompletionUsage,
ValidateModelRequest,
ValidateModelResponse,
TransformersUpgradeInfo,
InstallLatestTransformersRequest,
InstallLatestTransformersResponse,
TextContentPart,
ImageContentPart,
ImageUrl,
@ -3836,11 +3839,25 @@ async def load_model(
GGUF models load via llama-server (llama.cpp) instead of Unsloth.
"""
# A sidecar install that has reserved the swap must not lose to a load that
# then gets unloaded by the pre-swap teardown. Rechecked under the gate: an
# install can reserve while this request queues on the gate, so the pre-gate
# check alone is only a fast path.
from core.inference.llama_keepwarm import inference_lifecycle_gate
from utils.transformers_version import sidecar_swap_in_progress
_swap_409 = HTTPException(
status_code = 409,
detail = "A transformers installation is in progress. Retry when it completes.",
)
if sidecar_swap_in_progress():
raise _swap_409
# Hold the lifecycle gate across the load so idle auto-unload can't unload the
# model mid-load. Auto-switch calls _load_model_impl directly since it already
# holds this gate.
from core.inference.llama_keepwarm import inference_lifecycle_gate
async with inference_lifecycle_gate():
if sidecar_swap_in_progress():
raise _swap_409
return await _load_model_impl(request, fastapi_request, current_subject)
@ -4037,6 +4054,17 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
f"Resolved load_in_4bit={effective_load_in_4bit} for '{model_log_label}' "
f"from adapter_config.json / base model (requested {request.load_in_4bit})"
)
# Latest-sidecar models load 16-bit (worker refuses bnb 4-bit); size the guard
# to match. Off-loop: tier resolution reads configs.
if effective_load_in_4bit and not config.is_gguf:
from utils.transformers_version import latest_tier_active_for
if await asyncio.to_thread(latest_tier_active_for, config.identifier, request.hf_token):
effective_load_in_4bit = False
logger.info(
f"Latest-transformers sidecar active for '{model_log_label}' - "
"sizing and loading in 16-bit (4-bit is disabled for brand-new "
"architectures)"
)
# Refuse a load that would OOM active training, before the unload step below
# frees the resident model. Off-loop: guard does sync nvidia-smi / HF work.
@ -4470,6 +4498,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
logger.warning("GGUF runtime missing while loading '%s': %s", model_log_label, e)
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Lost the spawn-time race to a sidecar install/repair: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
# Friendlier message for models Unsloth cannot load.
if native_grant_backed:
redacted_msg = redact_native_paths(str(e))
@ -4598,16 +4631,6 @@ async def validate_model(
detail = "gpu_ids is not supported for GGUF models yet.",
)
effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
# Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
)
# Both checks cover the [adapter, base] set (matching the scan route and workers):
# either repo can ship auto_map code or a poisoned pickle.
@ -4624,16 +4647,69 @@ async def validate_model(
security_targets = list(dict.fromkeys(security_targets))
is_gguf = getattr(config, "is_gguf", False)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
# run the remote-code/security preflight for non-GGUF loads.
# Does a newer transformers ship this model_type? Static overlay first, cached
# PyPI/main snapshot only for unknown types. Never fails validation; run before
# the training guard so an installable upgrade sizes as 16-bit.
transformers_upgrade: Optional[TransformersUpgradeInfo] = None
if not is_gguf:
from utils.transformers_latest import check_upgrade_for_model
# Cover [adapter, base]: the worker activates transformers for the base model.
for _target in security_targets:
_upgrade = await asyncio.to_thread(
check_upgrade_for_model, _target, request.hf_token
)
if _upgrade is not None:
transformers_upgrade = TransformersUpgradeInfo(**_upgrade)
break
# Whether the model can load on the CURRENT transformers through its own remote
# code (auto_map, or the YAML trust default). Computed before the 16-bit flip
# because a model with this fallback still loads 4-bit without the offered install,
# exactly as /load does.
requires_trust_remote_code = False
requires_security_review = False
if not is_gguf:
requires_trust_remote_code = any(
_requires_trust_remote_code_for_model(_t, request.hf_token)
for _t in security_targets
)
# Mirror /load's latest-sidecar 16-bit flip so the guard sizes it the same way. An
# ALREADY-ACTIVE latest sidecar always forces 16-bit (the worker will). A merely
# OFFERED (not yet installed) upgrade forces 16-bit only when the model has NO
# custom-code fallback: with auto_map it still loads 4-bit on the current
# transformers (as /load does without a successful install), and the install route
# refuses while training is active, so sizing 16-bit here would 409 the only viable
# 4-bit path. /load re-sizes 16-bit after a successful install and re-guards there.
if effective_load_in_4bit and not is_gguf:
from utils.transformers_version import latest_tier_active_for
_install_only_upgrade = (
transformers_upgrade is not None
and transformers_upgrade.supported_in_pypi
and transformers_upgrade.pypi_version
and not requires_trust_remote_code
)
if _install_only_upgrade or await asyncio.to_thread(
latest_tier_active_for, config.identifier, request.hf_token
):
effective_load_in_4bit = False
# Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
# run the security preflight for non-GGUF loads (requires_trust_remote_code was
# already resolved above for the sizing flip).
requires_security_review = False
if not is_gguf:
requires_security_review = any(
_requires_security_review_for_model(_t, request.hf_token) for _t in security_targets
)
@ -4676,6 +4752,8 @@ async def validate_model(
requires_trust_remote_code = requires_trust_remote_code,
requires_security_review = requires_security_review,
context_length = context_length,
requires_transformers_upgrade = transformers_upgrade is not None,
transformers_upgrade = transformers_upgrade,
)
except HTTPException:
@ -4720,6 +4798,217 @@ async def validate_model(
)
# studio_router only: admin action, kept off the OpenAI-compatible /v1 mount.
@studio_router.post(
"/install-latest-transformers", response_model = InstallLatestTransformersResponse
)
async def install_latest_transformers_route(
request: InstallLatestTransformersRequest, current_subject: str = Depends(get_current_subject)
):
"""
Consented install of the latest transformers release into the persistent
.venv_t5_latest sidecar.
Called after the user confirms the transformers-upgrade dialog raised by /validate
(requires_transformers_upgrade). The requested version must match the current latest
PyPI release (re-verified server-side); the sidecar then participates in routing on
this and every future start. A pip install runs off-loop, so this can take a minute.
"""
from utils.transformers_latest import install_latest_transformers
from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap
# The install stage-and-swaps .venv_t5_latest in place; a live worker would
# lazy-import from the new version mid-run, mixing incompatible modules. Gate on
# worker LIVENESS not tier (no HF token here, so tier re-resolution is unreliable
# for gated repos): training and export are refused, the chat model unloaded.
# Reserve the swap FIRST, before any await: training/export starts check this
# reservation, so raising it after the gate wait would let a worker slip in.
if not try_begin_sidecar_swap():
raise HTTPException(
status_code = 409,
detail = "A transformers installation is already in progress.",
)
# Until the installer thread takes over, this coroutine owns the reservation
# and must release it on any early exit (the 409 refusals below).
owns_reservation = True
try:
from core.export import get_export_backend
from core.training import get_training_backend
if get_training_backend().is_training_active():
raise HTTPException(
status_code = 409,
detail = (
"A training run is active. Wait for it to finish before "
"installing a new transformers version."
),
)
_export = get_export_backend()
if _export.is_export_active():
raise HTTPException(
status_code = 409,
detail = (
"An export is running. Wait for it to finish before "
"installing a new transformers version."
),
)
# A loaded (idle) export checkpoint would be torn down by the pre-swap
# cleanup; if the swap then failed, that state would be silently lost
# with no rollback signal. Make the user unload it deliberately first.
if getattr(_export, "current_checkpoint", None):
raise HTTPException(
status_code = 409,
detail = (
"An export checkpoint is loaded. Unload it from the Export "
"page before installing a new transformers version."
),
)
# In-flight streams passed the middleware already, so the lifecycle gate can't
# protect them and the swap's unload would kill them mid-stream; mirror the
# auto-switch busy check. This route is not middleware-counted and pending
# requests stay blocked in the middleware, so neither is subtracted here.
from core.inference.llama_keepwarm import (
inference_lifecycle_gate,
note_model_unloaded,
other_inference_request_count,
)
if other_inference_request_count(current_request_counted = False, include_pending = False) > 0:
raise HTTPException(
status_code = 409,
detail = (
"Another inference request is in progress. Wait for it to "
"finish before installing a new transformers version."
),
)
# Hold the lifecycle gate /load holds so no HF worker can start (or be mid-load
# with active_model_name unset) while the sidecar is swapped. Teardown runs via
# before_swap, only once the staged install succeeded: a failed pip/compat check
# must not leave the user with their model gone. GGUF stays loaded (llama-server
# never imports transformers).
backend = get_inference_backend()
export_backend = get_export_backend()
unloaded_chat = {"v": False}
def _unload_before_swap() -> None:
# Runs on the install thread, inside the gate held by _gated_install. Any
# failure raises so the previous sidecar stays untouched (a worker that did
# not tear down cleanly may still lazy-import from it). Export teardown runs
# FIRST so its failure aborts while the chat model is still loaded;
# cleanup_memory shuts the subprocess down even when its command fails, so
# judge by worker liveness, not its return value.
export_backend.cleanup_memory()
export_alive = getattr(export_backend, "is_worker_alive", None)
if callable(export_alive) and export_alive():
raise RuntimeError("Export worker still alive before the transformers swap")
active = getattr(backend, "active_model_name", None)
if active:
if not backend.unload_model(active):
# A failed unload still clears the orchestrator's model state,
# so the model is gone from the parent's view even though the
# swap aborts: report it so the client rolls back instead of
# pointing at an unloaded model.
if getattr(backend, "active_model_name", None) != active:
unloaded_chat["v"] = True
note_model_unloaded()
raise RuntimeError(f"Could not unload '{active}' before the transformers swap")
note_model_unloaded()
unloaded_chat["v"] = True
logger.info(
"Unloaded '%s' before swapping in transformers %s",
active,
request.version,
)
# A failed load can leave a live worker with no active model that
# still holds sidecar modules (and blocks the rename on Windows).
worker_alive = getattr(backend, "is_worker_alive", None)
if callable(worker_alive) and worker_alive():
# _shutdown_subprocess keeps the handle when the worker outlives SIGKILL,
# so both its False result and the liveness recheck catch a survivor
# rather than the recheck being fooled by a nulled handle.
stopped = backend._shutdown_subprocess()
if not stopped or worker_alive():
raise RuntimeError("Inference worker still alive before the transformers swap")
def _run_install() -> dict:
# Owns the reservation from here: releasing in the thread, not the route,
# keeps it held if the request is cancelled while the install still stages.
try:
return install_latest_transformers(request.version, _unload_before_swap, True)
finally:
end_sidecar_swap()
# Snapshot before waiting on the gate: a /load already holding it can
# complete meanwhile (including a same-model reload with new settings),
# and the installer must not unload a model whose successful LoadResponse
# the client is about to render. The generation counter catches reloads
# the name alone would miss.
active_before_gate = (
getattr(backend, "active_model_name", None),
getattr(backend, "load_generation", 0),
)
async def _gated_install() -> dict:
# Held by THIS task, not the request coroutine: a cancelled POST unwinding an
# `async with` here would drop the only guard /load honors mid-install.
async with inference_lifecycle_gate():
_active_now = (
getattr(backend, "active_model_name", None),
getattr(backend, "load_generation", 0),
)
if _active_now != active_before_gate:
end_sidecar_swap()
raise HTTPException(
status_code = 409,
detail = (
"A model load completed while the install was waiting. "
"Retry the install."
),
)
# Recheck under the gate: new streams bump their in-flight count while
# holding it, so once held nothing slips past (the pre-gate check is only
# a fast path and can be outlasted by a wait on a long /load).
if (
other_inference_request_count(
current_request_counted = False, include_pending = False
)
> 0
):
end_sidecar_swap()
raise HTTPException(
status_code = 409,
detail = (
"Another inference request is in progress. Wait for "
"it to finish before installing a new transformers "
"version."
),
)
return await asyncio.to_thread(_run_install)
install_task = asyncio.ensure_future(_gated_install())
owns_reservation = False
# shield: a cancelled request stops waiting, but the installer runs to
# completion (holding the gate) instead of being torn down mid-swap.
result = await asyncio.shield(install_task)
finally:
if owns_reservation:
end_sidecar_swap()
if not result["success"]:
if result.get("latest_version"):
# Structured failure so the dialog can update to the newer release
# and offer a retry that can actually succeed.
return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"])
if unloaded_chat["v"]:
# The chat model is already gone even though the swap failed; return a
# structured failure (not a bare 400) so the client can restore its
# model state instead of pointing at an unloaded model.
return InstallLatestTransformersResponse(**result, model_unloaded = True)
raise HTTPException(status_code = 400, detail = result["message"])
return InstallLatestTransformersResponse(**result, model_unloaded = unloaded_chat["v"])
@router.post("/unload", response_model = UnloadResponse)
async def unload_model(request: UnloadRequest, current_subject: str = Depends(get_current_subject)):
"""

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

@ -146,6 +146,16 @@ async def start_training(
# No in-process ensure_transformers_version(): the subprocess
# (worker.py) activates the correct version before importing ML libs.
# A consented latest-transformers install stage-and-swaps .venv_t5_latest;
# a worker spawned mid-swap could activate a half-replaced sidecar.
from utils.transformers_latest import is_install_in_progress
if is_install_in_progress():
raise HTTPException(
status_code = 409,
detail = ("A transformers installation is in progress. Retry when it completes."),
)
backend = get_training_backend()
# S3 dataset loading needs the optional boto3 dependency. Reject early
@ -341,6 +351,24 @@ async def start_training(
"s3_config": request.s3_config.model_dump() if request.s3_config else None,
}
# Latest-sidecar models size and train 16-bit (same flip as chat load):
# 4-bit is disabled for brand-new architectures, so VRAM coexistence
# checks must not underestimate against a load the worker will refuse.
if training_kwargs["load_in_4bit"]:
from utils.transformers_version import latest_tier_active_for
if await asyncio.to_thread(
latest_tier_active_for,
training_kwargs["model_name"],
training_kwargs["hf_token"] or None,
):
training_kwargs["load_in_4bit"] = False
logger.info(
"Latest-transformers sidecar active for %s - sizing and "
"training in 16-bit (4-bit is disabled for brand-new "
"architectures)",
training_kwargs["model_name"],
)
# Training page has no trust_remote_code toggle, so honor the YAML default
# -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a
# local path or a name merely starting with "unsloth/".
@ -426,9 +454,16 @@ async def start_training(
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
# The hook runs only once start guards pass -> VRAM freed iff training starts.
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
)
from utils.transformers_version import SidecarSwapInProgress
try:
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
)
except SidecarSwapInProgress as exc:
# Expected loss of the race against a sidecar install: a retryable
# 409 matching the route-entry guard, not an internal error.
raise HTTPException(status_code = 409, detail = str(exc))
if not success:
progress_error = backend.trainer.training_progress.error

10
studio/backend/tests/test_chat_load_during_training.py
View file

@ -651,11 +651,19 @@ class TestLoadModelGuardIntegration(unittest.TestCase):
inf._shutdown_subprocess = MagicMock()
llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None)
llama.unload_model = MagicMock()
cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None)
cfg = SimpleNamespace(
is_gguf = False,
is_lora = False,
path = None,
base_model = None,
identifier = "unsloth/Qwen3-1.7B",
)
request = LoadRequest(model_path = "unsloth/Qwen3-1.7B")
info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"}
with (
# Pin the latest-sidecar tier check so the guard path stays offline.
patch("utils.transformers_version.latest_tier_active_for", return_value = False),
patch.object(self.route, "validate_extra_args", return_value = None),
patch.object(
self.route,

29
studio/backend/tests/test_orchestrator_unload_cancel.py
View file

@ -874,6 +874,35 @@ def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch):
assert o.models == {}
def test_load_model_aborts_when_old_worker_survives_shutdown(monkeypatch):
# A wedged worker that outlives terminate/kill makes _shutdown_subprocess return
# False. load_model must not spawn a second worker over it (double GPU allocation +
# the survivor's handle is lost); it aborts so the load can retry once it exits.
import types
from utils import transformers_version as tv
o = _bare_orchestrator()
o.active_model_name = "old"
o.models = {"old": {}}
o.loading_models = set()
monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False)
monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([0], "sel"))
monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None)
monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True)
monkeypatch.setattr(o, "_cancel_generation", lambda: None)
monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: False) # survivor
monkeypatch.setattr(
o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn over a live survivor")
)
with pytest.raises(RuntimeError, match = "did not exit"):
o.load_model(types.SimpleNamespace(identifier = "new", gguf_variant = None))
# The except path cleared the loading marker and mirrors.
assert "new" not in o.loading_models
assert o.active_model_name is None
def test_load_model_proceeds_when_not_cancelled(monkeypatch):
# Guard against a false abort: an uncancelled load keeps its marker and spawns.
o = _bare_orchestrator()

145
studio/backend/tests/test_shutdown_preserves_live_worker.py Normal file
View file

@ -0,0 +1,145 @@
# SPDX-License-Identifier: AGPL-3.0-only
"""_shutdown_subprocess returns whether the worker actually died, and preserves the
live handle when it survives terminate/kill.
A GPU worker wedged in an uninterruptible CUDA syscall can outlive SIGKILL. If shutdown
nulled its handle anyway, is_worker_alive() would report False and the pre-swap liveness
guard would let the destructive .venv_t5_latest rename proceed while a live worker still
holds sidecar transformers modules (breaking the rename on Windows). The methods must keep
the handle and return False so callers can refuse the swap.
"""
import pytest
from core.export.orchestrator import ExportOrchestrator
from core.inference.orchestrator import InferenceOrchestrator
class _FakeProc:
"""A subprocess handle that dies only on the requested step (or never)."""
def __init__(self, dies_on = None):
self._alive = True
self._dies_on = dies_on # None | "join" | "terminate" | "kill"
def is_alive(self):
return self._alive
def join(self, timeout = None):
if self._dies_on == "join":
self._alive = False
def terminate(self):
if self._dies_on == "terminate":
self._alive = False
def kill(self):
if self._dies_on == "kill":
self._alive = False
def _bare_inference():
o = InferenceOrchestrator.__new__(InferenceOrchestrator)
o._stop_dispatcher = lambda: None
o._cancel_generation = lambda: None
o._drain_queue = lambda: []
class _Q:
def put(self, *a, **k):
pass
o._cmd_queue = _Q()
o._resp_queue = _Q()
o._cancel_event = None
o._drain_event = None
return o
def _bare_export():
o = ExportOrchestrator.__new__(ExportOrchestrator)
o._drain_queue = lambda: []
class _Q:
def put(self, *a, **k):
pass
o._cmd_queue = _Q()
o._resp_queue = _Q()
return o
@pytest.fixture(autouse = True)
def _no_sleep(monkeypatch):
# _shutdown_subprocess sleeps 0.5s after cancelling; keep the tests instant.
import core.inference.orchestrator as inf_mod
monkeypatch.setattr(inf_mod.time, "sleep", lambda *_a, **_k: None)
class TestInferenceShutdownReturn:
def test_worker_that_dies_returns_true_and_clears_handle(self):
o = _bare_inference()
o._proc = _FakeProc(dies_on = "terminate")
assert o._shutdown_subprocess(timeout = 0.01) is True
assert o._proc is None
assert o.is_worker_alive() is False
def test_survivor_returns_false_and_keeps_handle(self):
o = _bare_inference()
o._proc = _FakeProc(dies_on = None) # outlives terminate AND kill
assert o._shutdown_subprocess(timeout = 0.01) is False
assert o._proc is not None
# is_worker_alive stays truthful, so the pre-swap guard can refuse the swap.
assert o.is_worker_alive() is True
def test_already_dead_returns_true(self):
o = _bare_inference()
o._proc = _FakeProc(dies_on = "join")
o._proc._alive = False
assert o._shutdown_subprocess(timeout = 0.01) is True
assert o._proc is None
class TestExportShutdownReturn:
def test_worker_that_dies_returns_true_and_clears_handle(self):
o = _bare_export()
o._proc = _FakeProc(dies_on = "terminate")
assert o._shutdown_subprocess(timeout = 0.01) is True
assert o._proc is None
assert o.is_worker_alive() is False
def test_survivor_returns_false_and_keeps_handle(self):
o = _bare_export()
o._proc = _FakeProc(dies_on = None)
assert o._shutdown_subprocess(timeout = 0.01) is False
assert o._proc is not None
assert o.is_worker_alive() is True
class TestSpawnPathsHonorFailedShutdown:
"""A fresh-load path must not spawn a second worker over one that outlived
terminate/kill: the survivor still holds GPU memory and its handle would be lost."""
def test_export_load_checkpoint_aborts_when_worker_survives(self, monkeypatch):
import threading
import utils.transformers_version as tv
o = ExportOrchestrator.__new__(ExportOrchestrator)
o._lock = threading.RLock()
o._proc = _FakeProc(dies_on = None) # survivor
o.clear_logs = lambda: None
o._cancel_requested = False
o._active_op_kind = None
o._export_active = False
o._ensure_subprocess_alive = lambda: True
o._shutdown_subprocess = lambda *a, **k: False
o._spawn_subprocess = lambda cfg: pytest.fail("must not spawn over a live survivor")
o._record_op_finished = lambda *a, **k: None
monkeypatch.setattr(tv, "sidecar_swap_in_progress", lambda: False)
ok, msg = o.load_checkpoint(checkpoint_path = "ckpt")
assert ok is False
assert "did not exit" in msg
# The finally cleared the op flags even though we returned early.
assert o._export_active is False

1099
studio/backend/tests/test_transformers_latest.py Normal file
View file

File diff suppressed because it is too large Load diff

646
studio/backend/tests/test_transformers_version.py
View file

@ -2550,3 +2550,649 @@ class TestHfEndpointUnreachable:
t0 = time.time()
result = hf_endpoint_unreachable(timeout = 2)
assert result is True and (time.time() - t0) < 6.0
class TestLatestTierActiveFor:
"""latest_tier_active_for: the 16-bit guard for the consented latest sidecar."""
@staticmethod
def _pin(
monkeypatch,
tv,
version = "5.13.1",
):
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: version)
monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: None)
def test_true_when_tier_latest(self, monkeypatch):
import utils.transformers_version as tv
self._pin(monkeypatch, tv)
monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, **k: "latest")
assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is True
def test_false_for_fixed_tiers(self, monkeypatch):
import utils.transformers_version as tv
self._pin(monkeypatch, tv)
for tier in ("default", "530", "550", "510"):
monkeypatch.setattr(tv, "get_transformers_tier", lambda *a, _t = tier, **k: _t)
assert tv.latest_tier_active_for("some/model") is False
def test_false_without_pin_and_no_resolution(self, monkeypatch):
"""No sidecar pin returns False before any tier or network resolution."""
import utils.transformers_version as tv
def _boom(*a, **k):
raise AssertionError("must not resolve without a pin")
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None)
monkeypatch.setattr(tv, "_remote_lora_base", _boom)
monkeypatch.setattr(tv, "get_transformers_tier", _boom)
assert tv.latest_tier_active_for("Zyphra/ZAYA1-8B") is False
def test_never_raises(self, monkeypatch):
import utils.transformers_version as tv
def _boom(*a, **k):
raise RuntimeError("tier resolution exploded")
self._pin(monkeypatch, tv)
monkeypatch.setattr(tv, "get_transformers_tier", _boom)
assert tv.latest_tier_active_for("some/model") is False
def test_remote_lora_base_is_resolved(self, monkeypatch):
"""A remote adapter is judged by its base model, like worker activation."""
import utils.transformers_version as tv
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.13.1")
monkeypatch.setattr(tv, "_remote_lora_base", lambda name, hf_token = None: "Zyphra/ZAYA1-8B")
tiers = {"Zyphra/ZAYA1-8B": "latest"}
monkeypatch.setattr(
tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default")
)
assert tv.latest_tier_active_for("someuser/zaya-lora") is True
def test_local_checkpoint_config_upgrades(self, monkeypatch, tmp_path):
"""An adapter dir with its own config.json merges tiers like activation does."""
import utils.transformers_version as tv
adapter = tmp_path / "ckpt"
adapter.mkdir()
(adapter / "adapter_config.json").write_text("{}")
(adapter / "adapter_model.safetensors").write_text("x")
(adapter / "config.json").write_text("{}")
self._pin(monkeypatch, tv)
monkeypatch.setattr(tv, "_resolve_base_model", lambda name: "base/model")
tiers = {"base/model": "default", str(adapter): "latest"}
monkeypatch.setattr(
tv, "get_transformers_tier", lambda name, *a, **k: tiers.get(name, "default")
)
assert tv.latest_tier_active_for(str(adapter)) is True
class TestLatestTierForces16Bit:
"""The inference worker and load route refuse bnb 4-bit on the latest sidecar."""
def _read(self, rel):
backend_dir = Path(__file__).resolve().parent.parent
return (backend_dir / rel).read_text()
def test_worker_guard_present(self):
src = self._read("core/inference/worker.py")
assert "latest_tier_active_for" in src, (
"core/inference/worker.py must force load_in_4bit=False when "
"latest_tier_active_for(model) is true: transformers' grouped-MoE "
"kernels crash on bnb-quantized expert weights for brand-new "
"architectures."
)
def test_route_guard_present(self):
src = self._read("routes/inference.py")
assert "latest_tier_active_for" in src, (
"routes/inference.py must size the VRAM guard with the same 16-bit "
"flip the worker applies for latest-sidecar models."
)
def test_validate_route_mirrors_16bit_flip(self):
# Without the same flip, /validate sizes 4-bit and /load then 409s.
src = self._read("routes/inference.py")
body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0]
assert "latest_tier_active_for" in body, (
"validate_model must apply the latest-sidecar 16-bit flip before "
"_guard_chat_load_against_training so /validate and /load agree."
)
# First-time loads have no pin yet, so an installable upgrade must also size 16-bit.
assert body.index("check_upgrade_for_model") < body.index(
"_guard_chat_load_against_training"
), "the upgrade check must run before the training guard"
assert (
"supported_in_pypi" in body.split("_guard_chat_load_against_training")[0]
), "an installable upgrade must force 16-bit sizing for the guard"
def test_validate_offered_upgrade_preserves_custom_code_4bit(self):
# A merely-offered (not installed) upgrade must NOT force 16-bit sizing when the
# model has a custom-code (auto_map) fallback: /load loads it 4-bit without the
# install, and the install route refuses during active training, so 16-bit sizing
# here would 409 the only viable 4-bit path.
src = self._read("routes/inference.py")
body = src.split("async def validate_model", 1)[1].split("\nasync def ", 1)[0]
flip = body.split("Mirror /load's latest-sidecar 16-bit flip", 1)[1].split(
"_guard_chat_load_against_training", 1
)[0]
assert "not requires_trust_remote_code" in flip, (
"the offered-upgrade 16-bit flip must be gated on the absence of a custom-code "
"fallback so /validate does not 409 a 4-bit load /load would allow"
)
# requires_trust_remote_code must be resolved before the flip consumes it.
assert body.index("requires_trust_remote_code = any(") < body.index(
"not requires_trust_remote_code"
)
def test_install_route_guards_active_latest_workers(self):
# Stage-and-swap replaces .venv_t5_latest in place, so a live worker on the
# old sidecar would lazy-import files from the new version.
src = self._read("routes/inference.py")
body = src.split("async def install_latest_transformers_route", 1)[1].split(
"\nasync def ", 1
)[0]
assert (
"is_training_active" in body
and "is_export_active" in body
and "inference_lifecycle_gate" in body
), (
"install_latest_transformers_route must refuse while training or export "
"runs, and hold the lifecycle gate while unloading the chat model and "
"swapping the sidecar."
)
# The unload (via before_swap so failed installs keep the model), the export-worker
# teardown, and the install must all sit INSIDE the gate so no /load interleaves.
assert "unload_model(active)" in body
assert "cleanup_memory()" in body
# Export teardown precedes the chat unload so its failure aborts with the model still loaded.
assert body.index("cleanup_memory()") < body.index("unload_model(active)")
assert "install_latest_transformers(" in body and "_unload_before_swap" in body
# The gate must be owned by the shielded task, not the request coroutine: a cancelled
# POST unwinding an async-with would release the only guard /load honors mid-install.
gated_task = body.split("async def _gated_install", 1)[1]
assert "inference_lifecycle_gate():" in gated_task
assert "asyncio.to_thread(_run_install)" in gated_task
# The reservation must be taken BEFORE the (awaitable) gate wait, or a
# training/export start could slip in while this request queues on the gate.
assert body.index("try_begin_sidecar_swap()") < body.index(
"inference_lifecycle_gate():"
), "the swap reservation must be raised before waiting on the lifecycle gate"
# A failed teardown must abort the swap (raise), not fall through to it.
assert body.count("raise RuntimeError") >= 3, (
"export, chat-unload, and idle-worker teardown failures must raise so "
"the staged install never swaps under a live worker"
)
# The installer thread owns (and releases) the reservation, shielded from
# request cancellation, so a cancelled POST cannot unlock a live swap.
assert "asyncio.shield" in body and "end_sidecar_swap()" in body
# In-flight generation streams predate the gate; the route refuses rather than kill them
# via the before_swap unload. The count is rechecked UNDER the gate, since a wait on a
# long /load outlasts the pre-gate fast path and streams take this same gate.
assert "other_inference_request_count" in body
gated_task = body.split("async def _gated_install", 1)[1]
assert "other_inference_request_count" in gated_task
def test_start_routes_refuse_during_install(self):
# A worker spawned mid-swap could activate a half-replaced sidecar.
training = self._read("routes/training.py")
start = training.split("async def start_training", 1)[1].split("\nasync def ", 1)[0]
assert (
"is_install_in_progress" in start
), "training /start must refuse while a transformers install is in progress"
export = self._read("routes/export.py")
helper = export.split("def _ensure_export_supported", 1)[1].split("\ndef ", 1)[0]
assert (
"is_install_in_progress" in helper
), "mutating export routes must refuse while a transformers install is in progress"
def test_spawn_sites_recheck_reservation(self):
# The route-level guards are one-shot; validation between them and the
# actual spawn can outlast an install's start, so the spawn itself rechecks.
training = self._read("core/training/training.py")
assert (
training.count("sidecar_swap_in_progress()") >= 2
), "both training spawn sites must recheck the sidecar swap reservation"
export = self._read("core/export/orchestrator.py")
spawn = export.split("def _spawn_subprocess", 1)[1].split("\n def ", 1)[0]
assert (
"sidecar_swap_kind()" in spawn
), "the export subprocess spawn must recheck the sidecar swap reservation"
# Training marks the spawn active BEFORE its recheck, so either side sees the other:
# is_training_active covers the window between proc.start() and the _proc assignment.
assert training.index("self._spawn_in_progress = True") < training.index(
"if sidecar_swap_in_progress():"
)
active = training.split("def is_training_active", 1)[1].split("\n def ", 1)[0]
assert "_spawn_in_progress" in active
# Export load-checkpoint refuses BEFORE tearing down the old worker, so a
# lost race against an install keeps the loaded checkpoint (no bare 500).
loadck = export.split("def load_checkpoint", 1)[1].split("\n def ", 1)[0]
assert loadck.index("sidecar_swap_in_progress()") < loadck.index("_shutdown_subprocess()")
# The training handshake precedes the VRAM-freeing before_spawn hook, so
# losing the race never tears down chat/export for a run that won't spawn.
assert training.index("self._spawn_in_progress = True") < training.index("before_spawn()")
# The spawn-time export check is op-aware for installs (the install side
# aborts on is_export_active) but always refuses for repairs, which have
# no such abort and can be rebuilding the sidecar right now.
assert (
'_swap_kind == "repair" or (_swap_kind is not None and not self._export_active)'
in spawn
)
class TestSidecarSwapReservation:
"""The lazy repair takes the same reservation the install route and worker starts use."""
def _repair_setup(self, monkeypatch, tmp_path):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
monkeypatch.setattr(
tv,
"_latest_pin_data",
lambda: {
"version": "5.99.0",
"packages": ["transformers==5.99.0"],
},
)
monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: False)
monkeypatch.setattr(tv, "_env_offline", lambda: False)
return tv
def test_repair_holds_reservation_during_swap(self, monkeypatch, tmp_path):
tv = self._repair_setup(monkeypatch, tmp_path)
seen = {}
def _fake_swap(
version,
packages,
before_swap = None,
):
seen["active_during_swap"] = tv.sidecar_swap_in_progress()
return True
monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _fake_swap)
assert tv._ensure_venv_t5_latest_exists() is True
assert seen["active_during_swap"] is True
assert tv.sidecar_swap_in_progress() is False
def test_foreign_process_lock_file_visible(self, monkeypatch, tmp_path):
"""A repair in a LIVE worker subprocess is seen (via the lock file) by this
process, and its lock is never broken while the owner is alive."""
import os
import time
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
lock = tv._swap_lock_path()
lock.parent.mkdir(parents = True, exist_ok = True)
# A live owner (this process): visible and never reclaimed, even once aged past
# the cutoff -- a slow but live pip install must keep its lock.
lock.write_text('{"pid": %d}' % os.getpid())
assert tv.sidecar_swap_in_progress() is True
assert tv.try_begin_sidecar_swap() is False
old_ts = time.time() - 3 * 60 * 60
os.utime(lock, (old_ts, old_ts))
assert tv.sidecar_swap_in_progress() is True
assert tv.try_begin_sidecar_swap() is False
def test_dead_owner_lock_reclaimed_promptly(self, monkeypatch, tmp_path):
"""A fresh lock whose recorded owner is dead is reclaimed at once, not after the
long cutoff: a crash mid-install must not wedge loads/training/export for hours."""
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
lock = tv._swap_lock_path()
lock.parent.mkdir(parents = True, exist_ok = True)
# 999999 is not a live PID: a fresh dead-owner lock is immediately stale.
lock.write_text('{"pid": 999999, "kind": "install"}')
assert tv._pid_alive(999999) is False
assert tv.sidecar_swap_in_progress() is False
assert tv.try_begin_sidecar_swap() is True
try:
assert lock.is_file()
finally:
tv.end_sidecar_swap()
assert not lock.exists()
def test_unreadable_pid_lock_uses_age_cutoff(self, monkeypatch, tmp_path):
"""A lock with no readable owner PID (mid create-before-write, or corrupt) is not
reclaimed while fresh -- only after the long cutoff -- so a lock a live owner just
created is not stolen before its PID lands."""
import os
import time
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
lock = tv._swap_lock_path()
lock.parent.mkdir(parents = True, exist_ok = True)
lock.write_text("") # created but metadata not yet written
assert tv.sidecar_swap_in_progress() is True
old_ts = time.time() - (tv._SWAP_LOCK_STALE_SECS + 60)
os.utime(lock, (old_ts, old_ts))
assert tv.sidecar_swap_in_progress() is False
def test_repair_refused_while_install_holds_reservation(self, monkeypatch, tmp_path):
tv = self._repair_setup(monkeypatch, tmp_path)
def _must_not_run(*a, **k):
raise AssertionError("repair must not swap while an install is in progress")
monkeypatch.setattr(tv, "_stage_and_swap_latest_venv", _must_not_run)
assert tv.try_begin_sidecar_swap() is True
try:
assert tv._ensure_venv_t5_latest_exists() is False
finally:
tv.end_sidecar_swap()
class TestRecoverStrandedSidecar:
"""A swap whose activation rename AND rollback both fail strands the previous sidecar
at .old with no live dir (its pin marker went with it). Reading the pin self-heals it,
but never while a swap legitimately holds the reservation."""
def _setup(self, monkeypatch, tmp_path):
import utils.transformers_version as tv
live = str(tmp_path / "venv_t5_latest")
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", live)
# Stranded state: live gone, previous sidecar (with its marker) sits at .old.
retired = Path(live + ".old")
retired.mkdir(parents = True)
(retired / tv._LATEST_PIN_MARKER).write_text(
'{"version": "5.99.0", "packages": ["transformers==5.99.0"]}'
)
return tv, Path(live), retired
def test_stranded_old_recovered_on_pin_read(self, monkeypatch, tmp_path):
tv, live, retired = self._setup(monkeypatch, tmp_path)
data = tv._latest_pin_data()
assert live.is_dir()
assert not retired.exists()
assert data is not None and data["version"] == "5.99.0"
def test_stranded_recovery_skipped_during_swap(self, monkeypatch, tmp_path):
tv, live, retired = self._setup(monkeypatch, tmp_path)
assert tv.try_begin_sidecar_swap() is True
try:
# A swap holds the reservation and may be mid-rename; do not race it.
assert tv._latest_pin_data() is None
assert not live.exists()
assert retired.is_dir()
finally:
tv.end_sidecar_swap()
# Once the swap is done, the next pin read recovers the stranded sidecar.
assert tv._latest_pin_data() is not None
assert live.is_dir()
class TestCachedLatestMappingRevalidated:
"""A cached 'latest' mapping is dropped and re-resolved when the sidecar since broke
in-process, so routing self-heals instead of trusting a mapping parsed from a sidecar
that no longer exists (which would keep routing latest-only models to a broken tier)."""
def test_broken_sidecar_drops_cached_latest_mapping(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})})
monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: False)
seen = {"n": 0}
def _fake_overlay(tier):
seen["n"] += 1
return None # broken/unavailable -> empty, uncached
monkeypatch.setattr(tv, "_overlay_transformers_dir", _fake_overlay)
assert tv._config_model_types("latest") == frozenset()
assert seen["n"] == 1 # re-resolved, not served from the stale cache
assert "latest" not in tv._config_mapping_cache
def test_intact_sidecar_serves_cached_latest_mapping(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})})
monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True)
monkeypatch.setattr(
tv,
"_overlay_transformers_dir",
lambda tier: pytest.fail("intact sidecar must serve the cache without re-resolving"),
)
assert tv._config_model_types("latest") == frozenset({"brandnew"})
def test_non_latest_cache_not_revalidated(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_config_mapping_cache", {"530": frozenset({"gemma3"})})
monkeypatch.setattr(
tv,
"_latest_sidecar_intact",
lambda: pytest.fail("non-latest tiers must not pay the sidecar-intact check"),
)
assert tv._config_model_types("530") == frozenset({"gemma3"})
def test_deleted_pin_drops_cached_latest_mapping(self, monkeypatch, tmp_path):
# A pin marker deleted after the mapping was cached makes _latest_pin_data None;
# the cache must be dropped (not trusted), so routing re-resolves to no latest tier
# rather than routing to a latest tier that then fails worker activation.
import utils.transformers_version as tv
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(tmp_path / "venv_t5_latest"))
monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False)
monkeypatch.setattr(tv, "_config_mapping_cache", {"latest": frozenset({"brandnew"})})
# No pin marker on disk -> _latest_pin_data() is None -> not intact.
assert tv._latest_sidecar_intact() is False
assert tv._config_model_types("latest") == frozenset()
assert "latest" not in tv._config_mapping_cache
class TestOverlayRepairsIncompleteSidecar:
"""Routing self-heals a pinned latest sidecar that is present but incomplete,
not only one whose transformers/ dir vanished: workers refuse parent-only
repairs, so a sidecar missing a pinned package would fail every load."""
def _setup(self, monkeypatch, tmp_path, valid):
import utils.transformers_version as tv
live = tmp_path / "venv_t5_latest"
(live / "transformers").mkdir(parents = True)
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live))
monkeypatch.setattr(tv, "_latest_tier_disabled", lambda: False)
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0")
monkeypatch.setattr(
tv,
"_latest_pin_data",
lambda: {"version": "5.99.0", "packages": ["transformers==5.99.0", "tiktoken"]},
)
monkeypatch.setattr(tv, "_venv_dir_is_valid", lambda d, p: valid)
monkeypatch.setattr(tv, "_latest_repair_failed_at", 0.0)
return tv
def test_incomplete_sidecar_triggers_repair(self, monkeypatch, tmp_path):
tv = self._setup(monkeypatch, tmp_path, valid = False)
called = {"n": 0}
def _fake_repair():
called["n"] += 1
return True
monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair)
src = tv._overlay_transformers_dir("latest")
assert called["n"] == 1
assert src == str(tmp_path / "venv_t5_latest" / "transformers")
def test_intact_sidecar_skips_repair(self, monkeypatch, tmp_path):
tv = self._setup(monkeypatch, tmp_path, valid = True)
def _must_not_run():
raise AssertionError("intact sidecar must not trigger a repair")
monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _must_not_run)
assert tv._overlay_transformers_dir("latest") == str(
tmp_path / "venv_t5_latest" / "transformers"
)
def test_failed_repair_backs_off(self, monkeypatch, tmp_path):
tv = self._setup(monkeypatch, tmp_path, valid = False)
called = {"n": 0}
def _fake_repair():
called["n"] += 1
return False
monkeypatch.setattr(tv, "_ensure_venv_t5_latest_exists", _fake_repair)
# A failed repair must not route through the broken sidecar, neither on
# the failing attempt nor while the backoff suppresses the next attempt.
assert tv._overlay_transformers_dir("latest") is None
assert tv._overlay_transformers_dir("latest") is None
assert called["n"] == 1
class TestStageAndSwapBeforeSwap:
"""before_swap fires only when the staged install succeeded and the swap is next."""
def _setup(self, monkeypatch, tmp_path, build_ok):
import utils.transformers_version as tv
live = tmp_path / "venv_latest"
monkeypatch.setattr(tv, "_VENV_T5_LATEST_DIR", str(live))
def _fake_build(target, packages, label):
if build_ok:
Path(target).mkdir(parents = True, exist_ok = True)
return build_ok
monkeypatch.setattr(tv, "_ensure_venv_dir", _fake_build)
return tv, live
def test_called_after_successful_staging(self, monkeypatch, tmp_path):
tv, live = self._setup(monkeypatch, tmp_path, build_ok = True)
calls = []
assert tv._stage_and_swap_latest_venv(
"5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1)
)
assert calls == [1] and live.is_dir()
def test_not_called_when_staging_fails(self, monkeypatch, tmp_path):
tv, live = self._setup(monkeypatch, tmp_path, build_ok = False)
calls = []
assert not tv._stage_and_swap_latest_venv(
"5.99.0", ("transformers==5.99.0",), before_swap = lambda: calls.append(1)
)
assert calls == [] and not live.exists()
def test_failure_in_before_swap_keeps_previous_sidecar(self, monkeypatch, tmp_path):
tv, live = self._setup(monkeypatch, tmp_path, build_ok = True)
live.mkdir()
(live / "sentinel").write_text("old")
def _boom():
raise RuntimeError("worker teardown failed")
assert not tv._stage_and_swap_latest_venv(
"5.99.0", ("transformers==5.99.0",), before_swap = _boom
)
assert (live / "sentinel").read_text() == "old"
class TestKillSwitchBeatsMappingCache:
def test_cached_latest_probe_ignored_when_disabled(self, monkeypatch):
import utils.transformers_version as tv
key = tv._probe_cache_key("some/model")
monkeypatch.setitem(tv._probe_tier_cache, key, "latest")
monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1")
# With the switch set, the cached latest entry must not short-circuit;
# the probe re-resolves against the non-latest order (stub it to 530).
monkeypatch.setattr(tv, "_probe_tier_venvs", lambda: {})
monkeypatch.setattr(tv, "_probe_tier_order", lambda: ())
assert tv._probe_tier("some/model", None, "test") != "latest"
# Cached non-latest entries and the unset switch still short-circuit.
monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS")
assert tv._probe_tier("some/model", None, "test") == "latest"
def test_cached_latest_mapping_ignored_when_disabled(self, monkeypatch):
import utils.transformers_version as tv
monkeypatch.setitem(tv._config_mapping_cache, "latest", frozenset({"brandnew"}))
# The cache is trusted only when the sidecar is intact; hold it intact so this
# test isolates the kill switch, not the sidecar-revalidation path.
monkeypatch.setattr(tv, "_latest_sidecar_intact", lambda: True)
monkeypatch.setenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "1")
assert tv._config_model_types("latest") == frozenset()
monkeypatch.delenv("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS")
assert tv._config_model_types("latest") == frozenset({"brandnew"})
class TestRaiseTierForNested:
"""_raise_tier_for_nested: a wrapper's nested model_type can raise a fast-path tier."""
def _patch_types(self, monkeypatch, per_tier):
import utils.transformers_version as tv
monkeypatch.setattr(
tv, "_config_model_types", lambda tier: frozenset(per_tier.get(tier, ()))
)
def test_nested_latest_only_type_raises(self, monkeypatch):
import utils.transformers_version as tv
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}})
cfg = {"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}}
assert tv._raise_tier_for_nested(cfg, "550") == "latest"
def test_never_lowers_a_fast_path_tier(self, monkeypatch):
import utils.transformers_version as tv
# Mapping alone would say 530, but the fast path (e.g. a name override) said 550.
self._patch_types(monkeypatch, {"530": {"qwen3_5"}, "550": {"qwen3_5"}})
assert tv._raise_tier_for_nested({"model_type": "qwen3_5"}, "550") == "550"
def test_no_config_keeps_tier(self):
import utils.transformers_version as tv
assert tv._raise_tier_for_nested(None, "550") == "550"
def test_unknown_nested_type_never_vetoes(self, monkeypatch):
import utils.transformers_version as tv
# A nested type unknown everywhere (not even latest) keeps the fast path.
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4"}})
cfg = {"model_type": "gemma4", "text_config": {"model_type": "unreleased"}}
assert tv._raise_tier_for_nested(cfg, "550") == "550"
def test_name_fast_path_folds_when_latest_pinned(self, monkeypatch):
"""A fixed-tier name match with a latest-only model_type routes to latest
once the sidecar is pinned; without a pin the name tier stands (no I/O)."""
import utils.transformers_version as tv
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"brandnew_arch"}})
monkeypatch.setattr(tv, "_tier_from_name", lambda name: ("550", "gemma-4"))
monkeypatch.setattr(
tv, "_load_config_json", lambda name, tok = None: {"model_type": "brandnew_arch"}
)
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: "5.99.0")
assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "latest"
monkeypatch.setattr(tv, "latest_venv_pinned_version", lambda: None)
monkeypatch.setattr(
tv,
"_load_config_json",
lambda name, tok = None: (_ for _ in ()).throw(AssertionError("no I/O without a pin")),
)
assert tv.get_transformers_tier("org/gemma-4-new", probe = False) == "550"
def test_fast_path_folds_nested_tier(self, monkeypatch, tmp_path):
"""End to end: a local wrapper config on a fixed fast path routes to latest
when its nested type only exists in the installed latest sidecar."""
import utils.transformers_version as tv
ckpt = tmp_path / "wrapper"
ckpt.mkdir()
(ckpt / "config.json").write_text(
json.dumps({"model_type": "gemma4", "text_config": {"model_type": "brandnew_arch"}})
)
self._patch_types(monkeypatch, {"550": {"gemma4"}, "latest": {"gemma4", "brandnew_arch"}})
monkeypatch.setattr(tv, "_config_needs_510", lambda cfg: False)
monkeypatch.setattr(tv, "_config_needs_550", lambda cfg: True)
assert tv.get_transformers_tier(str(ckpt), probe = False) == "latest"

42
studio/backend/utils/models/model_config.py
View file

@ -698,6 +698,25 @@ if backend_dir not in sys.path:
try:
from transformers import AutoConfig
# Union the ACTIVE sidecar's registry into the inlined parent-process sets
# so architectures only the sidecar knows still classify correctly.
try:
from transformers.models.auto import modeling_auto as _ma
for _attr in ("MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES",
"MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES"):
_d = dict(getattr(_ma, _attr, None) or {})
_VLM_MODEL_TYPES |= set(_d)
_VLM_CLASS_NAMES |= set(_d.values())
for _attr in ("MODEL_FOR_CTC_MAPPING_NAMES",
"MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES",
"MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES",
"MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES",
"MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES"):
_AUDIO_ONLY_MODEL_TYPES |= set(dict(getattr(_ma, _attr, None) or {}))
except Exception:
pass
# Capability detection never executes model repo code.
kwargs = {"trust_remote_code": False}
if token:
@ -727,13 +746,23 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
"""
token_arg = hf_token or ""
# Latest-only architectures need the latest sidecar for AutoConfig;
# other tiers keep the 5.5 sidecar.
sidecar_dir = _VENV_T5_DIR
try:
from utils.transformers_version import _VENV_T5_LATEST_DIR, get_transformers_tier
if get_transformers_tier(model_name, hf_token, probe = False) == "latest":
sidecar_dir = _VENV_T5_LATEST_DIR
except Exception:
pass
try:
result = subprocess.run(
[
sys.executable,
"-c",
_VISION_CHECK_SCRIPT,
_VENV_T5_DIR,
sidecar_dir,
_BACKEND_DIR,
model_name,
token_arg,
@ -876,6 +905,17 @@ def _is_vision_model_uncached(
model_name, hf_token = hf_token, local_files_only = local_files_only
)
if raw is not None:
if raw is False and not local_files_only:
# Raw heuristics predate latest-only architectures; on the latest tier,
# trust that sidecar's AutoConfig probe over the heuristic False. An
# inconclusive probe (sidecar mid-repair, timeout) is transient: return
# None so the heuristic False is not cached and the model is re-probed.
try:
from utils.transformers_version import get_transformers_tier
if get_transformers_tier(model_name, hf_token, probe = False) == "latest":
return _is_vision_model_subprocess(model_name, hf_token = hf_token)
except Exception:
pass
return raw
# Raw read failed transiently: fall back to AutoConfig (remote code DISABLED), via a

607
studio/backend/utils/transformers_latest.py Normal file
View file

@ -0,0 +1,607 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Latest-transformers support check for brand-new model architectures.
When a model's ``model_type`` is absent from every installed transformers overlay
(base 4.57.x plus the .venv_t5_530/550/510 sidecars and, if provisioned, .venv_t5_latest),
Studio cannot load it today. This module answers, without authentication, code execution,
or trust_remote_code:
1. Does the LATEST transformers release on PyPI ship this ``model_type``?
2. Does transformers ``main`` on GitHub ship it (dev-only, not yet installable)?
Sources (all unauthenticated; raw.githubusercontent.com is not API rate-limited and
api.github.com is deliberately never used):
- https://pypi.org/pypi/transformers/json -> latest release version
- https://raw.githubusercontent.com/huggingface/transformers/{ref}/src/transformers/
models/auto/configuration_auto.py + auto_mappings.py -> CONFIG_MAPPING_NAMES
The fetched sources are parsed with the same AST extractor the static router uses
(:func:`utils.transformers_version._model_types_from_source`), so the remote answer is
computed exactly like the local overlay answer.
Results are cached in memory and in a small JSON snapshot under ``studio_root()/cache``
(ttl ~1 day) so repeated tier resolutions never re-fetch; failures are backed off in
memory. Every fetch is bounded (<=5s, one retry), so a hung network cannot block model
loading. Fully offline-safe: offline env vars or the kill switch
``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS=1`` make every check return None (current
behavior preserved).
The consented install path (:func:`install_latest_transformers`) provisions the
persistent ``.venv_t5_latest`` sidecar via
:func:`utils.transformers_version.ensure_latest_transformers_venv`.
"""
import json
import os
import threading
import time
from pathlib import Path
from loggers import get_logger
from utils.paths.storage_roots import studio_root as _studio_root
from utils.transformers_version import (
_env_offline,
_load_config_json,
_model_types_from_source,
_tier_from_config_mapping,
_config_model_types,
_NESTED_CONFIG_KEYS,
_TIER_RANK,
_model_types_from_config,
_TRANSFORMERS_510_MODEL_TYPES,
_TRANSFORMERS_530_MODEL_TYPES,
_TRANSFORMERS_550_MODEL_TYPES,
ensure_latest_transformers_venv,
latest_venv_pinned_version,
)
logger = get_logger(__name__)
_PYPI_JSON_URL = "https://pypi.org/pypi/transformers/json"
_RAW_URL = (
"https://raw.githubusercontent.com/huggingface/transformers/{ref}"
"/src/transformers/models/auto/{name}"
)
_AUTO_FILES = ("configuration_auto.py", "auto_mappings.py")
_FETCH_TIMEOUT_SECONDS = 5.0
_FETCH_RETRIES = 1
_CACHE_TTL_SECONDS = 24 * 60 * 60
_FAILURE_BACKOFF_SECONDS = 300
_CACHE_FILE_NAME = "transformers_latest_check.json"
_SNAPSHOT_SCHEMA = 1
# Snapshot: {"schema", "fetched_at", "pypi_version", "pypi_model_types", "main_model_types"}.
# Install-in-progress state lives in utils.transformers_version (the sidecar swap reservation).
_lock = threading.Lock()
_memory_snapshot: dict | None = None
_last_failure_at: float = 0.0
_is_fetching: bool = False
_TRUE_VALUES = {"1", "true", "yes", "on"}
def _disabled() -> bool:
"""True if the operator disabled the latest-transformers check entirely."""
return (
os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in _TRUE_VALUES
)
def _cache_file() -> Path:
return _studio_root() / "cache" / _CACHE_FILE_NAME
# Sentinel for HTTP 404 (absent at ref), distinct from transient failures.
_FETCH_MISSING = "__unsloth_fetch_missing__"
def _fetch_text(url: str) -> str | None:
"""GET *url* with a bounded timeout and one retry; None on any failure.
Returns ``_FETCH_MISSING`` (without retrying) on HTTP 404 so callers can tell
"absent at this ref" apart from "network flaked".
"""
import urllib.error
import urllib.request
for attempt in range(1 + _FETCH_RETRIES):
try:
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _FETCH_TIMEOUT_SECONDS) as resp:
return resp.read().decode("utf-8", "replace")
except urllib.error.HTTPError as exc:
if exc.code == 404:
return _FETCH_MISSING
logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc)
except Exception as exc:
logger.debug("Fetch failed (attempt %d) for %s: %s", attempt + 1, url, exc)
return None
def _fetch_latest_pypi_version() -> str | None:
"""Latest transformers release version from PyPI's unauthenticated JSON API."""
body = _fetch_text(_PYPI_JSON_URL)
if body is None or body == _FETCH_MISSING:
return None
try:
version = json.loads(body).get("info", {}).get("version")
except Exception as exc:
logger.debug("Could not parse PyPI JSON: %s", exc)
return None
return version if isinstance(version, str) and version else None
def _fetch_remote_model_types(ref: str) -> frozenset[str] | None:
"""CONFIG_MAPPING_NAMES keys at *ref* (a release tag like ``v5.12.0`` or ``main``).
Fetches configuration_auto.py plus auto_mappings.py (the 5.10+ split) from
raw.githubusercontent.com and parses them with the shared AST extractor. A file
that 404s (auto_mappings.py on pre-5.10 tags) is skipped, but a transient fetch
or parse failure of EITHER file fails the whole lookup: most model types live in
auto_mappings.py on current releases, so a partial map cached for the TTL would
make /validate skip the upgrade prompt for architectures the release does ship.
An empty result is likewise a failure so it is never cached as "supports nothing".
"""
keys: set[str] = set()
fetched_any = False
for name in _AUTO_FILES:
source = _fetch_text(_RAW_URL.format(ref = ref, name = name))
if source is None:
return None
if source == _FETCH_MISSING:
continue
fetched_any = True
try:
keys |= _model_types_from_source(source)
except Exception as exc:
logger.debug("Could not parse %s at %s: %s", name, ref, exc)
return None
if not fetched_any or not keys:
return None
return frozenset(keys)
def _load_snapshot_file() -> dict | None:
"""Persisted snapshot from disk, or None (missing/corrupt/old schema)."""
try:
with open(_cache_file(), encoding = "utf-8") as f:
data = json.load(f)
except Exception:
return None
if not isinstance(data, dict) or data.get("schema") != _SNAPSHOT_SCHEMA:
return None
if not isinstance(data.get("fetched_at"), (int, float)):
return None
if not isinstance(data.get("pypi_version"), str):
return None
for key in ("pypi_model_types", "main_model_types"):
value = data.get(key)
if not isinstance(value, list) or not all(isinstance(v, str) for v in value):
return None
return data
def _save_snapshot_file(snapshot: dict) -> None:
"""Atomic best-effort write (tmp + os.replace, Windows-safe); failures only log."""
path = _cache_file()
tmp = path.with_name(path.name + ".tmp")
try:
path.parent.mkdir(parents = True, exist_ok = True)
tmp.write_text(json.dumps(snapshot), encoding = "utf-8")
os.replace(tmp, path)
except Exception as exc:
logger.debug("Could not persist %s: %s", path, exc)
try:
tmp.unlink(missing_ok = True)
except Exception:
pass
def _snapshot_is_fresh(snapshot: dict | None) -> bool:
return (
snapshot is not None
and (time.time() - float(snapshot.get("fetched_at", 0))) < _CACHE_TTL_SECONDS
)
def _refresh_snapshot() -> dict | None:
"""Fetch a fresh snapshot from PyPI + raw.githubusercontent.com; None on failure.
The PyPI version and its tagged mapping are required; the ``main`` mapping is
best-effort (recorded as an empty list plus ``main_checked=False`` when unavailable,
so a dev-only architecture is reported as "unknown" rather than "unsupported").
"""
version = _fetch_latest_pypi_version()
if version is None:
return None
pypi_types = _fetch_remote_model_types(f"v{version}")
if pypi_types is None:
return None
main_types = _fetch_remote_model_types("main")
return {
"schema": _SNAPSHOT_SCHEMA,
"fetched_at": time.time(),
"pypi_version": version,
"pypi_model_types": sorted(pypi_types),
"main_model_types": sorted(main_types) if main_types is not None else [],
"main_checked": main_types is not None,
}
def _get_snapshot() -> dict | None:
"""Current support snapshot: memory -> disk -> network, with TTL and failure backoff.
The network refresh runs outside the lock so a slow fetch cannot stall other
threads in the ASGI pool; _is_fetching deduplicates concurrent refreshes
(losers return None, the graceful fallthrough, rather than waiting).
"""
global _memory_snapshot, _last_failure_at, _is_fetching
with _lock:
if _snapshot_is_fresh(_memory_snapshot):
return _memory_snapshot
disk = _load_snapshot_file()
if _snapshot_is_fresh(disk):
_memory_snapshot = disk
return disk
if _disabled() or _env_offline():
return None
if time.time() - _last_failure_at < _FAILURE_BACKOFF_SECONDS:
return None
if _is_fetching:
return None
_is_fetching = True
fresh = None
try:
fresh = _refresh_snapshot()
finally:
with _lock:
_is_fetching = False
if fresh is None:
_last_failure_at = time.time()
else:
_memory_snapshot = fresh
if fresh is None:
# A stale positive could offer a version PyPI no longer serves; be strict.
return None
_save_snapshot_file(fresh)
return fresh
def clear_caches() -> None:
"""Test helper: drop the in-memory snapshot, failure backoff, and busy flags."""
global _memory_snapshot, _last_failure_at, _is_fetching
with _lock:
_memory_snapshot = None
_last_failure_at = 0.0
_is_fetching = False
from utils.transformers_version import end_sidecar_swap
end_sidecar_swap()
def latest_transformers_supports(model_type: str) -> dict | None:
"""Whether the newest transformers (PyPI release and/or GitHub main) ships *model_type*.
Returns ``{"pypi_version": str, "supported_in_pypi": bool, "supported_in_main": bool}``
or None when the answer is unavailable (offline, kill switch, network failure) — the
caller must then fall through to current behavior. Cached (memory + JSON snapshot on
disk, ttl ~1 day) so repeated tier resolutions never re-fetch.
"""
if not isinstance(model_type, str) or not model_type:
return None
if _disabled() or _env_offline():
return None
snapshot = _get_snapshot()
if snapshot is None:
return None
return {
"pypi_version": snapshot["pypi_version"],
"supported_in_pypi": model_type in set(snapshot["pypi_model_types"]),
"supported_in_main": model_type in set(snapshot["main_model_types"]),
}
# model_types the hardcoded tier tables already route; never remote-check these.
def _hardcoded_model_types() -> frozenset[str]:
return frozenset(
_TRANSFORMERS_530_MODEL_TYPES
| _TRANSFORMERS_550_MODEL_TYPES
| _TRANSFORMERS_510_MODEL_TYPES
)
def check_upgrade_for_model(model_name: str, hf_token: str | None = None) -> dict | None:
"""Upgrade signal for *model_name*, or None when current routing already handles it.
The tier hook for the pre-load ``/validate`` path: fires ONLY when the model's
``model_type`` is absent from every installed overlay (and from the hardcoded tier
tables), i.e. exactly when today's load would fail with an unrecognized-architecture
error. Returns ``{"model_type", "pypi_version", "supported_in_pypi",
"supported_in_main"}`` when the newest transformers knows the type, else None.
Never raises; every network touch is bounded and cached. Offline or with the
``UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS`` kill switch it returns None immediately.
"""
try:
if _disabled() or _env_offline():
return None
cfg = _load_config_json(model_name, hf_token)
if not isinstance(cfg, dict):
return None
candidates = _model_types_from_config(cfg)
if not candidates:
return None
# Without a readable base mapping every type looks brand new; bail out.
if not _config_model_types("default"):
return None
hardcoded = _hardcoded_model_types()
missing = [
candidate
for candidate in candidates
if candidate not in hardcoded
and not any(candidate in _config_model_types(tier) for tier in _TIER_RANK)
]
if not missing:
return None
# Latest must load EVERY missing type (wrappers build nested sub-configs
# through CONFIG_MAPPING) or the load still fails.
supports = [latest_transformers_supports(candidate) for candidate in missing]
if any(
s is None or not (s["supported_in_pypi"] or s["supported_in_main"]) for s in supports
):
return None
# Offer the PyPI install only if the release ships every missing type; a
# main-only type in the mix surfaces as dev-only.
model_type = missing[0]
supported_in_pypi = all(s["supported_in_pypi"] for s in supports)
supported_in_main = all(s["supported_in_pypi"] or s["supported_in_main"] for s in supports)
logger.info(
"Model %s has model_type=%s unknown to every installed transformers "
"(latest PyPI %s: %s, main: %s)",
model_name,
model_type,
supports[0]["pypi_version"],
"supported" if supported_in_pypi else "unsupported",
"supported" if supported_in_main else "unsupported",
)
return {
"model_type": model_type,
"pypi_version": supports[0]["pypi_version"],
"supported_in_pypi": supported_in_pypi,
"supported_in_main": supported_in_main,
}
except Exception as exc:
logger.debug("Latest-transformers check failed for '%s': %s", model_name, exc)
return None
# --- Dependency compatibility preflight ------------------------------------------------------
# Sidecars install transformers --no-deps atop the base env. Before installing, compare
# requires_dist: unsatisfied shadowable deps become exact --target pins, anything else blocks.
# Safe to shadow inside the sidecar dir (pure wheels, no torch coupling).
_SHADOWABLE_DEPS = frozenset({"tokenizers", "safetensors"})
# Provided by the sidecar recipe; checked against its pin, not the base env.
_SIDECAR_PROVIDED = {"huggingface-hub": "1.8.0", "hf-xet": "1.4.2"}
# CLI-only; never imported at runtime in Studio's workers.
_IGNORED_DEPS = frozenset({"typer"})
def _canonical_dep_name(name: str) -> str:
return name.lower().replace("_", "-")
def _fetch_requires_dist(version: str) -> list[str] | None:
"""Core (marker-free, non-extra) requires_dist of transformers *version* from PyPI."""
body = _fetch_text(f"https://pypi.org/pypi/transformers/{version}/json")
if body is None or body == _FETCH_MISSING:
return None
try:
reqs = json.loads(body).get("info", {}).get("requires_dist")
except Exception:
return None
if not isinstance(reqs, list):
return None
return [r for r in reqs if isinstance(r, str)]
def _resolve_exact_version(name: str, specifier) -> str | None:
"""Newest PyPI release of *name* satisfying *specifier* (exact pin for the shadow)."""
body = _fetch_text(f"https://pypi.org/pypi/{name}/json")
if body is None or body == _FETCH_MISSING:
return None
try:
from packaging.version import InvalidVersion, Version
releases = json.loads(body).get("releases", {})
best = None
for candidate in releases:
try:
parsed = Version(candidate)
except InvalidVersion:
continue
if parsed.is_prerelease or not specifier.contains(candidate):
continue
if best is None or parsed > Version(best):
best = candidate
return best
except Exception as exc:
logger.debug("Could not resolve an exact %s version: %s", name, exc)
return None
def compat_plan(version: str) -> tuple[tuple[str, ...], list[str]]:
"""(extra exact pins to shadow-install, blocking requirement strings) for *version*.
Compares the release's core requires_dist against the running base env (the env the
workers overlay the sidecar onto). A requirement the base env satisfies needs nothing;
an unsatisfied shadowable dep becomes an exact pin inside the sidecar; any other
unsatisfied requirement is a blocker. An unavailable requires_dist BLOCKS the
install: proceeding unverified could pin a sidecar whose imports then crash the
workers, and the caller just reached PyPI for the version check so a retry is cheap.
"""
reqs = _fetch_requires_dist(version)
if reqs is None:
return (), ["dependency metadata for this release (could not be fetched from PyPI; retry)"]
try:
from importlib.metadata import PackageNotFoundError
from importlib.metadata import version as _installed_version
from packaging.requirements import InvalidRequirement, Requirement
except Exception:
return (), []
extras: list[str] = []
blockers: list[str] = []
for raw in reqs:
try:
req = Requirement(raw)
except InvalidRequirement:
continue
if req.extras or (req.marker is not None and not req.marker.evaluate()):
continue
name = _canonical_dep_name(req.name)
if name in _IGNORED_DEPS:
continue
if name in _SIDECAR_PROVIDED:
if not req.specifier.contains(_SIDECAR_PROVIDED[name], prereleases = True):
blockers.append(raw)
continue
try:
installed = _installed_version(req.name)
except PackageNotFoundError:
installed = None
if installed is not None and req.specifier.contains(installed, prereleases = True):
continue
if name in _SHADOWABLE_DEPS:
exact = _resolve_exact_version(name, req.specifier)
if exact is None:
blockers.append(raw)
else:
extras.append(f"{name}=={exact}")
else:
blockers.append(raw)
return tuple(extras), blockers
def is_install_in_progress() -> bool:
"""True while a latest-transformers install or lazy repair holds the sidecar swap
reservation. Training and export starts check this so a fresh worker never
activates the sidecar mid-swap."""
from utils.transformers_version import sidecar_swap_in_progress
return sidecar_swap_in_progress()
def install_latest_transformers(
version: str,
before_swap = None,
reserved: bool = False,
) -> dict:
"""Consented install of the latest transformers sidecar; returns a structured result.
Guards: the requested *version* must match the current PyPI latest from the (cached)
snapshot, so a client cannot pin an arbitrary package version through this endpoint.
On success ``.venv_t5_latest`` is provisioned and pinned; routing then resolves the
new tier automatically on this and every future start. *before_swap* is forwarded
to the stage-and-swap: it runs only after the staged install succeeded, right
before the live sidecar is replaced. *reserved* means the caller already holds the
sidecar swap reservation (the install route takes it before waiting on the
inference lifecycle gate, so worker starts see it for the whole window).
"""
from utils.transformers_version import end_sidecar_swap, try_begin_sidecar_swap
if not reserved and not try_begin_sidecar_swap():
return {
"success": False,
"version": version,
"message": "A transformers installation is already in progress.",
}
try:
return _install_latest_transformers_locked(version, before_swap = before_swap)
finally:
if not reserved:
end_sidecar_swap()
def _install_latest_transformers_locked(version: str, before_swap = None) -> dict:
"""Body of install_latest_transformers; runs with the in-progress flag held."""
if _disabled():
return {
"success": False,
"version": version,
"message": "Latest-transformers installs are disabled "
"(UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS).",
}
if _env_offline():
return {
"success": False,
"version": version,
"message": "Cannot install: Studio is in offline mode.",
}
# Re-verify against a LIVE snapshot (a release may land inside the cache TTL);
# fall back to the cached one on fetch failure.
global _memory_snapshot
snapshot = _refresh_snapshot()
if snapshot is not None:
with _lock:
_memory_snapshot = snapshot
_save_snapshot_file(snapshot)
else:
snapshot = _get_snapshot()
if snapshot is None:
return {
"success": False,
"version": version,
"message": "Could not verify the latest transformers release on PyPI.",
}
if version != snapshot["pypi_version"]:
return {
"success": False,
"version": version,
"message": f"Requested version {version!r} is not the latest transformers "
f"release ({snapshot['pypi_version']}).",
# Lets the consent dialog retry with the release that superseded the
# one /validate saw, instead of re-sending the stale version forever.
"latest_version": snapshot["pypi_version"],
}
extra_packages, blockers = compat_plan(version)
if blockers:
return {
"success": False,
"version": version,
"message": "Cannot install transformers "
f"{version}: this environment does not satisfy {', '.join(blockers)}. "
"A Studio update is required first.",
}
if not ensure_latest_transformers_venv(version, extra_packages, before_swap = before_swap):
return {
"success": False,
"version": version,
"message": f"Installing transformers {version} failed; see the Studio logs.",
}
_invalidate_capability_caches()
return {
"success": True,
"version": version,
"message": f"Installed transformers {version} into the latest sidecar "
f"(pinned: {latest_venv_pinned_version()}).",
}
def _invalidate_capability_caches():
"""Drop caches computed before the new sidecar existed: tier probes and the
latest tier's model_type mapping (stale on upgrade) plus vision detection
(a raw-heuristic False may now defer to the sidecar AutoConfig probe)."""
try:
from utils import transformers_version as tv
tv._probe_tier_cache.clear()
tv._config_mapping_cache.pop("latest", None)
except Exception:
pass
try:
from utils.models import model_config as mc
mc._vision_detection_cache.clear()
except Exception:
pass

812
studio/backend/utils/transformers_version.py
View file

@ -35,9 +35,12 @@ import json
import structlog
from loggers import get_logger
import os
import re
import shutil
import subprocess
import sys
import threading
import time
from pathlib import Path
from utils.native_path_leases import child_env_without_native_path_secret
@ -235,8 +238,12 @@ _VENV_T5_DIR = _VENV_T5_550_DIR
# reuses the workspace torch (torch-agnostic).
_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor")
# Tier precedence: higher rank wins in _higher_tier.
_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3}
# User-consented "latest transformers" sidecar (utils/transformers_latest.py); pinned version in a marker file.
_VENV_T5_LATEST_DIR = str(_studio_root() / ".venv_t5_latest")
_LATEST_PIN_MARKER = ".unsloth_pinned_transformers"
# Tier precedence: higher rank wins in _higher_tier. "latest" outranks every fixed tier.
_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3, "latest": 4}
def _higher_tier(a: str, b: str) -> str:
@ -254,20 +261,40 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None =
``hf_token`` is forwarded to tier detection so a gated/private model whose only 5.x
signal is an authenticated config/tokenizer reaches the right sidecar, not the default.
"""
# Pre-resolve only LoRA adapters; full checkpoints go to get_transformers_tier so their
# local config.json drives the tier (a full checkpoint with a private/offline
# _name_or_path must not resolve to an unreachable HF id and skip its own config).
# Pre-resolve LoRA adapters (local dir or remote adapter repo); full checkpoints
# go to get_transformers_tier so their local config.json drives the tier (a full
# checkpoint with a private/offline _name_or_path must not resolve to an
# unreachable HF id and skip its own config). Remote adapters activate for their
# BASE model, matching latest_tier_active_for and the inference worker.
if _is_lora_adapter_dir(Path(model_name)):
resolved = _resolve_base_model(model_name)
else:
resolved = model_name
resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name
tier = get_transformers_tier(resolved, hf_token)
if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"):
# Gate on a real local config.json: a checkpoint carries config the base may not
# surface, but path names alone must not upgrade a plain adapter.
tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token))
if tier == "510":
if tier == "latest":
pinned = latest_venv_pinned_version()
if pinned is None or not _ensure_venv_t5_latest_exists():
raise RuntimeError(
f"Cannot activate the latest-transformers sidecar: "
f".venv_t5_latest missing or unpinned at {_VENV_T5_LATEST_DIR}"
)
if _VENV_T5_LATEST_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_LATEST_DIR)
logger.info(
"Prepended transformers %s venv to sys.path from %s "
"(path only; the loaded version is confirmed later by "
"'Subprocess loaded transformers ...' on first import)",
pinned,
_VENV_T5_LATEST_DIR,
)
_pp = os.environ.get("PYTHONPATH", "")
os.environ["PYTHONPATH"] = _VENV_T5_LATEST_DIR + (os.pathsep + _pp if _pp else "")
elif tier == "510":
if not _ensure_venv_t5_510_exists():
raise RuntimeError(
f"Cannot activate transformers {TRANSFORMERS_510_VERSION}: "
@ -322,6 +349,34 @@ def activate_transformers_for_subprocess(model_name: str, hf_token: str | None =
logger.info("Using default transformers (4.57.x) for %s", model_name)
def latest_tier_active_for(model_name: str, hf_token: str | None = None) -> bool:
"""True when *model_name* routes to the consented latest-transformers sidecar.
Mirrors the inference worker's pre-activation resolution (local adapter dir,
then a remote adapter's Hub adapter_config.json). ``latest`` only wins when
the sidecar exists with a valid pin, i.e. exactly the loads that will import
the newest release. Never raises: any resolution failure returns False so
callers treat the model as a known tier.
"""
try:
# No consented sidecar pin means nothing routes to latest; return before
# any resolution so the common case costs no config or network reads.
if latest_venv_pinned_version() is None:
return False
if _is_lora_adapter_dir(Path(model_name)):
resolved = _resolve_base_model(model_name)
else:
# A remote LoRA activates the sidecar for its BASE model; sizing and the
# worker's 4-bit guard must see that base too, not the adapter repo.
resolved = _remote_lora_base(model_name, hf_token = hf_token) or model_name
tier = get_transformers_tier(resolved, hf_token)
if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"):
tier = _higher_tier(tier, get_transformers_tier(model_name, hf_token))
return tier == "latest"
except Exception:
return False
def _has_adapter_weights(path: Path) -> bool:
"""True if *path* holds LoRA adapter weight files (``adapter_model.*``)."""
try:
@ -881,17 +936,85 @@ def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None:
_config_mapping_cache: dict[str, frozenset[str]] = {}
def _latest_tier_disabled() -> bool:
"""Kill switch shared with utils.transformers_latest: lets operators roll
back a provisioned latest sidecar without deleting files."""
return os.environ.get("UNSLOTH_STUDIO_NO_LATEST_TRANSFORMERS", "").strip().lower() in (
"1",
"true",
"yes",
"on",
)
# Failed lazy repairs back off so a broken sidecar can't turn every routing
# call into a pip install attempt.
_latest_repair_failed_at: float = 0.0
_LATEST_REPAIR_BACKOFF_SECS = 5 * 60
def _latest_sidecar_intact() -> bool:
"""The pinned latest sidecar exists with its transformers dir and every pinned
package. False when the pin itself is gone: a cached 'latest' mapping must then be
dropped (routing re-resolves to no latest tier), not trusted, and a sidecar that kept
transformers/ but lost a pinned package must self-heal rather than route models to a
latest tier that fails activation in workers, which refuse parent-only repairs.
(_overlay_transformers_dir only calls this after gating on a present pin, so the
pin-missing case here is the cache-revalidation caller whose pin was deleted after
the mapping was first cached.)"""
pin = _latest_pin_data()
if pin is None:
return False
return _venv_dir_is_valid(_VENV_T5_LATEST_DIR, tuple(pin["packages"]))
def _overlay_transformers_dir(tier: str) -> str | None:
"""transformers source dir for a tier, located without importing it."""
global _latest_repair_failed_at
if tier != "default":
root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier)
# latest requires a valid pin and the kill switch off.
if tier == "latest" and (_latest_tier_disabled() or latest_venv_pinned_version() is None):
return None
root = {
"530": _VENV_T5_530_DIR,
"550": _VENV_T5_550_DIR,
"510": _VENV_T5_510_DIR,
"latest": _VENV_T5_LATEST_DIR,
}.get(tier)
src = os.path.join(root, "transformers") if root else None
if src and tier == "latest" and not _latest_sidecar_intact():
# A valid pin whose sidecar vanished or lost a pinned package (partial
# deletion, disk issue, interrupted external edits) must self-heal, or
# latest-only models either silently route to older tiers or reach a
# worker that cannot repair, failing every load until a manual
# reinstall. Repair under the swap reservation; back off after a
# failure so routing calls don't hammer pip.
repaired = False
if time.time() - _latest_repair_failed_at >= _LATEST_REPAIR_BACKOFF_SECS:
if _ensure_venv_t5_latest_exists():
_latest_repair_failed_at = 0.0
repaired = True
else:
_latest_repair_failed_at = time.time()
if not repaired:
# Still broken: treat the overlay as unavailable rather than route
# models to a tier whose worker activation is known to fail. Models
# an older tier supports keep loading there until a repair succeeds,
# matching the behavior when the sidecar dir is missing entirely.
return None
return src if src and _safe_is_dir(Path(src)) else None
# default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one
# is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir.
sidecars = tuple(
os.path.abspath(d) + os.sep
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR)
for d in (
_VENV_T5_530_DIR,
_VENV_T5_550_DIR,
_VENV_T5_510_DIR,
_VENV_T5_LATEST_DIR,
_VENV_LLMCOMPRESSOR_DIR,
)
)
candidates = []
try:
@ -930,11 +1053,47 @@ def _mapping_first_keys(value: ast.AST) -> set[str]:
return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)}
def _model_types_from_source(source: str) -> set[str]:
"""model_type keys of CONFIG_MAPPING_NAMES in *source* (AST only, no execution).
Handles the direct ``CONFIG_MAPPING_NAMES = ...`` binding (dict literal or
OrderedDict/dict call over 2-tuple lists and **{...} unpacking) and any
``CONFIG_MAPPING_NAMES.update({...})`` mutation. Shared by the on-disk overlay
reader below and the remote latest-release checker (utils/transformers_latest.py).
"""
keys: set[str] = set()
tree = ast.parse(source)
for node in ast.walk(tree):
if isinstance(node, ast.Assign) and any(
isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets
):
keys |= _mapping_first_keys(node.value)
elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
fn = node.value.func
if (
isinstance(fn, ast.Attribute)
and fn.attr == "update"
and isinstance(fn.value, ast.Name)
and fn.value.id == "CONFIG_MAPPING_NAMES"
):
keys |= _mapping_first_keys(node.value)
return keys
def _config_model_types(tier: str) -> frozenset[str]:
"""model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py)."""
# Kill switch beats the cache: a stale mapping must not keep routing latest-only models until restart.
if tier == "latest" and _latest_tier_disabled():
return frozenset()
cached = _config_mapping_cache.get(tier)
if cached is not None:
return cached
# A cached 'latest' mapping can outlive the sidecar it was parsed from: if the
# pinned sidecar was since deleted or lost a package in this process, drop the
# cache so routing re-resolves through _overlay_transformers_dir (which self-heals)
# instead of routing latest-only models to a broken tier until restart.
if tier != "latest" or _latest_sidecar_intact():
return cached
_config_mapping_cache.pop("latest", None)
tdir = _overlay_transformers_dir(tier)
if tdir is None:
return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads
@ -944,22 +1103,7 @@ def _config_model_types(tier: str) -> frozenset[str]:
if not _safe_is_file(path):
continue
try:
tree = ast.parse(path.read_text(encoding = "utf-8"))
for node in ast.walk(tree):
# direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation
if isinstance(node, ast.Assign) and any(
isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets
):
keys |= _mapping_first_keys(node.value)
elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
fn = node.value.func
if (
isinstance(fn, ast.Attribute)
and fn.attr == "update"
and isinstance(fn.value, ast.Name)
and fn.value.id == "CONFIG_MAPPING_NAMES"
):
keys |= _mapping_first_keys(node.value)
keys |= _model_types_from_source(path.read_text(encoding = "utf-8"))
except Exception:
continue
result = frozenset(keys)
@ -967,23 +1111,73 @@ def _config_model_types(tier: str) -> frozenset[str]:
return result
def _tier_from_config_mapping(cfg: dict) -> str | None:
"""Lowest tier whose transformers ships cfg's model_type, or None if unknown."""
model_type = cfg.get("model_type")
if not isinstance(model_type, str):
for key in _NESTED_CONFIG_KEYS:
sub = cfg.get(key)
if isinstance(sub, dict) and isinstance(sub.get("model_type"), str):
model_type = sub["model_type"]
break
if not isinstance(model_type, str):
return None
def _model_types_from_config(cfg: dict) -> list[str]:
"""All model_types in the config: the primary (top-level, else first nested)
first, then every other nested sub-config. Wrappers instantiate sub-configs
through CONFIG_MAPPING, so nested types matter for routing too."""
seen: list[str] = []
def add(value):
if isinstance(value, str) and value and value not in seen:
seen.append(value)
add(cfg.get("model_type"))
for key in _NESTED_CONFIG_KEYS:
sub = cfg.get(key)
if isinstance(sub, dict):
add(sub.get("model_type"))
for value in cfg.values():
if isinstance(value, dict):
add(value.get("model_type"))
return seen
def _lowest_tier_for(model_type: str) -> str | None:
for tier in sorted(_TIER_RANK, key = _TIER_RANK.get):
if model_type in _config_model_types(tier):
return tier
return None
def _tier_from_config_mapping(cfg: dict) -> str | None:
"""Lowest tier able to load every model_type in cfg, or None when the
primary type is unknown everywhere. A nested type can raise the tier (its
sub-config is built through CONFIG_MAPPING); an unknown nested type never
vetoes, since no installed tier could load it either way (the latest
checker handles surfacing the install prompt for it)."""
types = _model_types_from_config(cfg)
if not types:
return None
best = _lowest_tier_for(types[0])
if best is None:
return None
for model_type in types[1:]:
tier = _lowest_tier_for(model_type)
if tier is not None and _TIER_RANK[tier] > _TIER_RANK[best]:
best = tier
return best
def _raise_tier_for_nested(cfg: dict | None, tier: str) -> str:
"""Raise *tier* when the mapping resolver needs a higher one for *cfg*.
A wrapper's top-level model_type can match a hardcoded fast path while a
nested text/vision config's type only exists in a newer sidecar (e.g. the
installed latest); its sub-config is built through CONFIG_MAPPING, so the
fast-path tier would fail to load it. Raise-only: never lowers a fast-path
match, so name overrides (Qwen3.6) keep their tier. Never raises an
exception: a resolution failure keeps the fast-path tier."""
if not isinstance(cfg, dict):
return tier
try:
mapped = _tier_from_config_mapping(cfg)
if mapped is not None and _TIER_RANK.get(mapped, 0) > _TIER_RANK.get(tier, 0):
return mapped
except Exception:
pass
return tier
# --- AutoConfig probe: general tier resolution for ambiguous models ----------
# When the cheap signals only say "needs some 5.x", parse config.json with the built-in
# parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond
@ -1039,9 +1233,19 @@ def _probe_tier_venvs():
"530": (_VENV_T5_530_DIR, _ensure_venv_t5_530_exists),
"550": (_VENV_T5_550_DIR, _ensure_venv_t5_550_exists),
"510": (_VENV_T5_510_DIR, _ensure_venv_t5_510_exists),
"latest": (_VENV_T5_LATEST_DIR, _ensure_venv_t5_latest_exists),
}
def _probe_tier_order() -> tuple[str, ...]:
"""Sidecar probe order. The consented "latest" sidecar joins only once it is
provisioned (pin marker present): an absent optional tier must not flip the probe's
skipped-tier bookkeeping, keeping pre-latest behavior byte-identical."""
if not _latest_tier_disabled() and latest_venv_pinned_version() is not None:
return _PROBE_TIER_ORDER + ("latest",)
return _PROBE_TIER_ORDER
def _probe_autoconfig(target_dir: str, model_name: str, hf_token: str | None) -> bool | None:
"""Parse config.json with the built-in parser inside *target_dir*'s sidecar.
True = parses, False = parse/version failure (escalate), None = transient
@ -1119,7 +1323,7 @@ def _probe_tier(
stays on the default. Cached per _probe_cache_key (process lifetime). No Hub sha is
resolved: that would import huggingface_hub before the sidecar is on sys.path.
"""
if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes"):
if os.environ.get("UNSLOTH_DISABLE_TIER_PROBE", "").lower() in ("1", "true", "yes", "on"):
return floor
key = _probe_cache_key(model_name)
# Key by probe mode: the default-first path can return 'default', which must not be
@ -1127,7 +1331,10 @@ def _probe_tier(
if include_default or floor != "530":
key = f"{key}\0floor={floor}:def={int(include_default)}"
if key in _probe_tier_cache:
return _probe_tier_cache[key]
cached = _probe_tier_cache[key]
# Kill switch beats the cache (like _config_model_types): a stale 'latest' probe must not keep activating it.
if cached != "latest" or not _latest_tier_disabled():
return cached
def _cache(tier: str, *, skipped: bool) -> str:
# Do not pin a result that depended on a skipped lower tier: once that sidecar is
@ -1137,7 +1344,8 @@ def _probe_tier(
return tier
venvs = _probe_tier_venvs()
order = (("default",) + _PROBE_TIER_ORDER) if include_default else _PROBE_TIER_ORDER
sidecar_order = _probe_tier_order()
order = (("default",) + sidecar_order) if include_default else sidecar_order
probed_count = 0
skipped_any = False
for tier in order:
@ -1264,17 +1472,21 @@ def get_transformers_tier(
cfg = _load_config_json(model_name, hf_token)
if cfg is not None:
if _config_needs_510(cfg):
tier = _raise_tier_for_nested(cfg, "510")
logger.info(
"Transformers tier 510 selected for %s (local config.json check)",
"Transformers tier %s selected for %s (local config.json check)",
tier,
model_name,
)
return "510"
return tier
if _config_needs_550(cfg):
tier = _raise_tier_for_nested(cfg, "550")
logger.info(
"Transformers tier 550 selected for %s (local config.json check)",
"Transformers tier %s selected for %s (local config.json check)",
tier,
model_name,
)
return "550"
return tier
if _config_needs_530(cfg):
# Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name. Only a real
# Hub id (or the folder basename) may override 530, so a stale local
@ -1287,17 +1499,20 @@ def get_transformers_tier(
)
override = _higher_tier_name_override(hint_src)
if override is not None:
override = _raise_tier_for_nested(cfg, override)
logger.info(
"Transformers tier %s selected for %s (name overrides 530 config)",
override,
model_name,
)
return override
tier = _raise_tier_for_nested(cfg, "530")
logger.info(
"Transformers tier 530 selected for %s (local config.json check)",
"Transformers tier %s selected for %s (local config.json check)",
tier,
model_name,
)
return "530"
return tier
# Unknown arch: resolve the base id from config. A resolved local dir
# recurses (config check); a Hub id uses name rules only (no network).
resolved = _resolve_base_model(model_name)
@ -1359,6 +1574,13 @@ def get_transformers_tier(
result = _tier_from_name(model_name)
if result is not None:
tier, match = result
# With a consented latest sidecar pinned, a name that matches a fixed
# tier can still carry a latest-only model_type (e.g. a newer variant
# reusing a family name); consult the config so an accepted upgrade
# actually routes to the sidecar it installed. Costs a config read only
# in the pinned case, keeping the pre-latest path I/O-free.
if latest_venv_pinned_version() is not None:
tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), tier)
logger.info(
"Transformers tier %s selected for %s (substring match: %s)",
tier,
@ -1369,11 +1591,13 @@ def get_transformers_tier(
# --- Slow config fallbacks (network for HF IDs; authenticated with hf_token) --------
if _check_config_needs_510(model_name, hf_token):
logger.info("Transformers tier 510 selected for %s (config.json check)", model_name)
return "510"
tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "510")
logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
return tier
if _check_config_needs_550(model_name, hf_token):
logger.info("Transformers tier 550 selected for %s (config.json check)", model_name)
return "550"
tier = _raise_tier_for_nested(_load_config_json(model_name, hf_token), "550")
logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
return tier
if _check_config_needs_530(model_name, hf_token):
# Qwen3.6 reuses Qwen3.5 config ids but needs 5.5 by name; honor a real Hub-id name
# hint from _name_or_path before selecting 530.
@ -1383,14 +1607,16 @@ def get_transformers_tier(
base if isinstance(base, str) and base != model_name else None
)
if override is not None:
override = _raise_tier_for_nested(remote_cfg, override)
logger.info(
"Transformers tier %s selected for %s (name overrides 530 config)",
override,
model_name,
)
return override
logger.info("Transformers tier 530 selected for %s (config.json check)", model_name)
return "530"
tier = _raise_tier_for_nested(remote_cfg, "530")
logger.info("Transformers tier %s selected for %s (config.json check)", tier, model_name)
return tier
# _load_config_json (not the cache-only reader) so a config served from the hub
# cache during a transient outage still feeds the mapping resolver.
remote_cfg = _load_config_json(model_name, hf_token)
@ -1657,6 +1883,471 @@ def _ensure_venv_t5_exists() -> bool:
return _ensure_venv_t5_550_exists()
# --- User-consented "latest transformers" sidecar (.venv_t5_latest) --------------------------
# Provisioned via ensure_latest_transformers_venv() after the user confirms the upgrade popup
# (utils/transformers_latest.py); pinned in a marker file so restarts revalidate and routing auto-picks it.
# PEP 440-ish release strings only (guards the pip install spec against injection).
_LATEST_VERSION_RE = r"[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0-9]+)?"
def _is_valid_version_string(version: str) -> bool:
import re
return isinstance(version, str) and re.fullmatch(_LATEST_VERSION_RE, version) is not None
# Only the sidecar recipe's own packages, as plain (optionally ==pinned) specs, may
# come from the on-disk pin marker; anything else (URLs, extras, options) is rebuilt.
_PIN_SPEC_RE = re.compile(r"^[A-Za-z0-9_.-]+(==[A-Za-z0-9_.+-]+)?$")
_PIN_ALLOWED_NAMES = frozenset(
{
"transformers",
"huggingface_hub",
"huggingface-hub",
"hf_xet",
"hf-xet",
"tiktoken",
"tokenizers",
"safetensors",
}
)
def _is_safe_pin_spec(spec: str) -> bool:
if not _PIN_SPEC_RE.match(spec):
return False
name = spec.split("==", 1)[0].lower().replace("_", "-")
return name in {n.replace("_", "-") for n in _PIN_ALLOWED_NAMES}
def _recover_stranded_latest_sidecar() -> None:
"""Restore a sidecar stranded at ``.old`` by a swap whose activation rename AND its
rollback both failed (e.g. a lingering worker file handle on Windows blocked both).
That double failure leaves no live dir and the pin marker gone with it, so the
sidecar reads as unprovisioned and never self-heals. Recover only when no live dir
exists and no swap is in flight: the reservation is held throughout the swap, so the
transient live-absent window of a legitimate swap never triggers a restore."""
live = Path(_VENV_T5_LATEST_DIR)
retired = Path(_VENV_T5_LATEST_DIR + ".old")
try:
if live.exists() or not retired.is_dir() or sidecar_swap_in_progress():
return
os.rename(retired, live)
logger.info("Recovered .venv_t5_latest from a stranded .old after a failed swap")
except OSError:
pass
def _latest_pin_data() -> dict | None:
"""Parsed pin marker: {"version": str, "packages": [specs...]}, or None.
The marker is JSON; a plain version string (older/simpler writers) is tolerated and
expanded with the default package set.
"""
_recover_stranded_latest_sidecar()
marker = Path(_VENV_T5_LATEST_DIR) / _LATEST_PIN_MARKER
try:
if not marker.is_file():
return None
raw = marker.read_text(encoding = "utf-8").strip()
except Exception:
return None
try:
data = json.loads(raw)
except ValueError:
data = raw
if isinstance(data, str):
if not _is_valid_version_string(data):
return None
return {"version": data, "packages": list(_venv_t5_latest_packages(data))}
if not isinstance(data, dict):
return None
version = data.get("version")
if not _is_valid_version_string(version):
return None
packages = data.get("packages")
if not (
isinstance(packages, list)
and packages
and all(isinstance(p, str) and _is_safe_pin_spec(p) for p in packages)
):
# Malformed or unexpected specs (the pin is user-writable on disk) never
# reach pip: rebuild the canonical set for the pinned version instead.
packages = list(_venv_t5_latest_packages(version))
return {"version": version, "packages": packages}
def latest_venv_pinned_version() -> str | None:
"""Exact transformers version pinned in .venv_t5_latest's marker, or None if the
sidecar was never provisioned (or the marker is unreadable/invalid)."""
data = _latest_pin_data()
return data["version"] if data else None
def _venv_t5_latest_packages(version: str, extra_packages: tuple[str, ...] = ()) -> tuple[str, ...]:
"""Package set for the latest sidecar; mirrors the fixed .venv_t5_* sidecars.
*extra_packages* carries dep-compat shadows (e.g. a newer tokenizers) computed by
utils.transformers_latest before install."""
return (
f"transformers=={version}",
"huggingface_hub==1.8.0",
"hf_xet==1.4.2",
"tiktoken",
) + tuple(extra_packages)
# Single reservation for ANY .venv_t5_latest replacement (consented install or lazy repair),
# checked by training/export starts so no worker spawns mid-swap. Backed by a lock FILE (not just
# this flag) so a lazy repair running in a worker subprocess stays visible to the parent's route
# checks; the in-process flag marks ownership (only the owner unlinks the file).
_sidecar_swap_lock = threading.Lock()
_sidecar_swap_active = False
_sidecar_swap_token: str | None = None
_sidecar_swap_kind: str | None = None
# An install is minutes; a lock this old is a crashed owner, not a live swap.
_SWAP_LOCK_STALE_SECS = 2 * 60 * 60
def _swap_lock_path() -> Path:
return Path(_VENV_T5_LATEST_DIR + ".swaplock")
def _pid_alive(pid) -> bool:
if not isinstance(pid, int) or pid <= 0:
return False
try:
import psutil
return psutil.pid_exists(pid)
except Exception:
pass
if os.name == "nt":
# os.kill(pid, 0) is NOT a POSIX signal-0 liveness probe on Windows: signal 0
# is CTRL_C_EVENT, so CPython routes it through GenerateConsoleCtrlEvent (a real
# Ctrl+C to that console group) rather than a harmless check. Probe via OpenProcess.
try:
import ctypes
from ctypes import wintypes
kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
kernel32.OpenProcess.restype = wintypes.HANDLE
# PROCESS_QUERY_LIMITED_INFORMATION: minimal right, granted across integrity levels.
handle = kernel32.OpenProcess(0x1000, False, pid)
if handle:
kernel32.CloseHandle(handle)
return True
# ERROR_ACCESS_DENIED means the process exists but we may not query it.
return ctypes.get_last_error() == 5
except Exception:
return False
try:
os.kill(pid, 0)
return True
except OSError:
return False
except Exception:
return False
def _swap_lock_is_stale(path: Path) -> bool:
"""Stale when the recorded owner is provably dead: a crashed installer is reclaimed
at once, not after the long cutoff, so `/load`, training, export, and repair are not
wedged for hours after a crash. A live but slow pip install keeps its lock (its PID
is alive), so breaking it and racing two swaps on the same staging dirs stays
impossible. Only a lock whose PID can't be read (mid-write or corrupt) falls back to
the age cutoff, so the create-before-metadata-write window is never mistaken for dead."""
try:
age = time.time() - path.stat().st_mtime
except OSError:
return False
data = _read_swap_lock(path) or {}
pid = data.get("pid")
if not isinstance(pid, int) or pid <= 0:
return age > _SWAP_LOCK_STALE_SECS
return not _pid_alive(pid)
class SidecarSwapInProgress(RuntimeError):
"""A worker start lost the race to a .venv_t5_latest install/repair; retryable."""
def _read_swap_lock(path: Path) -> dict | None:
try:
data = json.loads(path.read_text(encoding = "utf-8"))
return data if isinstance(data, dict) else {}
except FileNotFoundError:
return None
except OSError:
return {}
except Exception:
return {}
def try_begin_sidecar_swap(kind: str = "install") -> bool:
"""Reserve the sidecar swap window; False when one is already reserved
(in this process or, via the lock file, in any worker subprocess).
*kind* is "install" (consented route) or "repair" (lazy venv repair)."""
global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind
with _sidecar_swap_lock:
if _sidecar_swap_active:
return False
token = f"{os.getpid()}-{time.time_ns()}"
path = _swap_lock_path()
try:
path.parent.mkdir(parents = True, exist_ok = True)
except OSError:
pass
for attempt in range(2):
try:
fd = os.open(str(path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
break
except FileExistsError:
if attempt or not _swap_lock_is_stale(path):
return False
try:
path.unlink()
except OSError:
return False
except OSError:
# Lock file not creatable (odd filesystem): fall back to the process-local reservation.
fd = None
break
if fd is not None:
try:
with os.fdopen(fd, "w") as f:
f.write(
json.dumps(
{"pid": os.getpid(), "at": time.time(), "token": token, "kind": kind}
)
)
except OSError:
pass
_sidecar_swap_active = True
_sidecar_swap_token = token
_sidecar_swap_kind = kind
return True
def end_sidecar_swap() -> None:
"""Release the reservation taken by :func:`try_begin_sidecar_swap`."""
global _sidecar_swap_active, _sidecar_swap_token, _sidecar_swap_kind
with _sidecar_swap_lock:
if _sidecar_swap_active:
# Only the file WE wrote is removed: if this reservation was declared
# stale and superseded, unlinking blindly would drop the new owner's
# live lock and unguard its in-flight swap.
path = _swap_lock_path()
data = _read_swap_lock(path)
if data is not None and data.get("token", _sidecar_swap_token) == _sidecar_swap_token:
try:
path.unlink()
except OSError:
pass
_sidecar_swap_active = False
_sidecar_swap_token = None
_sidecar_swap_kind = None
def sidecar_swap_in_progress() -> bool:
"""True while a .venv_t5_latest install or repair holds the reservation,
in this process or any other Studio process (lock file)."""
return sidecar_swap_kind() is not None
def sidecar_swap_kind() -> str | None:
"""The active reservation's kind ("install" / "repair"), or None when idle.
Lets guards that rely on the install route's own abort-on-active-worker
checks keep refusing for repairs, which have no such checks."""
with _sidecar_swap_lock:
if _sidecar_swap_active:
return _sidecar_swap_kind or "install"
path = _swap_lock_path()
try:
if not path.is_file() or _swap_lock_is_stale(path):
return None
except OSError:
return None
data = _read_swap_lock(path) or {}
kind = data.get("kind")
return kind if kind in ("install", "repair") else "install"
def _stage_and_swap_latest_venv(
version: str,
packages: tuple[str, ...],
before_swap = None,
) -> bool:
"""Stage-and-swap: build the new sidecar next to the live one and swap only
once complete, so a failed install or marker write never destroys a
previously working .venv_t5_latest or its pin. Shared by the consented
install and the lazy repair path. *before_swap* (optional callable) runs
after the staging build succeeds and immediately before the live dir is
replaced, so callers can tear down workers only when the swap is certain;
if it raises, the previous sidecar is left untouched."""
staging = _VENV_T5_LATEST_DIR + ".staging"
retired = _VENV_T5_LATEST_DIR + ".old"
shutil.rmtree(staging, ignore_errors = True)
try:
if not _ensure_venv_dir(staging, packages, f"transformers {version} (latest)"):
# No exception, so the except cleanup below never runs; drop the partial dir.
shutil.rmtree(staging, ignore_errors = True)
return False
(Path(staging) / _LATEST_PIN_MARKER).write_text(
json.dumps({"version": version, "packages": list(packages)}), encoding = "utf-8"
)
if before_swap is not None:
before_swap()
shutil.rmtree(retired, ignore_errors = True)
if os.path.isdir(_VENV_T5_LATEST_DIR):
os.rename(_VENV_T5_LATEST_DIR, retired)
try:
os.rename(staging, _VENV_T5_LATEST_DIR)
except OSError:
# Restore the previous sidecar if the final swap fails.
if not os.path.isdir(_VENV_T5_LATEST_DIR) and os.path.isdir(retired):
os.rename(retired, _VENV_T5_LATEST_DIR)
raise
except Exception as exc:
logger.error("Could not provision transformers %s into .venv_t5_latest: %s", version, exc)
shutil.rmtree(staging, ignore_errors = True)
return False
shutil.rmtree(retired, ignore_errors = True)
# CONFIG_MAPPING_NAMES may have changed: drop the cached key set.
_config_mapping_cache.pop("latest", None)
logger.info("Provisioned .venv_t5_latest with transformers %s", version)
return True
def _workers_active_for_repair() -> bool:
"""Best-effort: any parent-visible chat/training/export worker alive. Never
raises; unavailable backends (worker subprocess, early startup) count idle."""
try:
from core.training import get_training_backend
if get_training_backend().is_training_active():
return True
except Exception:
pass
try:
from core.export import get_export_backend
_export = get_export_backend()
if _export.is_export_active():
return True
_alive = getattr(_export, "is_worker_alive", None)
if callable(_alive) and _alive():
return True
except Exception:
pass
try:
from core.inference import get_inference_backend
backend = get_inference_backend()
if getattr(backend, "active_model_name", None):
return True
# An in-flight load counts too: its worker spawns moments later.
if getattr(backend, "loading_models", None):
return True
_alive = getattr(backend, "is_worker_alive", None)
if callable(_alive) and _alive():
return True
except Exception:
pass
return False
def _ensure_venv_t5_latest_exists() -> bool:
"""Ensure .venv_t5_latest/ holds its pinned transformers version.
Never installs without a pin: an unprovisioned sidecar (no marker) returns False so
routing and probing behave exactly as before the feature existed. With a pin present
it repairs a broken dir the same way the fixed sidecars do.
"""
pin = _latest_pin_data()
if pin is None:
return False
version = pin["version"]
packages = tuple(pin["packages"])
if _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages):
return True
if _env_offline():
logger.warning(
".venv_t5_latest (transformers %s) is incomplete and offline mode is set; "
"cannot repair it.",
version,
)
return False
# Repairs are a parent-process action: a worker child's backend singletons are
# empty, so it cannot see live siblings that may still lazy-import from the
# sidecar. Fail activation in the child instead; the parent's routing
# self-heal (guarded below) performs the actual repair.
try:
import multiprocessing as _mp
if _mp.parent_process() is not None:
logger.warning(
".venv_t5_latest is incomplete; repairs run in the parent process. "
"Retry after the parent repairs the sidecar."
)
return False
except Exception:
pass
# Same stage-and-swap as the install, under the same reservation so training/export starts
# (which check sidecar_swap_in_progress) wait out a lazy repair; a failed repair keeps the pin.
if not try_begin_sidecar_swap(kind = "repair"):
logger.warning(
"Cannot repair .venv_t5_latest: another sidecar install or repair is in progress."
)
return False
try:
# Worker check UNDER the reservation (the install route quiesces workers;
# a repair has none): worker starts set their active markers BEFORE
# rechecking the reservation, so either this check sees them and aborts,
# or their recheck sees this reservation and aborts -- no interleaving
# lets a worker spawn against a mid-swap sidecar.
if _workers_active_for_repair():
logger.warning(
"Cannot repair .venv_t5_latest: active chat/training/export workers "
"may be importing from it. Retry when they are idle."
)
return False
return _stage_and_swap_latest_venv(version, packages)
finally:
end_sidecar_swap()
def ensure_latest_transformers_venv(
version: str,
extra_packages: tuple[str, ...] = (),
before_swap = None,
) -> bool:
"""Provision .venv_t5_latest/ pinned to *version* (user-consented install path).
Reuses the same --target/--no-deps installer as the fixed sidecars, then writes the pin
marker (version + full package set) so the venv persists across restarts and
:func:`latest_venv_pinned_version` / routing pick it up automatically.
*extra_packages* carries dep-compat shadows (see utils.transformers_latest).
Returns True on success.
"""
if not _is_valid_version_string(version):
logger.error("Refusing to install invalid transformers version %r", version)
return False
if _env_offline():
logger.warning(
"Cannot install transformers %s: HF/transformers offline mode is set.", version
)
return False
packages = _venv_t5_latest_packages(version, extra_packages)
pin = _latest_pin_data()
if (
pin is not None
and pin["version"] == version
and tuple(pin["packages"]) == packages
and _venv_dir_is_valid(_VENV_T5_LATEST_DIR, packages)
):
return True
return _stage_and_swap_latest_venv(version, packages, before_swap = before_swap)
# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) ---------------------
# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize
# Qwen3.5 / Gemma-4 / Llama.
@ -1819,7 +2510,7 @@ def _activate_venv(venv_dir: str, label: str) -> None:
def _deactivate_5x() -> None:
"""Remove all .venv_t5_*/ dirs from sys.path, purge stale modules, reimport."""
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR):
for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_T5_LATEST_DIR):
while d in sys.path:
sys.path.remove(d)
logger.info("Removed venv_t5 dirs from sys.path")
@ -1853,14 +2544,25 @@ def ensure_transformers_version(model_name: str) -> None:
if _is_lora_adapter_dir(Path(model_name)):
resolved = _resolve_base_model(model_name)
else:
resolved = model_name
# A remote adapter's tier is its BASE model's (see activation above).
resolved = _remote_lora_base(model_name) or model_name
tier = get_transformers_tier(resolved)
if model_name != resolved and _safe_is_file(Path(model_name) / "config.json"):
# Gate on a real local config.json: a checkpoint carries config the base may not
# surface, but path names alone must not upgrade a plain adapter.
tier = _higher_tier(tier, get_transformers_tier(model_name))
if tier == "510":
if tier == "latest":
pinned = latest_venv_pinned_version()
if pinned is None:
raise RuntimeError(
f"Cannot activate the latest-transformers sidecar: "
f"no pin marker at {_VENV_T5_LATEST_DIR}"
)
target_version = pinned
venv_dir = _VENV_T5_LATEST_DIR
ensure_fn = _ensure_venv_t5_latest_exists
elif tier == "510":
target_version = TRANSFORMERS_510_VERSION
venv_dir = _VENV_T5_510_DIR
ensure_fn = _ensure_venv_t5_510_exists

2
studio/frontend/src/app/routes/__root.tsx
View file

@ -16,6 +16,7 @@ import {
type ChatSearch,
} from "@/features/chat";
import { RemoteCodeConsentDialog } from "@/features/security";
import { TransformersUpgradeDialog } from "@/features/transformers-upgrade";
import { useTrainingUnloadGuard } from "@/features/training";
import { useExportRuntimeLifecycle } from "@/features/export";
import { hasAuthToken } from "@/features/auth";
@ -230,6 +231,7 @@ function RootLayout() {
<PersonalizationSyncMount />
{!isAuthFlowRoute && <SettingsDialog />}
<RemoteCodeConsentDialog />
<TransformersUpgradeDialog />
{hideNavbar ? (
<main className="flex-1 pt-[var(--studio-hidden-route-top-inset,0px)] [--studio-titlebar-height:var(--studio-hidden-route-top-inset,0px)]">
<Suspense fallback={<RouteFallback />}>

5
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -1454,6 +1454,11 @@ async function autoLoadSmallestModel(): Promise<{
blockedByTrustRemoteCode = true;
return false;
}
// Never install packages from a background load; explicit loads raise the upgrade dialog.
if (validation.requires_transformers_upgrade) {
hadNonTrustFailure = true;
return false;
}
return true;
}

32
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -4,6 +4,10 @@
import { createElement, useCallback, useRef, useState } from "react";
import { toast } from "@/lib/toast";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import {
confirmTransformersUpgradeIfNeeded,
useTransformersUpgradeDialogStore,
} from "@/features/transformers-upgrade";
import { consumeNativePathToken } from "@/features/native-intents/api";
import {
notifyNative,
@ -245,6 +249,10 @@ function getTrustRemoteCodeRequiredMessage(modelName: string): string {
return `${modelName} was not loaded because its custom code was not approved. Load it again to review the code and approve it.`;
}
function getTransformersUpgradeRequiredMessage(modelName: string): string {
return `${modelName} was not loaded because it needs a newer transformers release that was not installed. Load it again to install it.`;
}
/**
* Reconcile the chat runtime store against `/api/inference/status`: refresh the
* models/loras catalogs and either re-pin the active checkpoint or clear the
@ -626,6 +634,30 @@ export function useChatModelRuntime() {
is_lora: isLora,
gguf_variant: ggufVariant ?? null,
});
// Upgrade consent runs before the security dialogs; Accept installs and the load continues.
if (validation.requires_transformers_upgrade) {
const upgraded = await confirmTransformersUpgradeIfNeeded({
modelName: modelId,
upgrade: validation.transformers_upgrade,
// No installable release: custom-code models may fall back to the trust_remote_code gate below.
trustRemoteCodeFallback: validation.requires_trust_remote_code,
});
// The install unloads the previous model before the swap (even when
// the swap then fails), so any exit after this point must roll back.
// False for the custom-code fallback, which resolves without installing.
if (
useTransformersUpgradeDialogStore
.getState()
.consumeServerUnloadedChat()
&& currentCheckpoint
) {
previousWasUnloaded = true;
}
if (!upgraded) {
throw new Error(getTransformersUpgradeRequiredMessage(displayName));
}
}
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
// Open the consent dialog when the model needs custom-code consent or has a
// flagged unsafe file. Fires even when trustRemoteCode is preset on, since the
// worker requires a matching fingerprint that only the dialog produces.

39
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -67,6 +67,10 @@ import { KnowledgeBaseComposerButton } from "@/features/rag/components/knowledge
import { NewProjectDialog } from "./components/new-project-dialog";
import { useChatProjects } from "./hooks/use-chat-projects";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import {
confirmTransformersUpgradeIfNeeded,
useTransformersUpgradeDialogStore,
} from "@/features/transformers-upgrade";
import { loadModel, validateModel } from "./api/chat-api";
import {
parseExternalModelId,
@ -929,6 +933,10 @@ export function SharedComposer({
return parts[parts.length - 1] || id;
}
// Set when an accepted transformers install unloaded the active model
// server-side; a later failure must then clear the stale checkpoint.
let upgradeUnloadedActive = false;
// Helper: load a model and update store checkpoint
async function ensureModelLoaded(
sel: CompareModelSelection,
@ -955,6 +963,31 @@ export function SharedComposer({
trust_remote_code: loadTrustRemoteCode,
chat_template_override: effectiveChatTemplateOverride,
});
// Upgrade dialog first (mirrors the primary load path).
if (validation.requires_transformers_upgrade) {
const upgraded = await confirmTransformersUpgradeIfNeeded({
modelName: sel.id,
upgrade: validation.transformers_upgrade,
// No installable release: custom-code models may fall back to the trust_remote_code gate below.
trustRemoteCodeFallback: validation.requires_trust_remote_code,
});
// The install unloads the active model before the swap (even when the
// swap then fails); if a later gate cancels or the load fails, the UI
// must stop pointing at that unloaded model.
if (
useTransformersUpgradeDialogStore
.getState()
.consumeServerUnloadedChat()
&& currentStore.params.checkpoint
) {
upgradeUnloadedActive = true;
}
if (!upgraded) {
throw new Error(
`${modelDisplayName(sel.id)} needs a newer transformers release to load.`,
);
}
}
if (
validation.requires_trust_remote_code ||
validation.requires_security_review
@ -990,6 +1023,7 @@ export function SharedComposer({
tensor_parallel: currentStore.tensorParallel,
});
saveSpeculativeType(specSettings.speculativeType);
upgradeUnloadedActive = false;
const store = useChatRuntimeStore.getState();
store.setCheckpoint(
resp.model,
@ -1097,6 +1131,11 @@ export function SharedComposer({
toast.success("Compare complete", { id: toastId, duration: 2000 });
} catch (err) {
compareStepSucceededRef.current = false;
// The install already unloaded the previously active model; drop the
// checkpoint so the UI does not keep pointing at an unloaded model.
if (upgradeUnloadedActive) {
useChatRuntimeStore.getState().clearCheckpoint();
}
toast.error("Compare failed", {
id: toastId,
description: err instanceof Error ? err.message : "Unknown error",

6
studio/frontend/src/features/chat/types/api.ts
View file

@ -1,6 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { TransformersUpgradeInfo } from "@/features/transformers-upgrade";
export interface BackendModelDetails {
id: string;
name?: string | null;
@ -78,6 +80,10 @@ export interface ValidateModelResponse {
requires_security_review?: boolean;
/** Native context length from the local GGUF header; null until downloaded. */
context_length?: number | null;
/** Architecture only shipped by a newer transformers; UI pauses on the upgrade dialog. */
requires_transformers_upgrade?: boolean;
/** Set only when requires_transformers_upgrade. */
transformers_upgrade?: TransformersUpgradeInfo | null;
}
export interface GgufVariantDetail {

32
studio/frontend/src/features/transformers-upgrade/api/transformers-upgrade-api.ts Normal file
View file

@ -0,0 +1,32 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { readFastApiError } from "@/lib/format-fastapi-error";
interface InstallLatestTransformersResponse {
success: boolean;
version: string;
message: string;
/** The server unloaded the active chat model before the swap (set even on a
* structured failure, so callers can restore their model state). */
model_unloaded?: boolean;
/** On a version-mismatch failure: the release that superseded the requested
* one, so Retry can use it. */
latest_version?: string | null;
}
/** Consented install of the latest transformers into the sidecar; synchronous, can take minutes. */
export async function installLatestTransformers(
version: string,
): Promise<InstallLatestTransformersResponse> {
const response = await authFetch("/api/inference/install-latest-transformers", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ version }),
});
if (!response.ok) {
throw new Error(await readFastApiError(response));
}
return (await response.json()) as InstallLatestTransformersResponse;
}

167
studio/frontend/src/features/transformers-upgrade/components/transformers-upgrade-dialog.tsx Normal file
View file

@ -0,0 +1,167 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Spinner } from "@/components/ui/spinner";
import { cn } from "@/lib/utils";
import { PackageIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store";
function modelDisplayName(modelName: string | null): string {
if (!modelName) return "This model";
return modelName.split("/").pop() || modelName;
}
/** Root-mounted consent dialog for models needing a newer transformers;
* Install runs the sidecar install and resumes the paused load on success. */
export function TransformersUpgradeDialog() {
const open = useTransformersUpgradeDialogStore((s) => s.open);
const modelName = useTransformersUpgradeDialogStore((s) => s.modelName);
const upgrade = useTransformersUpgradeDialogStore((s) => s.upgrade);
const phase = useTransformersUpgradeDialogStore((s) => s.phase);
const errorMessage = useTransformersUpgradeDialogStore((s) => s.errorMessage);
const trustRemoteCodeFallback = useTransformersUpgradeDialogStore(
(s) => s.trustRemoteCodeFallback,
);
const install = useTransformersUpgradeDialogStore((s) => s.install);
const resolve = useTransformersUpgradeDialogStore((s) => s.resolve);
const displayName = modelDisplayName(modelName);
const modelType = upgrade?.model_type ?? "unknown";
const version = upgrade?.pypi_version ?? null;
// Only released PyPI versions are installable; dev (main) builds are never offered.
const installable = Boolean(upgrade?.supported_in_pypi && version);
const devOnly = !installable && Boolean(upgrade?.supported_in_main);
const installing = phase === "installing";
return (
<AlertDialog
open={open}
onOpenChange={(next) => {
// Escape/overlay dismiss must not abandon an in-flight install.
if (!next && !installing) resolve(false);
}}
>
<AlertDialogContent className="max-w-lg">
<AlertDialogHeader className="min-w-0">
<div className="flex w-full min-w-0 items-start gap-3">
<div className="flex size-9 shrink-0 items-center justify-center rounded-full bg-blue-500/10 text-blue-600 dark:text-blue-400">
<HugeiconsIcon icon={PackageIcon} className="size-5" />
</div>
<div className="min-w-0 flex-1 space-y-3">
<div className="space-y-1">
<AlertDialogTitle>New model architecture</AlertDialogTitle>
<AlertDialogDescription>
<span className="font-medium text-foreground">
{displayName}
</span>{" "}
uses the{" "}
<span className="font-mono text-foreground">{modelType}</span>{" "}
architecture, which your installed transformers does not
support yet.{" "}
{installable ? (
<>
Install transformers{" "}
<span className="font-medium text-foreground">
{version}
</span>{" "}
from PyPI to load it. The install runs once and can take
a minute; loading continues automatically afterwards.
</>
) : devOnly ? (
<>
Even the latest transformers release on PyPI does not
support it yet: the architecture is only available on the
transformers development branch (main), and Studio does
not install development builds. Support arrives with the
next transformers release on PyPI.
</>
) : (
<>
No released transformers version supports it yet, so it
cannot be loaded.
</>
)}
{!installable && trustRemoteCodeFallback ? (
<>
{" "}
This model also ships its own modeling code; you can
continue and review enabling that custom code instead.
</>
) : null}
</AlertDialogDescription>
</div>
{phase === "error" && errorMessage ? (
<p className="rounded-lg border border-red-500/30 bg-red-500/10 px-3 py-2 text-xs text-red-600 dark:text-red-400">
{errorMessage}
</p>
) : null}
{installing ? (
<p className="flex items-center gap-2 text-xs text-muted-foreground">
<Spinner className="size-3.5" />
Installing transformers {version}... This can take a minute.
</p>
) : null}
</div>
</div>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel disabled={installing}>Cancel</AlertDialogCancel>
{installable ? (
<>
{phase === "error" && trustRemoteCodeFallback ? (
// Install failed but the model ships custom code: offer the
// caller's trust_remote_code gate instead of forcing a retry.
<AlertDialogAction
className="bg-transparent text-foreground hover:bg-accent"
onClick={() => resolve(true)}
>
Continue with custom code
</AlertDialogAction>
) : null}
<AlertDialogAction
disabled={installing}
className={cn(installing && "pointer-events-none")}
onClick={(event) => {
// Keep the dialog open; the store closes it on success.
event.preventDefault();
void install();
}}
>
{installing ? (
<>
<Spinner className="size-4" />
Installing...
</>
) : phase === "error" ? (
"Retry install"
) : (
`Install transformers ${version}`
)}
</AlertDialogAction>
</>
) : trustRemoteCodeFallback ? (
// No installable release but the model ships custom code: continue
// into the caller's trust_remote_code gate as the last resort.
<AlertDialogAction onClick={() => resolve(true)}>
Continue with custom code
</AlertDialogAction>
) : null}
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}

28
studio/frontend/src/features/transformers-upgrade/hooks/use-transformers-upgrade-consent.ts Normal file
View file

@ -0,0 +1,28 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useTransformersUpgradeDialogStore } from "../stores/transformers-upgrade-dialog-store";
import type { TransformersUpgradeInfo } from "../types";
interface ConfirmArgs {
modelName: string;
/** validate's transformers_upgrade payload; null/undefined skips the dialog. */
upgrade: TransformersUpgradeInfo | null | undefined;
/** When no release is installable, offer continuing into the caller's custom-code gate. */
trustRemoteCodeFallback?: boolean;
}
/** Pause a load needing a newer transformers on the consent dialog and run the install.
* Resolves true when the load can continue; false on cancel or not-installable with no fallback. */
export async function confirmTransformersUpgradeIfNeeded({
modelName,
upgrade,
trustRemoteCodeFallback,
}: ConfirmArgs): Promise<boolean> {
if (!upgrade) return true;
return useTransformersUpgradeDialogStore
.getState()
.requestConsent(modelName, upgrade, {
trustRemoteCodeFallback: Boolean(trustRemoteCodeFallback),
});
}

8
studio/frontend/src/features/transformers-upgrade/index.ts Normal file
View file

@ -0,0 +1,8 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { TransformersUpgradeDialog } from "./components/transformers-upgrade-dialog";
export { confirmTransformersUpgradeIfNeeded } from "./hooks/use-transformers-upgrade-consent";
export { installLatestTransformers } from "./api/transformers-upgrade-api";
export { useTransformersUpgradeDialogStore } from "./stores/transformers-upgrade-dialog-store";
export type { TransformersUpgradeInfo } from "./types";

139
studio/frontend/src/features/transformers-upgrade/stores/transformers-upgrade-dialog-store.ts Normal file
View file

@ -0,0 +1,139 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { create } from "zustand";
import { installLatestTransformers } from "../api/transformers-upgrade-api";
import type { TransformersUpgradeInfo, TransformersUpgradePhase } from "../types";
type Resolver = (installed: boolean) => void;
// One in-flight consent; a new request resolves any prior pending one as declined.
let pendingResolver: Resolver | null = null;
interface TransformersUpgradeDialogStore {
open: boolean;
modelName: string | null;
upgrade: TransformersUpgradeInfo | null;
phase: TransformersUpgradePhase;
errorMessage: string | null;
/** Model ships custom code; without a PyPI install the load may fall back to trust_remote_code. */
trustRemoteCodeFallback: boolean;
/** True once this consent's install completed. The install unloads the previous
* model before swapping, so the caller must treat it as already unloaded; the
* custom-code fallback resolves true without installing and leaves it loaded. */
installRan: boolean;
/** True when the server unloaded the active chat model during this consent,
* including a swap that failed AFTER the unload: callers must then treat
* their previous model as gone and roll back on any later cancel. */
serverUnloadedChat: boolean;
/** Read-and-clear serverUnloadedChat: each waiter consumes the signal once,
* so a superseding consent can neither erase it before the old waiter reads
* it nor leak it into an unrelated later load. */
consumeServerUnloadedChat: () => boolean;
/** Open the dialog for a paused load; resolves true on install success or custom-code fallback. */
requestConsent: (
modelName: string,
upgrade: TransformersUpgradeInfo,
options?: { trustRemoteCodeFallback?: boolean },
) => Promise<boolean>;
/** Accept/Retry: run the install; on success resolve(true) and close. */
install: () => Promise<void>;
resolve: (installed: boolean) => void;
}
export const useTransformersUpgradeDialogStore =
create<TransformersUpgradeDialogStore>()((set, get) => ({
open: false,
modelName: null,
upgrade: null,
phase: "consent",
errorMessage: null,
trustRemoteCodeFallback: false,
installRan: false,
serverUnloadedChat: false,
requestConsent: (modelName, upgrade, options) =>
new Promise<boolean>((resolve) => {
pendingResolver?.(false);
pendingResolver = resolve;
set({
open: true,
modelName,
upgrade,
phase: "consent",
errorMessage: null,
trustRemoteCodeFallback: Boolean(options?.trustRemoteCodeFallback),
installRan: false,
});
}),
consumeServerUnloadedChat: () => {
const value = get().serverUnloadedChat;
if (value) set({ serverUnloadedChat: false });
return value;
},
install: async () => {
const { upgrade, phase } = get();
const version = upgrade?.pypi_version;
if (!version || phase === "installing") return;
const requestResolver = pendingResolver;
set({ phase: "installing", errorMessage: null });
let result: Awaited<ReturnType<typeof installLatestTransformers>>;
try {
result = await installLatestTransformers(version);
// Latch the server-side unload IMMEDIATELY, before any resolver-identity
// guard: even a superseded consent's install may have unloaded the chat
// model, and the signal must survive for whichever load consumes it next.
if (result.model_unloaded) {
set({ serverUnloadedChat: true });
}
} catch (error) {
// Ignore the failure if a newer request superseded this consent.
if (pendingResolver === requestResolver) {
set({
phase: "error",
errorMessage:
error instanceof Error && error.message
? error.message
: "Failed to install transformers.",
});
}
return;
}
if (pendingResolver === requestResolver) {
if (result.success) {
// serverUnloadedChat was latched above (and is never reset here): a
// retry after a failed-after-unload attempt reports false because the
// model is already gone, and a superseded install may have set it too.
set({ installRan: true });
get().resolve(true);
return;
}
// Structured failure: the swap failed but may have already unloaded the
// chat model; record that so a later cancel still rolls the caller back.
// A version mismatch also carries the superseding release, so Retry
// re-requests a version that can actually succeed.
const { upgrade } = get();
set({
phase: "error",
errorMessage: result.message || "Failed to install transformers.",
serverUnloadedChat:
get().serverUnloadedChat || Boolean(result.model_unloaded),
...(result.latest_version && upgrade
? { upgrade: { ...upgrade, pypi_version: result.latest_version } }
: {}),
});
}
},
resolve: (installed) => {
const resolver = pendingResolver;
pendingResolver = null;
set({
open: false,
modelName: null,
upgrade: null,
phase: "consent",
errorMessage: null,
trustRemoteCodeFallback: false,
});
resolver?.(installed);
},
}));

16
studio/frontend/src/features/transformers-upgrade/types.ts Normal file
View file

@ -0,0 +1,16 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/** Wire shape of `transformers_upgrade` from /api/inference/validate. */
export interface TransformersUpgradeInfo {
/** config.json model_type unknown to installed transformers. */
model_type: string;
/** Latest transformers release on PyPI at check time. */
pypi_version?: string | null;
/** Latest PyPI release ships this model_type (installable after consent). */
supported_in_pypi?: boolean;
/** Only transformers main ships it (dev-only; not installable). */
supported_in_main?: boolean;
}
export type TransformersUpgradePhase = "consent" | "installing" | "error";