Merge origin/main into studio-autoload-on-device-only

Takes main's wording for the GgufVariantDetail.partial doc comment,
which both sides added.
This commit is contained in:
Unsloth 2026-07-26 17:18:16 -07:00
commit b7b37e165c
62 changed files with 9917 additions and 436 deletions

View file

@ -655,6 +655,15 @@ _apt_distro_description() {
)
}
# ── Helper: can the controlling terminal actually be opened for reading? ──
# `test -r` only checks permission bits, which look fine in containers and
# systemd units where open() then fails with ENXIO. Probe with a real open.
# The subshell is required: in dash a failed redirection on the special
# builtin `:` exits the whole script.
_can_read_tty() {
( : </dev/tty ) >/dev/null 2>&1
}
# ── Helper: install packages via apt, escalating to sudo only if needed ──
# Usage: _smart_apt_install pkg1 pkg2 pkg3 ...
_smart_apt_install() {
@ -695,24 +704,63 @@ _smart_apt_install() {
echo " from your distro's official repositories (not a third-party tarball)."
echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
echo ""
printf " Accept? [Y/n] "
if [ -r /dev/tty ]; then
read -r REPLY </dev/tty || REPLY="y"
else
REPLY="y"
fi
case "$REPLY" in
[nN]*)
if _can_read_tty; then
printf " Accept? [Y/n] "
# The device opened, so a failed read is EOF, not consent: decline,
# as the autostart prompt below does. Enter is still yes (a
# successful read of an empty line).
read -r REPLY </dev/tty || REPLY="n"
case "$REPLY" in
[nN]*)
echo ""
echo " Please install these packages first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
;;
esac
# Mirror the headless branch: on a sudoers denial, a wrong password
# or an apt error, say what to run by hand instead of letting set -e
# abort on a bare sudo/apt message.
if sudo apt-get update -y </dev/null &&
sudo apt-get install -y $_STILL_MISSING </dev/null; then
:
else
echo ""
echo " Please install these packages first, then re-run Unsloth Studio setup:"
echo " Could not install these packages: $_STILL_MISSING"
echo " See the error above."
echo " Please install them first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
;;
*)
sudo apt-get update -y </dev/null
sudo apt-get install -y $_STILL_MISSING </dev/null
;;
esac
fi
else
# Nobody can answer a prompt or type a password here. -n makes sudo
# refuse rather than prompt into a closed stdin, which is how #7307
# died. Probe with the real commands: `sudo -l` answers whether they
# are *authorized*, not whether running them needs authentication.
# -k ignores any cached timestamp, so only a real NOPASSWD rule gets
# through, not someone's sudo in another shell minutes ago. Per
# sudo(8), -k alongside a command ignores the cached credentials and
# "will not update" them, so other sessions keep theirs.
echo " No terminal to confirm on; trying passwordless sudo."
if sudo -n -k apt-get update -y </dev/null &&
sudo -n -k apt-get install -y $_STILL_MISSING </dev/null; then
echo " Installed with passwordless sudo."
else
echo ""
echo " Could not install these packages: $_STILL_MISSING"
echo " Detected ${_ad_desc}."
# Either sudo refused, or apt failed on a bad repo, dpkg lock or
# network outage. sudo exits 1 on an auth/config problem and
# when the command cannot be executed, but otherwise passes the
# command's own status through, so state both causes.
echo " Either sudo needs a password here, or apt-get itself"
echo " failed; see the error above. With no terminal to"
echo " authenticate on, this cannot be done unattended."
echo " Please install them first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
exit 1
fi
fi
else
echo ""
echo " sudo is not available on this system."
@ -4055,9 +4103,11 @@ echo ""
# In non-interactive environments (Docker, CI, cloud-init) just print instructions.
if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then
echo ""
printf " Start Unsloth Studio now? [Y/n] "
# No readable answer (closed/EOF tty) defaults to no; Enter is still yes.
if [ -r /dev/tty ]; then
# Prompt only when something can answer: `test -r` passes on the unopenable
# /dev/tty found in containers, leaving a dangling question in the log.
if _can_read_tty; then
printf " Start Unsloth Studio now? [Y/n] "
read -r _reply </dev/tty || _reply="n"
else
_reply="n"

View file

@ -81,6 +81,82 @@ _PYTORCH_MISSING_MESSAGE = (
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
def _multi_gpu_device_map_kwargs() -> dict:
"""``device_map`` kwargs for sharding a checkpoint across every visible GPU.
unsloth's ``from_pretrained`` defaults to ``device_map="sequential"``, which stacks
the whole model on GPU0 and OOMs multi-GPU hosts whose other GPUs sit empty (#7053).
Returns ``{"device_map": "balanced"}`` only on a real multi-GPU CUDA/ROCm host
(mirroring the inference loader's ``get_device_map``), else empty so single-GPU, CPU
and MLX loads keep the loader default."""
if _IS_MLX:
return {}
try:
from utils.hardware import get_device_map, get_parent_visible_gpu_ids
visible = get_parent_visible_gpu_ids()
if len(visible) > 1:
device_map = get_device_map(visible)
elif not visible:
# UUID/MIG masks resolve to no numeric ids; get_device_map(None) falls back
# to the visible-GPU count, so a multi-GPU UUID/MIG host still shards.
device_map = get_device_map(None)
else:
return {}
if device_map == "balanced":
return {"device_map": device_map}
except Exception as exc:
logger.debug(f"multi-GPU device_map resolution failed; using loader default: {exc}")
return {}
def _is_oom_error(exc: BaseException) -> bool:
"""True for an accelerator OOM, however it is spelled.
accelerate and transformers re-raise it as a plain ``RuntimeError`` on several paths
and ROCm/XPU use their own classes, so match the message too.
"""
if torch is not None:
oom_types = tuple(
t
for t in (
getattr(torch, "OutOfMemoryError", None),
getattr(getattr(torch, "cuda", None), "OutOfMemoryError", None),
getattr(getattr(torch, "xpu", None), "OutOfMemoryError", None),
)
if isinstance(t, type)
)
if oom_types and isinstance(exc, oom_types):
return True
return "out of memory" in f"{type(exc).__name__}: {exc}".lower()
def _is_cpu_spill_rejection(exc: BaseException) -> bool:
"""bitsandbytes refuses a map that spills to CPU/disk with a plain ``ValueError``.
Busy secondary GPUs can make ``balanced`` spill to CPU even where the old sequential
load fit on GPU0, and that message says nothing about memory, so the retry has to
match it explicitly. See transformers ``quantizers/quantizer_bnb_4bit.py``.
"""
return "dispatched on the cpu or the disk" in str(exc).lower()
class _CpuSpillRetry(Exception):
"""A multi-GPU load that succeeded but left modules offloaded to CPU/disk."""
def _cpu_offloaded_modules(model) -> int:
"""Count the modules a load parked on CPU or disk.
Only bitsandbytes refuses such a map; a full-precision load accepts it, leaves the
parameters on meta and dies much later in safetensors with "Cannot copy out of meta
tensor". Nothing raises at load time, so inspect the map directly. PEFT re-dispatches
when attaching an adapter, so in practice this catches merged checkpoints.
"""
device_map = getattr(model, "hf_device_map", None) or {}
return sum(1 for target in device_map.values() if str(target) in ("cpu", "disk"))
def _supports_kwarg(fn, name):
"""True if `fn` accepts keyword `name` directly or via **kwargs."""
import inspect
@ -271,6 +347,7 @@ class ExportBackend:
load_in_4bit: bool = True,
trust_remote_code: bool = False,
hf_token: Optional[str] = None,
_device_map_override: Optional[dict] = None,
) -> Tuple[bool, str]:
"""
Load a checkpoint for export.
@ -303,6 +380,14 @@ class ExportBackend:
# Skip the Hub when offline so a no-internet export uses the local cache.
local_files_only = _hf_offline()
# Shard across every visible GPU instead of stacking on GPU0 (#7053); {} on
# single-GPU/CPU/MLX. _device_map_override is the single-device retry below.
_device_map_kw = (
_multi_gpu_device_map_kwargs()
if _device_map_override is None
else _device_map_override
)
# Run the type-detection probes in the forced-offline window (else a gated
# base 404s); it covers is_vision_model's Hub reads + the transformers-5
# subprocess, and local_files_only makes detect_audio_type's requests.get skip.
@ -328,6 +413,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "whisper":
@ -343,6 +429,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "snac":
@ -355,6 +442,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "bicodec":
@ -368,6 +456,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self._audio_type == "dac":
@ -380,6 +469,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
elif self.is_vision:
@ -392,6 +482,7 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
tokenizer = processor # vision: processor acts as tokenizer
@ -405,8 +496,16 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
**_device_map_kw,
)
# Only when we asked for the multi-GPU map: a single-GPU host has no second
# placement to retry on, so leave its behaviour untouched.
_offloaded = _cpu_offloaded_modules(model) if _device_map_kw else 0
if _device_map_override is None and _offloaded:
del model
raise _CpuSpillRetry(f"{_offloaded} module(s) offloaded to CPU/disk")
if _IS_MLX:
# MLX doesn't use PeftModel — detect LoRA via adapter_config.json
self.is_peft = adapter_config.exists()
@ -429,11 +528,41 @@ class ExportBackend:
return True, f"Loaded {model_type} model{peft_info} successfully"
except Exception as e:
logger.error(f"Error loading checkpoint: {e}")
import traceback
# Sharding is an optimisation, never a requirement. "balanced" budgets from the
# free memory read BEFORE this process opens a CUDA context on each GPU, so when
# a training or chat job already owns the others the shard can OOM, or spill to
# CPU and be refused by bitsandbytes, where the old single-device load succeeded.
# Fall back once before giving up.
if (
_device_map_override is None
and (
isinstance(e, _CpuSpillRetry) or _is_oom_error(e) or _is_cpu_spill_rejection(e)
)
and _multi_gpu_device_map_kwargs()
):
# Retry outside this block: the live traceback pins the half-built model's
# frames, so an in-block retry inherits the exhausted device.
retry_reason = str(e)
else:
logger.error(f"Error loading checkpoint: {e}")
import traceback
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
logger.error(traceback.format_exc())
return False, f"Failed to load checkpoint: {str(e)}"
logger.warning(
f"Multi-GPU export load unusable ({retry_reason}); retrying on "
f"the single-device loader default."
)
self.cleanup_memory()
return self.load_checkpoint(
checkpoint_path,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
hf_token = hf_token,
_device_map_override = {},
)
def _write_export_metadata(self, save_directory: str):
"""Write export_metadata.json with base model info for Chat page discovery."""

View file

@ -307,6 +307,9 @@ def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]":
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min
# A transport error can arrive before the child is reapable; a request path cannot
# afford the 5s the background MTP reload spends on the same race.
_RESPAWN_REAP_GRACE_S = 1.0
def _finalize_reasoning_only_cumulative(
@ -2099,6 +2102,9 @@ class LlamaCppBackend:
# Serialises mid-session respawns so many generations hitting a killed
# server trigger at most one reload (see _respawn_if_dead).
self._respawn_lock = threading.Lock()
# Bumped by every unload. load_model clears _cancel_event, so a respawn that
# raced an unload needs a signal that survives the clear (see _respawn_if_dead).
self._unload_epoch = 0
# Set by the in-app updater while it swaps prebuilt binaries; load_model()
# rejects fast so no server starts from a half-swapped binary.
self._llama_update_in_progress = False
@ -9308,6 +9314,7 @@ class LlamaCppBackend:
"""Terminate the subprocess and cancel any in-flight download."""
self._cancel_event.set()
with self._lock:
self._unload_epoch += 1
self._kill_process()
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
self._model_identifier = None
@ -10107,15 +10114,18 @@ class LlamaCppBackend:
return False
if not self._mtp_runtime_fallback_active:
return False
if not self._last_load_kwargs or self._process is None:
# Read before claiming: a raise after the claim strands the flag, and nothing
# else clears it, blocking every later respawn.
kwargs = self._last_load_kwargs
proc = self._process
if not kwargs or proc is None:
return False
# Single-flight: the first failure claims the reload.
with self._mtp_runtime_fallback_lock:
if self._mtp_runtime_fallback_in_progress:
return False
self._mtp_runtime_fallback_in_progress = True
snapshot = dict(self._last_load_kwargs)
proc = self._process
snapshot = dict(kwargs)
def _recover():
try:
@ -10163,7 +10173,14 @@ class LlamaCppBackend:
with self._mtp_runtime_fallback_lock:
self._mtp_runtime_fallback_in_progress = False
threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
try:
threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
except RuntimeError as exc:
# Release the claim: a reload that never started would block respawn forever.
with self._mtp_runtime_fallback_lock:
self._mtp_runtime_fallback_in_progress = False
logger.error(f"Could not start the MTP-crash reload: {exc}")
return False
return True
def _start_mtp_crash_watchdog(self) -> None:
@ -10635,6 +10652,21 @@ class LlamaCppBackend:
finally:
_cancel_closed.set()
def _server_socket_is_open(self, timeout_s: float = 0.15) -> bool:
"""True if anything still accepts on the server port.
The listening socket dies with the process, so this tells a live server
from a dead one without waiting for the child to become reapable.
"""
port = self._port
if not port:
return False
try:
with socket.create_connection(("127.0.0.1", port), timeout = timeout_s):
return True
except OSError:
return False
def _respawn_if_dead(self) -> bool:
"""Relaunch the llama-server if its process has exited.
@ -10644,28 +10676,114 @@ class LlamaCppBackend:
recover, returning True once healthy. Serialised on ``_respawn_lock`` so
many generations hitting the dead server trigger at most one reload.
"""
# Read outside the lock so a queued caller can tell the replacement from the child
# its own error came from; otherwise each burns the grace wait below, and that
# sleep is held under the lock, so the waits serialise.
served_by = self._process
with self._respawn_lock:
proc = self._process
if proc is None:
return False
if proc.poll() is None:
# Process is alive: either a concurrent caller already respawned
# it (healthy), or this connection error wasn't a dead server.
if self._cancel_event.is_set():
# unload_model sets this before it kills, so the child can still be
# accepting. Reporting it healthy would aim the retry at a server
# that is deliberately going away.
return False
if proc is not served_by:
# Replaced while we queued: this child never served our request.
return self._healthy
kwargs = self._last_load_kwargs
if not kwargs:
return False
logger.warning(
f"llama-server for '{self._model_identifier}' exited "
f"(code {proc.returncode}); respawning to recover the session"
)
with self._lock:
self._healthy = False
if proc.poll() is None:
# Still serving, so the error was transient. Charging it the grace below
# would cost a second per caller, serialised under this lock.
if self._server_socket_is_open():
return self._healthy
# A closing server can beat its own exit status: calling it alive returns
# the stale _healthy and spends the retry on the corpse.
deadline = time.monotonic() + _RESPAWN_REAP_GRACE_S
while proc.poll() is None and time.monotonic() < deadline:
time.sleep(0.05)
if proc.poll() is None:
# Alive: either a concurrent caller already respawned it (healthy), or
# this connection error wasn't a dead server.
return self._healthy
with self._mtp_runtime_fallback_lock:
if self._mtp_runtime_fallback_in_progress:
# An MTP-free reload owns this corpse; replaying the old kwargs
# restarts the crashing config and aborts that reload.
logger.info("Respawn skipped: an MTP-free reload is already recovering.")
return False
# The RLock lets the load_model below re-enter it.
with self._serial_load_lock:
if self._process is not proc:
logger.info("Respawn skipped: a newer load is already active.")
return self._healthy
# Snapshot under _lock, the one unload_model holds, so a teardown is
# either wholly before us (flag set) or wholly after (epoch bumped).
# _serial_load_lock alone would not exclude it: unload never takes it.
with self._lock:
if self._cancel_event.is_set():
logger.info("Respawn skipped: the model was unloaded.")
return False
kwargs = dict(self._last_load_kwargs or {})
if not kwargs:
return False
epoch = self._unload_epoch
self._healthy = False
logger.warning(
f"llama-server for '{self._model_identifier}' exited "
f"(code {proc.returncode}); respawning to recover the session"
)
try:
started = bool(self.load_model(**kwargs))
except Exception as exc:
logger.error(f"Failed to respawn llama-server: {exc}")
return False
if started and self._unload_epoch != epoch:
# An unload landed mid-reload. load_model cleared _cancel_event on
# the way in, so the epoch is the only surviving evidence; undo the
# replacement rather than leave a model the user stopped running.
logger.info("Respawn undone: the model was unloaded during the reload.")
self.unload_model()
return False
return started
@contextlib.contextmanager
def _open_chat_stream_with_respawn_retry(self, payload: dict, cancel_event):
"""Open a chat stream, respawning a dead llama-server once before streaming.
Retry only when opening the response fails: once it is open a consumer may
already have emitted content or tool events, so a replay could duplicate
output and side effects. ``base_url`` is resolved per attempt because a
respawn may pick a new port. The budget is one retry per model request, not
per chat turn, so a long tool loop never discards a completed tool.
A child dying after the accept but before the headers surfaces as
ReadError/WriteError/RemoteProtocolError rather than ConnectError, and which
one differs per OS. llama-server flushes its 200 at slot start, so that window
is an upload still in flight or a request behind busy slots; a death during
decode arrives with the response open and is not replayed. Timeouts are
excluded: the server is slow, not dead, and a replay would spend the
first-token budget twice.
"""
for attempt in range(2):
response_opened = False
try:
return bool(self.load_model(**kwargs))
except Exception as exc:
logger.error(f"Failed to respawn llama-server: {exc}")
return False
url = f"{self.base_url}/v1/chat/completions"
with self._open_stream(url, payload, cancel_event) as opened:
response_opened = True
yield opened
return
except (httpx.NetworkError, httpx.RemoteProtocolError) as exc:
if response_opened:
raise
if self._maybe_recover_from_mtp_crash(exc):
raise RuntimeError("Lost connection to llama-server") from exc
if attempt == 0 and self._respawn_if_dead():
logger.warning(
"llama-server was unreachable; respawned it and retrying the generation"
)
continue
raise
def generate_chat_completion(
self,
@ -10931,16 +11049,20 @@ class LlamaCppBackend:
build_rag_autoinject,
execute_tool,
is_always_safe_tool,
is_potentially_unsafe_tool_call,
is_high_risk_tool_call,
)
# Normalize the mode: "full" and bypass_permissions are the same
# switch, whichever arrives first wins toward the permissive side.
# "off" keeps the sandbox but never prompts.
# "full" and bypass_permissions are the same switch, whichever arrives
# first wins. "off" keeps the sandbox but never prompts. Unset defaults to
# "auto"; unknown falls back to the stricter "ask". An explicit
# confirm_tool_calls=True with no mode is already resolved to "ask" at the
# request layer, so it never arrives here as an ambiguous unset.
if permission_mode == "full":
bypass_permissions = True
elif bypass_permissions:
permission_mode = "full"
elif permission_mode is None:
permission_mode = "auto"
elif permission_mode not in ("ask", "auto", "off"):
permission_mode = "ask"
@ -10963,7 +11085,6 @@ class LlamaCppBackend:
yield _ev
conversation.extend(_auto["messages"])
url = f"{self.base_url}/v1/chat/completions"
_accumulated_completion_tokens = 0
_accumulated_predicted_ms = 0.0
_accumulated_predicted_n = 0
@ -11223,7 +11344,7 @@ class LlamaCppBackend:
_text_args_name = ""
_confirm_gated_iteration = bool(confirm_tool_calls) and not bypass_permissions
with self._open_stream(url, payload, cancel_event) as (
with self._open_chat_stream_with_respawn_retry(payload, cancel_event) as (
response,
first_token_deadline,
):
@ -12035,18 +12156,16 @@ class LlamaCppBackend:
decision.as_assistant_tool_call()
)
# Bypass wins over the confirm gate at the loop level too,
# so a direct internal caller with both flags never prompts.
# In "auto" mode only calls detected as potentially unsafe
# pause; read-only calls run straight through. "off" never
# prompts (sandbox stays on).
# Bypass wins here too, so a direct internal caller with both
# flags never prompts. "auto" pauses only high-risk calls;
# "off" never prompts (sandbox stays on).
needs_confirm = (
bool(confirm_tool_calls)
and not bypass_permissions
and permission_mode != "off"
)
if needs_confirm and permission_mode == "auto":
needs_confirm = is_potentially_unsafe_tool_call(
needs_confirm = is_high_risk_tool_call(
decision.tool_name, decision.arguments
)
approval_id = new_approval_id() if needs_confirm else ""
@ -12260,7 +12379,7 @@ class LlamaCppBackend:
_stream_done = False
try:
with self._open_stream(url, stream_payload, cancel_event) as (
with self._open_chat_stream_with_respawn_retry(stream_payload, cancel_event) as (
response,
first_token_deadline,
):

View file

@ -147,10 +147,16 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
)
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs
from utils.hf_cache_settings import known_hf_hub_caches
from core.inference.model_ids import public_model_id
index: dict[str, _LocalGgufEntry] = {}
seen_hf: set[str] = set()
try:
active_root = str(Path(_resolve_hf_cache_dir()).resolve())
except Exception:
active_root = None
def _scan_hf_once(directory) -> list:
if directory is None:
return []
@ -162,7 +168,13 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
if rp in seen_hf:
return []
seen_hf.add(rp)
return _scan_hf_cache(directory)
# Only the active cache loads by repo id. Say so, or an inactive repo is
# indexed under an id it cannot load by, and its snapshot basename (what
# /v1/models advertises once loaded by path) is never a key at all.
# No format classification here: nothing on this path reads model_format,
# and its recursive walk would duplicate the one _local_gguf_entry already
# does per snapshot, on the request path.
return _scan_hf_cache(directory, active_cache = rp == active_root, classify_format = False)
except Exception as exc: # a missing/malformed root must skip, never crash the index
logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc)
return []
@ -220,12 +232,61 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
continue
# Index every alias (including the path) so a client can resolve by any of
# them, even though only the non-path loader_id is advertised.
for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)):
for key in (
raw_id,
getattr(info, "model_id", None),
getattr(info, "display_name", None),
public_model_id(raw_id),
):
if key:
index.setdefault(key.strip().lower(), entry)
# Other revisions of the same repo resolve to their own weights, so a pin on
# one keeps working after Hugging Face writes a newer snapshot.
for name, sibling_entry in _sibling_revision_entries(raw_id, loader_id):
index.setdefault(name.strip().lower(), sibling_entry)
return index
def _sibling_revision_entries(raw_id: str, loader_id: str):
"""Yield ``(revision_name, entry)`` for the repo's OTHER cached revisions.
An inactive-cache repo carries its snapshot path as the id, and /v1/models
advertises only that directory's basename once loaded, so anything durable
pinned to it (a subagent config) holds one revision hash. Hugging Face writes a
new snapshot dir on every update, and the scan emits a single entry per repo
pointed at the newest one, so that pin would otherwise stop resolving and drop
through to whatever model is loaded.
Each revision gets an entry for its OWN directory rather than an alias onto the
scanned one: aliasing would redirect a pin that names an older complete revision
onto a newer half-downloaded snapshot and break a request that works today.
Incomplete revisions are skipped for the same reason.
Sibling names are only revisions inside a real cache repo
(``<root>/models--org--name/snapshots/<rev>``). A scan folder that merely happens
to be called ``snapshots`` holds unrelated models, and treating those as
revisions would silently serve one model in place of another.
"""
from pathlib import Path
from types import SimpleNamespace
snapshots = Path(raw_id).parent
if snapshots.name != "snapshots" or not snapshots.parent.name.startswith("models--"):
return
from routes.models import snapshot_variants_all_complete
try:
siblings = [p for p in snapshots.iterdir() if p.is_dir() and p.name != Path(raw_id).name]
except OSError:
return
for sibling in siblings:
if not snapshot_variants_all_complete(str(sibling)):
continue
entry = _local_gguf_entry(loader_id, SimpleNamespace(path = str(sibling)))
if entry is not None:
yield sibling.name, entry
def _index() -> dict[str, _LocalGgufEntry]:
global _scan
# Build under the lock so concurrent callers with an expired cache don't all

View file

@ -514,13 +514,17 @@ def run_safetensors_tool_loop(
"""
conversation = list(messages)
# Normalize the mode (mirrors the GGUF loop): "full" and
# bypass_permissions are the same switch; unset/unknown behaves as "ask".
# "off" keeps the sandbox but never prompts.
# Mirrors the GGUF loop: "full" and bypass_permissions are the same switch;
# unset defaults to "auto", unknown falls back to the stricter "ask"; "off"
# keeps the sandbox but never prompts. An explicit confirm_tool_calls=True with
# no mode is already resolved to "ask" at the request layer, so it never
# arrives here as an ambiguous unset.
if permission_mode == "full":
bypass_permissions = True
elif bypass_permissions:
permission_mode = "full"
elif permission_mode is None:
permission_mode = "auto"
elif permission_mode not in ("ask", "auto", "off"):
permission_mode = "ask"
@ -1189,18 +1193,15 @@ def run_safetensors_tool_loop(
else:
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
# Bypass wins over the confirm gate at the loop level too, so a
# direct internal caller passing both flags never prompts. In
# "auto" mode only calls detected as potentially unsafe pause.
# "off" never prompts (sandbox stays on).
# Bypass wins here too, so a direct internal caller with both flags
# never prompts. "auto" pauses only high-risk calls; "off" never
# prompts (sandbox stays on).
needs_confirm = (
bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off"
)
if needs_confirm and permission_mode == "auto":
from core.inference.tools import is_potentially_unsafe_tool_call
needs_confirm = is_potentially_unsafe_tool_call(
decision.tool_name, decision.arguments
)
from core.inference.tools import is_high_risk_tool_call
needs_confirm = is_high_risk_tool_call(decision.tool_name, decision.arguments)
approval_id = new_approval_id() if needs_confirm else ""
decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None
start_event = decision.tool_start_event()

File diff suppressed because it is too large Load diff

View file

@ -919,11 +919,11 @@ class ThinkingConfig(BaseModel):
# Recognized permission_mode values. The field accepts a plain string rather than
# a Literal so an unrecognized value from a newer UI/client degrades to the
# safest gate ("ask") instead of a 422; the tool loops apply the same unknown ->
# ask fallback, so normalizing here keeps that forward-compat path reachable at
# the API boundary. None stays unset ("behaves as 'ask'" without self-enabling
# the confirm gate).
# a Literal so an unrecognized value from a newer UI/client degrades to the safest
# gate ("ask") instead of a 422. None stays unset at the request boundary: the tool
# loops normalize it to the product default "auto", while the route's confirm-gate
# derivation keeps an unset mode lenient (a non-streaming request cannot prompt, so
# it runs) to keep non-streaming clients and health checks working.
_KNOWN_PERMISSION_MODES = ("ask", "auto", "off", "full")
@ -1086,11 +1086,13 @@ class ChatCompletionRequest(BaseModel):
"[x-unsloth] Permission level for local tool calls. 'ask' pauses every "
"call for approval; 'ask'/'auto' enable the confirmation gate on their "
"own (needs a streaming request to deliver prompts). 'auto' ('Approve for "
"me') only pauses calls detected as potentially unsafe (state-mutating "
"terminal/python/MCP calls); read-only calls run immediately, and the "
"sandbox stays on. 'full' is equivalent to bypass_permissions=true (no "
"confirmation, no sandbox). Unset behaves as 'ask'. An unrecognized value "
"(e.g. from a newer client) is treated as 'ask'."
"me') only pauses calls detected as high risk (credential reads, privilege "
"escalation, destructive/persistence, network exfil); ordinary calls run "
"immediately, and the sandbox stays on. 'full' is equivalent to "
"bypass_permissions=true (no confirmation, no sandbox). Unset defaults to "
"'auto' for the per-call gate; a non-streaming request without an explicit "
"mode cannot prompt and runs the loop. An unrecognized value (e.g. from a "
"newer client) is treated as 'ask'."
),
)
auto_heal_tool_calls: Optional[bool] = Field(
@ -1376,6 +1378,21 @@ class ChatCompletionRequest(BaseModel):
elif self.permission_mode == "off":
# "Off" never prompts, so route guards must see confirm disabled.
self.confirm_tool_calls = False
elif (
self.permission_mode is None
and self.confirm_tool_calls is True
and not (self.provider_id or self.provider_type)
):
# An explicit confirm_tool_calls=True with no mode opted into the
# pre-permission-mode contract of gating every call, so resolve it to
# "ask" rather than let the loop apply the "auto" default, which would
# silently weaken that opt-in to high-risk calls only. Unlike the "ask"
# branch below this only sets permission_mode, which is inert unless
# Unsloth's own tool loop runs, so it needs no enable_tools/mcp gate --
# deliberate, since a process-wide --enable-tools policy can force the
# loop when the request sets neither flag. A bare unset request
# (confirm_tool_calls is None) still defaults to auto.
self.permission_mode = "ask"
elif (
self.permission_mode == "ask"
and self.confirm_tool_calls is None
@ -2059,7 +2076,7 @@ class AnthropicMessagesRequest(BaseModel):
)
permission_mode: Optional[str] = Field(
None,
description = "[x-unsloth] Permission level for local tool calls: 'ask' pauses every call, 'auto' only pauses calls detected as potentially unsafe, 'off' never pauses (sandbox stays on), 'full' equals bypass_permissions=true. Unset behaves as 'ask'; an unrecognized value (e.g. from a newer client) is treated as 'ask'. Declared explicitly so omitted requests default to None instead of raising AttributeError.",
description = "[x-unsloth] Permission level for local tool calls: 'ask' pauses every call, 'auto' ('Approve for me') only pauses calls detected as high risk, 'off' never pauses (sandbox stays on), 'full' equals bypass_permissions=true. Unset defaults to 'auto' for the per-call gate; a non-streaming request without an explicit mode runs the loop. An unrecognized value (e.g. from a newer client) is treated as 'ask'. Declared explicitly so omitted requests default to None instead of raising AttributeError.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,

View file

@ -143,6 +143,12 @@ class GgufVariantDetail(BaseModel):
update_available: bool = Field(
False, description = "Whether a newer version of this variant is available on HF"
)
partial: bool = Field(
False,
description = "Whether this variant is an interrupted download. The hub service "
"already computes it; carry it through so callers can hide a quant whose shards "
"are incomplete instead of offering one that cannot load.",
)
class GgufVariantsResponse(BaseModel):

View file

@ -2137,14 +2137,13 @@ def _explicit_studio_tool_loop_requested(payload) -> bool:
def _permission_mode_confirm(payload) -> bool:
"""Effective confirm-gate intent for Unsloth's own local tool loop.
Honors the documented default that an unset permission_mode behaves as
"ask". An explicit confirm_tool_calls (True or False) wins; explicit
ask/auto always engage the gate (a non-streaming one is then rejected, since
it cannot prompt); off/full never prompt. An unset mode defaults to ask, but
that is only realizable on a streaming request, so a non-streaming unset
request keeps the legacy run-without-gate behavior instead of 400ing. Used
at the pre-switch guard and the per-backend tool paths so a forced tool loop
(CLI --enable-tools) with the default mode still gates streaming requests.
An explicit confirm_tool_calls (True or False) wins; explicit ask/auto always
engage the gate (a non-streaming one is then rejected, since it cannot prompt);
off/full never prompt. An unset mode stays lenient here even though the loop
defaults it to "auto": a non-streaming request keeps the legacy
run-without-gate behavior instead of 400ing, so non-streaming clients and
health checks keep working. Used at the pre-switch guard and the per-backend
tool paths so a forced tool loop (CLI --enable-tools) still gates streaming.
"""
if payload.confirm_tool_calls is not None:
return bool(payload.confirm_tool_calls)

View file

@ -314,7 +314,11 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
try:
if not child.is_dir():
continue
has_gguf = any(child.glob("*.gguf"))
gguf_names = [p.name for p in child.glob("*.gguf")]
has_gguf = bool(gguf_names)
# mmproj alone is a vision adapter, not servable weights, so it decides
# presence but never format (same rule as _dir_model_format).
has_main_gguf = any(_is_main_gguf_filename(n) for n in gguf_names)
has_non_gguf_weights = _has_non_gguf_weights(child)
has_config = (child / "config.json").exists() or (
child / "adapter_config.json"
@ -332,7 +336,7 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
# A folder whose only weights are .gguf is GGUF-format even when it also
# ships a config.json (common for HF GGUF repos); such folders often lack
# a -GGUF suffix, so surface the format for the UI's GGUF classification.
model_format = "gguf" if has_gguf and not has_non_gguf_weights else None
model_format = "gguf" if has_main_gguf and not has_non_gguf_weights else None
found.append(
LocalModelInfo(
id = str(child),
@ -348,7 +352,8 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
for gguf_file in models_dir.glob("*.gguf"):
if limit is not None and len(found) >= limit:
break
if gguf_file.is_file():
# A standalone mmproj is a vision adapter, not servable weights.
if gguf_file.is_file() and _is_main_gguf_filename(gguf_file.name):
try:
updated_at = gguf_file.stat().st_mtime
except OSError:
@ -367,7 +372,12 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca
return found
def _scan_hf_cache(cache_dir: Path, *, active_cache: bool = True) -> List[LocalModelInfo]:
def _scan_hf_cache(
cache_dir: Path,
*,
active_cache: bool = True,
classify_format: bool = True,
) -> List[LocalModelInfo]:
if not cache_dir.exists() or not cache_dir.is_dir():
return []
@ -392,13 +402,23 @@ def _scan_hf_cache(cache_dir: Path, *, active_cache: bool = True) -> List[LocalM
partial = partial or hf_cache_scan.is_gguf_repo_partial(model_id, repo_dir)
load_id = model_id
snapshot = _resolve_hf_cache_realpath(repo_dir)
if not active_cache:
load_id = _resolve_hf_cache_realpath(repo_dir) or str(repo_dir.resolve())
load_id = snapshot or str(repo_dir.resolve())
# Classify from the snapshot's own weights. A GGUF repo without a -GGUF
# suffix is common, and leaving this unset makes every consumer guess from
# the name; the snapshot is already resolved just above.
model_format = (
_dir_model_format(Path(snapshot), recursive = True)
if snapshot and classify_format
else None
)
found.append(
LocalModelInfo(
id = load_id,
model_id = model_id,
display_name = model_id.split("/")[-1],
model_format = model_format,
path = load_id if not active_cache else str(repo_dir),
source = "hf_cache",
active_cache = active_cache,
@ -409,16 +429,30 @@ def _scan_hf_cache(cache_dir: Path, *, active_cache: bool = True) -> List[LocalM
return found
def _dir_model_format(path: Path) -> Optional[str]:
def _dir_model_format(path: Path, recursive: bool = False) -> Optional[str]:
"""Return ``"gguf"`` for a directory whose only weights are ``.gguf`` files.
LM Studio and custom GGUF folders frequently lack a ``-GGUF`` name suffix,
so the UI relies on this hint to route them through the GGUF load path
rather than treating them as plain local checkpoints.
rather than treating them as plain local checkpoints. A directory whose only
``.gguf`` is an mmproj vision adapter is not one: the variant selector drops
mmproj, so that path would find nothing to serve.
``recursive`` is for HF cache snapshots, which keep split quants in per-quant
subdirectories: a flat glob sees no ``.gguf`` there and would report the
snapshot as non-GGUF, hiding every sharded repo from the GGUF pickers. It looks
one level down rather than walking the tree, because that is where split quants
live and ``/api/models/local`` is async: an unbounded ``rglob`` per repo would
have to exhaust every non-GGUF snapshot before concluding there is no GGUF,
blocking the event loop on a large cache.
"""
try:
if not any(path.glob("*.gguf")):
return None
found = path.glob("*.gguf")
if not any(_is_main_gguf_filename(p.name) for p in found):
if not recursive:
return None
if not any(_is_main_gguf_filename(p.name) for p in path.glob("*/*.gguf")):
return None
return None if _has_non_gguf_weights(path) else "gguf"
except OSError:
return None
@ -455,7 +489,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
for child in lm_dir.iterdir():
try:
if not child.is_dir():
if child.suffix == ".gguf" and child.is_file():
if _is_main_gguf_filename(child.name) and child.is_file():
try:
updated_at = child.stat().st_mtime
except OSError:
@ -518,7 +552,7 @@ def _scan_lmstudio_dir(lm_dir: Path) -> List[LocalModelInfo]:
updated_at = updated_at,
),
)
elif model_dir.suffix == ".gguf" and model_dir.is_file():
elif _is_main_gguf_filename(model_dir.name) and model_dir.is_file():
try:
updated_at = model_dir.stat().st_mtime
except OSError:
@ -2792,6 +2826,7 @@ async def get_gguf_variants(
),
downloaded = bool(v.downloaded),
update_available = bool(getattr(v, "update_available", False)),
partial = bool(getattr(v, "partial", False)),
)
for v in response.variants
],
@ -3016,11 +3051,80 @@ def _repo_gguf_last_modified(repo_info) -> float:
return latest
def snapshot_variants_all_complete(snapshot: str) -> bool:
"""True when every quant the variant lister would advertise from *snapshot* is
fully on disk.
One complete quant is not enough: the picker enumerates the whole directory, so a
half-downloaded split quant sitting beside a good one still gets offered and the
generated command asks llama-server for shards that are absent. Both sides derive
their labels from ``extract_quant_label`` over paths relative to the snapshot, so
the sets are directly comparable.
"""
from hub.utils import inventory_scan
from hub.utils.gguf import list_local_gguf_variants
try:
variants, _ = list_local_gguf_variants(snapshot)
offered = {v.quant for v in variants if getattr(v, "quant", None)}
if not offered:
return False
return offered <= inventory_scan._completed_gguf_variants(Path(snapshot))
except Exception:
return False
def _repo_gguf_load_id(repo_info, active_root: Optional[Path]) -> Optional[str]:
"""Snapshot dir holding the newest primary GGUF, for a repo outside the active
hub cache that does not resolve by id. ``None`` when the id works or no
snapshot is recorded, since the repo dir itself is not loadable.
"""
repo_path = getattr(repo_info, "repo_path", None)
if repo_path is None or active_root is None:
return None
try:
if repo_path.parent.resolve(strict = False) == active_root:
return None
except (OSError, RuntimeError, ValueError):
pass
# Order by snapshot directory mtime, matching hub.utils.gguf.iter_hf_cache_snapshots,
# which is what variant discovery reads. Blob mtimes would disagree with it whenever
# Hugging Face reuses an older blob in a newer snapshot, and the command would then
# name a snapshot that does not hold the quant the picker offered.
candidates: List[tuple[float, str]] = []
for revision in repo_info.revisions:
snapshot = getattr(revision, "snapshot_path", None)
if snapshot is None:
continue
if not any(_is_main_gguf_filename(f.file_name) for f in revision.files):
continue
try:
mtime = Path(snapshot).stat().st_mtime
except OSError:
mtime = 0.0
candidates.append((mtime, str(snapshot)))
candidates.sort(key = lambda c: c[0], reverse = True)
# Newest first, but skip one holding only part of a split quant: an interrupted
# download would otherwise beat an older snapshot that can still load. Scanning
# stops at the first usable snapshot, so the usual case walks one directory.
for _, snapshot in candidates:
if snapshot_variants_all_complete(snapshot):
return snapshot
# Nothing complete anywhere: publishing a half-downloaded snapshot would put that
# path in the copied command and fail on load. Drop the id so the repo id is used,
# which fetches the missing shards instead.
return None
@router.get("/cached-gguf")
async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
"""List GGUF repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
try:
cache_scans = _all_hf_cache_scans()
try:
active_root = _resolve_hf_cache_dir().resolve(strict = False)
except Exception:
active_root = None
seen_lower: dict[str, dict] = {}
for hf_cache in cache_scans:
@ -3046,6 +3150,9 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
"cache_path": str(repo_info.repo_path),
"has_vision": _repo_has_mmproj(repo_info),
}
load_id = _repo_gguf_load_id(repo_info, active_root)
if load_id:
row["load_id"] = load_id
# Keep the newest timestamp across duplicate caches;
# attach only when known so absent rows sort as oldest.
lm = max(last_modified, (existing or {}).get("last_modified", 0.0))

View file

@ -663,9 +663,9 @@ def test_bypass_env_does_not_add_unset_windows_profile_vars(monkeypatch, tmp_pat
@_POSIX_ONLY
def test_bypass_exec_hardens_parent_proc_env(monkeypatch, captured_popen):
# Stripping the child env is not enough: a same-UID child can read the
# parent's /proc environ. The exec paths must invoke the parent hardening
# when (and only when) the sandbox is disabled.
# Stripping the child env is not enough: a same-UID child can read the parent's
# /proc environ. Both exec paths harden the parent in bypass mode (fail closed)
# and in sandboxed mode too (best-effort backstop for a classifier miss).
calls = {"n": 0}
def fake_harden():
@ -680,7 +680,7 @@ def test_bypass_exec_hardens_parent_proc_env(monkeypatch, captured_popen):
calls["n"] = 0
_python_exec("print(1)", None, 5, "t", disable_sandbox = False)
_bash_exec("echo hi", None, 5, "t", disable_sandbox = False)
assert calls["n"] == 0 # never hardened on the sandboxed path
assert calls["n"] == 2 # sandboxed path now hardens too (best-effort)
def test_bypass_exec_fails_closed_when_hardening_fails(monkeypatch, captured_popen):

View file

@ -126,6 +126,185 @@ def test_collect_local_models_prefers_complete_previous_copy(monkeypatch, tmp_pa
assert row.active_cache is False
def test_list_cached_gguf_reports_snapshot_load_id_for_inactive_cache(monkeypatch, tmp_path):
"""Only a repo outside the active cache needs a snapshot load_id."""
active = tmp_path / "active"
snapshot = tmp_path / "legacy" / "models--Org--Away" / "snapshots" / "rev"
snapshot.mkdir(parents = True)
(snapshot / "Q4_K_M.gguf").write_bytes(b"\0")
away = _repo(
"Org/Away",
[],
tmp_path / "legacy" / "models--Org--Away",
revisions = [
SimpleNamespace(files = [_file("Q4_K_M.gguf", 5_000)], snapshot_path = snapshot),
],
)
here = _repo("Org/Here", [_file("Q4_K_M.gguf", 6_000)], active / "models--Org--Here")
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [away, here])]
)
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active)
rows = {
c["repo_id"]: c
for c in asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))["cached"]
}
assert rows["Org/Away"]["load_id"] == str(snapshot)
assert "load_id" not in rows["Org/Here"]
def test_list_cached_gguf_load_id_follows_snapshot_dir_mtime(monkeypatch, tmp_path):
"""Pick the snapshot variant discovery reads: newest directory, not newest blob."""
import os
active = tmp_path / "active"
repo_dir = tmp_path / "legacy" / "models--Org--Multi"
older, newer = repo_dir / "snapshots" / "rev-a", repo_dir / "snapshots" / "rev-b"
for path in (older, newer):
path.mkdir(parents = True)
(older / "Q4_K_M.gguf").write_bytes(b"\0")
(newer / "Q8_0.gguf").write_bytes(b"\0")
os.utime(older, (1_000, 1_000))
os.utime(newer, (2_000, 2_000))
repo = _repo(
"Org/Multi",
[],
repo_dir,
revisions = [
# The older directory holds the newer blob, which is what diverges.
SimpleNamespace(
files = [_file("Q4_K_M.gguf", 5_000, blob_path = "b1")], snapshot_path = older
),
SimpleNamespace(files = [_file("Q8_0.gguf", 6_000, blob_path = "b2")], snapshot_path = newer),
],
)
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])]
)
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active)
monkeypatch.setattr(
models_route, "_blob_mtime", lambda f: 9_000 if f.blob_path == "b1" else 1.0
)
rows = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))["cached"]
assert rows[0]["load_id"] == str(newer)
def test_list_cached_gguf_load_id_skips_partial_split_snapshot(monkeypatch, tmp_path):
"""A half-downloaded split quant must not beat an older snapshot that can load."""
import os
active = tmp_path / "active"
repo_dir = tmp_path / "legacy" / "models--Org--Split"
older, newer = repo_dir / "snapshots" / "rev-a", repo_dir / "snapshots" / "rev-b"
for path in (older, newer):
path.mkdir(parents = True)
(older / "Model-Q8_0.gguf").write_bytes(b"\0")
# Only part 1 of 3 landed before the download was interrupted.
(newer / "Model-Q4_K_M-00001-of-00003.gguf").write_bytes(b"\0")
os.utime(older, (1_000, 1_000))
os.utime(newer, (2_000, 2_000))
repo = _repo(
"Org/Split",
[],
repo_dir,
revisions = [
SimpleNamespace(files = [_file("Model-Q8_0.gguf", 5_000)], snapshot_path = older),
SimpleNamespace(
files = [_file("Model-Q4_K_M-00001-of-00003.gguf", 6_000)], snapshot_path = newer
),
],
)
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])]
)
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active)
rows = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))["cached"]
assert rows[0]["load_id"] == str(older)
def test_list_cached_gguf_omits_load_id_when_no_snapshot_is_complete(monkeypatch, tmp_path):
"""With only a half-downloaded split quant, fall back to the repo id, not a path."""
active = tmp_path / "active"
repo_dir = tmp_path / "legacy" / "models--Org--Torn"
snapshot = repo_dir / "snapshots" / "rev"
snapshot.mkdir(parents = True)
(snapshot / "Model-Q4_K_M-00001-of-00003.gguf").write_bytes(b"\0")
repo = _repo(
"Org/Torn",
[],
repo_dir,
revisions = [
SimpleNamespace(
files = [_file("Model-Q4_K_M-00001-of-00003.gguf", 6_000)], snapshot_path = snapshot
),
],
)
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])]
)
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active)
rows = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))["cached"]
assert "load_id" not in rows[0]
def test_list_cached_gguf_skips_snapshot_with_one_incomplete_variant(monkeypatch, tmp_path):
"""A good quant beside a half-downloaded one is still not a safe load target."""
import os
active = tmp_path / "active"
repo_dir = tmp_path / "legacy" / "models--Org--Mixed"
older, newer = repo_dir / "snapshots" / "rev-a", repo_dir / "snapshots" / "rev-b"
for path in (older, newer):
path.mkdir(parents = True)
(older / "Model-Q8_0.gguf").write_bytes(b"\0")
# rev-b has a complete Q8_0 AND a half-downloaded split Q4_K_M. The picker
# enumerates the whole directory, so it would offer the broken one.
(newer / "Model-Q8_0.gguf").write_bytes(b"\0")
(newer / "Model-Q4_K_M-00001-of-00003.gguf").write_bytes(b"\0")
os.utime(older, (1_000, 1_000))
os.utime(newer, (2_000, 2_000))
repo = _repo(
"Org/Mixed",
[],
repo_dir,
revisions = [
SimpleNamespace(files = [_file("Model-Q8_0.gguf", 5_000)], snapshot_path = older),
SimpleNamespace(
files = [
_file("Model-Q8_0.gguf", 5_000),
_file("Model-Q4_K_M-00001-of-00003.gguf", 6_000),
],
snapshot_path = newer,
),
],
)
monkeypatch.setattr(
models_route, "_all_hf_cache_scans", lambda: [SimpleNamespace(repos = [repo])]
)
monkeypatch.setattr(models_route, "_resolve_hf_cache_dir", lambda: active)
rows = asyncio.run(models_route.list_cached_gguf(current_subject = "test-user"))["cached"]
assert rows[0]["load_id"] == str(older)
def test_list_cached_gguf_includes_non_suffix_repo_when_cache_contains_gguf(monkeypatch, tmp_path):
repo = _repo(
"HauhauCS/Gemma-4-E4B-Uncensored-HauhauCS-Aggressive",

View file

@ -0,0 +1,242 @@
# 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 checkpoint loading must shard across every visible GPU (#7053): the
``device_map="sequential"`` loader default stacks the whole model on GPU0 and OOMs
while the other GPUs sit empty. The loader now passes ``device_map="balanced"``, but
only on a real multi-GPU CUDA/ROCm host, so single-GPU, CPU and MLX are untouched."""
from __future__ import annotations
import contextlib
import sys
import types
from pathlib import Path
_BACKEND_DIR = Path(__file__).resolve().parent.parent
if str(_BACKEND_DIR) not in sys.path:
sys.path.insert(0, str(_BACKEND_DIR))
_TESTS_DIR = Path(__file__).resolve().parent
if str(_TESTS_DIR) not in sys.path:
sys.path.insert(0, str(_TESTS_DIR))
# Reuse the absolute-paths test's stub harness for loading core/export/export.py
# without torch/unsloth.
from test_export_absolute_paths import ( # noqa: E402
_install_export_backend_stubs,
_load_module,
)
def _export_mod(monkeypatch):
_install_export_backend_stubs(monkeypatch)
return _load_module("test_core_export_backend_device_map", "core/export/export.py", monkeypatch)
def _stub_hardware(monkeypatch, visible, device_map):
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: visible, raising = False)
monkeypatch.setattr(hw, "get_device_map", lambda ids: device_map, raising = False)
# ── _multi_gpu_device_map_kwargs ──
def test_multi_gpu_host_gets_balanced(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0, 1, 2], "balanced")
assert mod._multi_gpu_device_map_kwargs() == {"device_map": "balanced"}
def test_single_gpu_host_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0], "sequential")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_non_balanced_resolution_keeps_loader_default(monkeypatch):
# >1 visible id but a non-CUDA device resolves to "sequential": pass nothing.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
_stub_hardware(monkeypatch, [0, 1], "sequential")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_uuid_mig_mask_falls_back_to_count_detection(monkeypatch):
# UUID/MIG masks resolve to NO numeric ids ([]), but get_device_map(None) still
# detects >1 GPU, so the empty list must route there, not to the loader default.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [], raising = False)
monkeypatch.setattr(
hw,
"get_device_map",
lambda ids: "balanced" if ids is None else "sequential",
raising = False,
)
assert mod._multi_gpu_device_map_kwargs() == {"device_map": "balanced"}
def test_no_visible_gpus_keeps_loader_default(monkeypatch):
# Empty mask / CPU host: get_device_map(None) resolves "sequential" -> {}.
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", lambda: [], raising = False)
monkeypatch.setattr(hw, "get_device_map", lambda ids: "sequential", raising = False)
assert mod._multi_gpu_device_map_kwargs() == {}
def test_mlx_host_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
# The stubs set _IS_MLX = True; even a multi-GPU view must yield no device_map.
_stub_hardware(monkeypatch, [0, 1], "balanced")
assert mod._multi_gpu_device_map_kwargs() == {}
def test_hardware_probe_failure_keeps_loader_default(monkeypatch):
mod = _export_mod(monkeypatch)
monkeypatch.setattr(mod, "_IS_MLX", False)
hw = sys.modules["utils.hardware"]
def _boom():
raise RuntimeError("no GPUs")
monkeypatch.setattr(hw, "get_parent_visible_gpu_ids", _boom, raising = False)
assert mod._multi_gpu_device_map_kwargs() == {}
# ── load_checkpoint forwards the kwargs to from_pretrained ──
class _RecordingLoader:
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
return types.SimpleNamespace(), types.SimpleNamespace()
def _load_text_checkpoint(monkeypatch, tmp_path, device_map_kwargs):
mod = _export_mod(monkeypatch)
_RecordingLoader.calls = []
monkeypatch.setattr(mod, "FastLanguageModel", _RecordingLoader)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: device_map_kwargs)
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
assert ok, message
assert len(_RecordingLoader.calls) == 1
return _RecordingLoader.calls[0]
def test_load_checkpoint_forwards_balanced_device_map(monkeypatch, tmp_path):
kwargs = _load_text_checkpoint(monkeypatch, tmp_path, {"device_map": "balanced"})
assert kwargs["device_map"] == "balanced"
def test_load_checkpoint_omits_device_map_on_single_gpu(monkeypatch, tmp_path):
kwargs = _load_text_checkpoint(monkeypatch, tmp_path, {})
assert "device_map" not in kwargs # loader default (sequential) untouched
# ── a load that succeeds but offloads to CPU/disk ──
def test_cpu_offloaded_modules_counts_cpu_and_disk(monkeypatch):
mod = _export_mod(monkeypatch)
model = types.SimpleNamespace(hf_device_map = {"a": 0, "b": "cpu", "c": 1, "d": "disk"})
assert mod._cpu_offloaded_modules(model) == 2
def test_cpu_offloaded_modules_ignores_gpu_only_and_missing_maps(monkeypatch):
mod = _export_mod(monkeypatch)
assert mod._cpu_offloaded_modules(types.SimpleNamespace(hf_device_map = {"a": 0})) == 0
assert mod._cpu_offloaded_modules(types.SimpleNamespace(hf_device_map = None)) == 0
assert mod._cpu_offloaded_modules(types.SimpleNamespace()) == 0
class _SpillThenCleanLoader:
"""First call offloads to CPU (bf16 accepts it silently), second is clean."""
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
device_map = {"model.layers.0": 0} if len(cls.calls) > 1 else {"model.layers.0": "cpu"}
return types.SimpleNamespace(hf_device_map = device_map), types.SimpleNamespace()
def _run_spill_loader(monkeypatch, tmp_path, device_map_kwargs):
mod = _export_mod(monkeypatch)
_SpillThenCleanLoader.calls = []
monkeypatch.setattr(mod, "FastLanguageModel", _SpillThenCleanLoader)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: device_map_kwargs)
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
return ok, message, _SpillThenCleanLoader.calls
def test_successful_load_that_offloads_to_cpu_retries_single_device(monkeypatch, tmp_path):
# Nothing raises, so only hf_device_map catches it; the parameters would otherwise
# stay on meta and kill the export inside safetensors.
ok, message, calls = _run_spill_loader(monkeypatch, tmp_path, {"device_map": "balanced"})
assert ok, message
assert len(calls) == 2
assert calls[0]["device_map"] == "balanced"
assert "device_map" not in calls[1]
def test_single_gpu_offload_is_left_alone(monkeypatch, tmp_path):
# No multi-GPU map was requested, so there is nothing to retry on.
ok, message, calls = _run_spill_loader(monkeypatch, tmp_path, {})
assert ok, message
assert len(calls) == 1
def test_retry_result_is_kept_even_if_it_also_offloads(monkeypatch, tmp_path):
# The retry runs with _device_map_override set, so it must never recurse again.
mod = _export_mod(monkeypatch)
class _AlwaysSpills:
calls: list[dict] = []
@classmethod
def from_pretrained(cls, **kwargs):
cls.calls.append(kwargs)
return types.SimpleNamespace(hf_device_map = {"a": "cpu"}), types.SimpleNamespace()
monkeypatch.setattr(mod, "FastLanguageModel", _AlwaysSpills)
monkeypatch.setattr(mod, "detect_audio_type", lambda *a, **k: None)
monkeypatch.setattr(mod, "is_vision_model", lambda *a, **k: False)
monkeypatch.setattr(mod, "_hf_offline", lambda *a, **k: False)
monkeypatch.setattr(mod, "_offline_window_if", lambda flag: contextlib.nullcontext())
monkeypatch.setattr(mod, "_multi_gpu_device_map_kwargs", lambda: {"device_map": "balanced"})
checkpoint = tmp_path / "checkpoint-100"
checkpoint.mkdir()
backend = mod.ExportBackend.__new__(mod.ExportBackend)
backend.cleanup_memory = lambda: None
ok, message = backend.load_checkpoint(str(checkpoint))
assert ok, message
assert len(_AlwaysSpills.calls) == 2

View file

@ -9,10 +9,14 @@ No GPU, network, or subprocesses are required.
from __future__ import annotations
import asyncio
import importlib.util
import logging
import sys
import threading
import types as _types
from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import pytest
@ -28,7 +32,12 @@ _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
# routes/inference.py binds structlog.get_logger at import time, and setdefault
# keeps a bare stub an earlier test left behind: repair it rather than rely on order.
_structlog_stub.get_logger = lambda *_args, **_kwargs: logging.getLogger("structlog_stub")
sys.modules.setdefault("structlog", _structlog_stub)
if not hasattr(sys.modules["structlog"], "get_logger"):
sys.modules["structlog"].get_logger = _structlog_stub.get_logger
try:
import httpx # noqa: F401
@ -120,6 +129,22 @@ def _fail_get_paths_info(*_args, **_kwargs):
raise AssertionError("cached reuse must return before the sizing preflight")
def _load_route_module(name: str, relative_path: str):
"""Import a route module under a private name so patches can't leak."""
spec = importlib.util.spec_from_file_location(name, Path(_BACKEND_DIR) / relative_path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
async def _inline_to_thread(func, /, *args, **kwargs):
return func(*args, **kwargs)
async def _no_gguf_gpu_ids(*_args, **_kwargs):
return None
class TestLoadReusesCachedCopy:
def test_download_uses_selected_cache_for_lookup_preflight_and_write(
self, tmp_path, monkeypatch
@ -809,3 +834,116 @@ class TestLoadHubDownloadExclusion:
Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py"
).read_text()
assert "@_with_gguf_load_marker\n def load_model(" in llama_source
def _capture_hub_guard_require_mmproj(
self,
stored_extra_args,
request_extra_args = None,
):
"""Drive /load's GGUF path and return the hub guard's require_mmproj.
The guard reports a conflicting download, so the 409 is the observation
point and no llama-server ever starts.
"""
import core.inference.llama_cpp as llama_cpp_module
from fastapi import HTTPException
from models.inference import LoadRequest
route = _load_route_module(
"inference_route_module_for_inherited_extra_args_test",
"routes/inference.py",
)
captured = {}
def _fake_blocks(
repo,
variant,
*,
require_mmproj,
hf_token = None,
):
captured["repo"] = repo
captured["variant"] = variant
captured["require_mmproj"] = require_mmproj
return True
# A vision GGUF: require_mmproj is True unless the extras say --no-mmproj.
config = SimpleNamespace(
is_gguf = True,
is_lora = False,
is_vision = True,
is_audio = False,
audio_type = None,
has_audio_input = False,
gguf_hf_repo = REPO,
gguf_variant = VARIANT,
gguf_file = None,
gguf_mmproj_file = None,
identifier = REPO,
display_name = REPO,
)
# Pass-through extras the running backend recorded for the last load.
llama_backend = SimpleNamespace(
is_loaded = False,
extra_args = list(stored_extra_args),
extra_args_source = (REPO, VARIANT),
hf_variant = VARIANT,
model_identifier = REPO,
)
request = LoadRequest(
model_path = REPO,
gguf_variant = VARIANT,
llama_extra_args = request_extra_args,
)
with (
patch.object(
route,
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: config),
),
patch.object(route, "get_llama_cpp_backend", lambda: llama_backend),
patch.object(
route,
"get_inference_backend",
lambda: SimpleNamespace(active_model_name = None),
),
patch.object(route, "_resolve_gguf_gpu_ids_for_request", _no_gguf_gpu_ids),
patch.object(route, "_guard_chat_load_against_training", return_value = None),
patch.object(route, "_effective_load_in_4bit", return_value = False),
patch.object(route, "_hf_offline_if_dns_dead", nullcontext),
patch.object(route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(llama_cpp_module, "_hub_download_blocks_gguf_load", _fake_blocks),
):
with pytest.raises(HTTPException) as exc_info:
asyncio.run(
route._load_model_impl(
request,
SimpleNamespace(
app = SimpleNamespace(
state = SimpleNamespace(llama_parallel_slots = 1),
),
),
current_subject = "test-user",
)
)
assert exc_info.value.status_code == 409
assert captured["repo"] == REPO
return captured["require_mmproj"]
def test_inherited_extra_args_shape_hub_guard_require_mmproj(self):
# Inheritance must resolve before the hub-download guard: an inherited
# --no-mmproj decides require_mmproj, so resolving later rejects a load
# over a download the effective arguments disable (#7251).
assert self._capture_hub_guard_require_mmproj(["--no-mmproj"]) is False
# Control: nothing to inherit, so a vision GGUF still needs its mmproj.
assert self._capture_hub_guard_require_mmproj([]) is True
# An explicit request list wins over the stored one, both ways.
assert (
self._capture_hub_guard_require_mmproj([], request_extra_args = ["--no-mmproj"]) is False
)
assert (
self._capture_hub_guard_require_mmproj(["--no-mmproj"], request_extra_args = []) is True
)

View file

@ -14,6 +14,7 @@ import contextlib
import copy
import json
import sys
import threading
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
@ -55,7 +56,12 @@ def _finish(reason: str) -> str:
)
def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
def _make_backend(
monkeypatch,
streams: list[object],
payloads: list[dict],
urls: list[str] | None = None,
):
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
@ -77,7 +83,12 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
if urls is not None:
urls.append(_url)
stream = streams.pop(0)
if isinstance(stream, BaseException):
raise stream
yield type("FakeResponse", (), {"status_code": 200, "chunks": stream})()
def fake_iter_text_cancellable(
response,
@ -88,9 +99,27 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry)
monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable)
monkeypatch.setattr(backend, "_maybe_recover_from_mtp_crash", lambda *_a, **_k: False)
return backend
def _patch_successful_respawn(
monkeypatch,
backend,
port: int | None = None,
) -> list[bool]:
calls: list[bool] = []
def fake_respawn():
calls.append(True)
if port is not None:
backend._port = port
return True
monkeypatch.setattr(backend, "_respawn_if_dead", fake_respawn)
return calls
def _tool_names(payload: dict) -> list[str]:
return [
(tool.get("function") or {}).get("name")
@ -1837,6 +1866,8 @@ def test_confirm_tool_calls_allow_executes_gguf_tool(monkeypatch):
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
confirm_tool_calls = True,
# Unset defaults to "auto", which would not prompt this safe print(1).
permission_mode = "ask",
session_id = "sess",
)
)
@ -1869,6 +1900,8 @@ def test_confirm_tool_calls_close_after_prompt_cleans_gguf_slot(monkeypatch):
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
confirm_tool_calls = True,
# Unset defaults to "auto", which would not prompt this safe print(1).
permission_mode = "ask",
session_id = "sess",
)
try:
@ -1902,6 +1935,9 @@ def test_confirm_tool_calls_skips_gguf_rag_autoinject(monkeypatch):
tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}],
max_tool_iterations = 1,
confirm_tool_calls = True,
# "ask" gates every call so autoinject waits; unset defaults to
# "auto", where this safe retrieval never gates.
permission_mode = "ask",
session_id = "sess",
rag_scope = {"thread_id": "t1"},
)
@ -1946,6 +1982,8 @@ def test_confirm_tool_calls_deny_skips_gguf_tool_and_retry_can_execute(monkeypat
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 2,
confirm_tool_calls = True,
# Unset defaults to "auto", which would not prompt this safe print(1).
permission_mode = "ask",
session_id = "sess",
)
)
@ -2239,7 +2277,13 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [raising_stream()], payloads)
respawn_calls: list[bool] = []
monkeypatch.setattr(
backend,
"_respawn_if_dead",
lambda: respawn_calls.append(True) or True,
)
monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "OK")
collected: list[dict] = []
@ -2270,6 +2314,271 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch):
# The closing card is marked as an error, not an empty success, so the UI
# renders it as failed.
assert "Error" in (closing[0].get("result") or "")
assert respawn_calls == []
def test_connect_error_before_tool_stream_respawns_and_retries(monkeypatch):
"""A dead server before the first tool-loop response is opened is safe to retry."""
import httpx
payloads: list[dict] = []
urls: list[str] = []
backend = _make_backend(
monkeypatch,
[
httpx.ConnectError("server is down"),
[_sse({"content": "Recovered."}), _done()],
],
payloads,
urls,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend, port = 49999)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert respawn_calls == [True]
assert len(payloads) == 2
assert payloads[0] == payloads[1]
assert urls == [
"http://127.0.0.1:48847/v1/chat/completions",
"http://127.0.0.1:49999/v1/chat/completions",
]
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_connect_error_after_tool_result_recovers_both_generation_paths(monkeypatch):
"""Recover either post-tool generation path without rerunning the tool."""
import httpx
for max_tool_iterations, final_text in (
(2, "The result is 1."),
(1, "Final answer."),
):
payloads: list[dict] = []
backend = _make_backend(
monkeypatch,
[
_structured_tool_call("python", {"code": "print(1)"}, "call_once"),
httpx.ConnectError("server died between turns"),
[_sse({"content": final_text}), _done()],
],
payloads,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
tool_calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
tool_calls.append((name, arguments))
return "1"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "print one"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = max_tool_iterations,
)
)
assert respawn_calls == [True]
assert tool_calls == [("python", {"code": "print(1)"})]
assert len(payloads) == 3
assert payloads[1] == payloads[2]
assert any(e.get("type") == "content" and e.get("text") == final_text for e in events)
def test_connect_error_retry_is_bounded(monkeypatch):
"""A failed retry surfaces the error without another respawn attempt."""
import httpx
payloads: list[dict] = []
backend = _make_backend(
monkeypatch,
[
httpx.ConnectError("server is down"),
httpx.ConnectError("replacement is also down"),
],
payloads,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
except RuntimeError as exc:
raised = True
assert "Lost connection" in str(exc)
assert raised
assert respawn_calls == [True]
assert len(payloads) == 2
def test_pre_header_transport_errors_also_respawn(monkeypatch):
"""A child that dies during prefill already accepted the socket, so it does
not surface as ConnectError. Nothing has streamed yet, so replay is safe."""
import httpx
for exc in (
httpx.RemoteProtocolError("server disconnected without sending a response"),
httpx.ReadError("connection reset by peer"),
httpx.WriteError("broken pipe"),
):
payloads: list[dict] = []
backend = _make_backend(
monkeypatch, [exc, [_sse({"content": "Recovered."}), _done()]], payloads
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert respawn_calls == [True], type(exc).__name__
assert len(payloads) == 2, type(exc).__name__
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_a_not_yet_reaped_child_does_not_burn_the_retry(monkeypatch):
"""A closing server can beat its own exit status, so poll() briefly reports it
alive. Without a grace wait _respawn_if_dead hands back the stale _healthy and the
single retry is spent on the corpse rather than on a replacement."""
import httpx
class _Dying:
# reapable only from the 4th poll, mimicking teardown lagging the socket close
def __init__(self):
self.polls = 0
self.returncode = None
def poll(self):
self.polls += 1
if self.polls > 3:
self.returncode = -9
return -9
return None
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [], payloads)
backend._process = _Dying()
backend._healthy = True
backend._respawn_lock = threading.RLock()
backend._lock = threading.RLock()
backend._mtp_runtime_fallback_lock = threading.Lock()
backend._serial_load_lock = threading.RLock()
backend._cancel_event = threading.Event()
backend._unload_epoch = 0
backend._mtp_runtime_fallback_in_progress = False
backend._mtp_runtime_fallback_active = False
backend._last_load_kwargs = {"gguf_path": "/m.gguf"}
backend._model_identifier = "m"
dying = backend._process
loads: list[dict] = []
@contextlib.contextmanager
def dead_until_respawned(
_c,
_url,
payload,
_ce,
headers = None,
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
if backend._process is dying:
raise httpx.ReadError("connection reset while shutting down")
yield type(
"FakeResponse",
(),
{"status_code": 200, "chunks": [_sse({"content": "Recovered."}), _done()]},
)()
def fake_load(**kwargs):
loads.append(kwargs)
backend._process = type("Live", (), {"poll": lambda self: None, "returncode": None})()
backend._healthy = True
return True
monkeypatch.setattr(backend, "_stream_with_retry", dead_until_respawned)
monkeypatch.setattr(backend, "load_model", fake_load)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert len(loads) == 1
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_prefill_timeout_is_not_retried(monkeypatch):
"""A slow-but-alive server must not have its first-token budget spent twice."""
import httpx
for exc in (httpx.ReadTimeout("no first token"), httpx.PoolTimeout("pool")):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [exc], payloads)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
except httpx.TimeoutException:
raised = True
assert raised, type(exc).__name__
assert respawn_calls == [], type(exc).__name__
assert len(payloads) == 1, type(exc).__name__
def test_mtp_crash_recovery_wins_over_respawn(monkeypatch):
"""An MTP crash reloads without MTP, so never respawn the same config on top."""
import httpx
for max_tool_iterations in (2, 1):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [httpx.ConnectError("mtp crash")], payloads)
monkeypatch.setattr(backend, "_maybe_recover_from_mtp_crash", lambda *_a, **_k: True)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = max_tool_iterations,
)
)
except RuntimeError as exc:
raised = True
assert "Lost connection" in str(exc)
assert raised
assert respawn_calls == []
assert len(payloads) == 1
def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch):
@ -2368,7 +2677,7 @@ def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(monkeypat
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda n, a, **_k: (calls.append((n, a)) or "x"),
lambda n, a, **_k: calls.append((n, a)) or "x",
)
events = list(
@ -2425,7 +2734,7 @@ def test_gguf_truncated_ordinary_json_with_name_key_is_shown_not_suppressed(monk
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda n, a, **_k: (calls.append((n, a)) or "x"),
lambda n, a, **_k: calls.append((n, a)) or "x",
)
events = list(
@ -2452,7 +2761,7 @@ def test_gguf_truncated_disabled_name_json_is_preserved_when_tools_active(monkey
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda n, a, **_k: (calls.append((n, a)) or "x"),
lambda n, a, **_k: calls.append((n, a)) or "x",
)
events = list(
@ -2509,7 +2818,7 @@ def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda n, a, **_k: (calls.append((n, a)) or "x"),
lambda n, a, **_k: calls.append((n, a)) or "x",
)
events = list(
@ -2692,7 +3001,7 @@ def test_gguf_initial_buffer_flush_holds_split_rehearsal_name(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
events = list(
@ -2729,7 +3038,7 @@ def test_gguf_rehearsal_name_after_prose_in_streaming_is_not_leaked(monkeypatch)
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
events = list(
@ -2762,7 +3071,7 @@ def test_gguf_plain_answer_ending_with_tool_name_word_is_preserved(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
events = list(
@ -2797,7 +3106,7 @@ def test_gguf_long_tool_name_split_rehearsal_is_not_capped_and_executes(monkeypa
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda n, a, **_k: (calls.append((n, a)) or "result"),
lambda n, a, **_k: calls.append((n, a)) or "result",
)
events = list(
@ -2831,7 +3140,7 @@ def test_gguf_streaming_keeps_bare_args_before_think_block(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
events = list(
@ -2863,7 +3172,7 @@ def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
events = list(
@ -2897,7 +3206,7 @@ def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose(mon
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
events = list(
@ -2957,7 +3266,7 @@ def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "OK",
)
events = list(
@ -3021,7 +3330,7 @@ def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "OK",
)
list(
@ -3048,7 +3357,7 @@ def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "OK",
)
list(
@ -3073,7 +3382,7 @@ def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disable
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
events = list(
backend.generate_chat_completion_with_tools(
@ -3108,7 +3417,7 @@ def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch):
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
lambda name, arguments, **_k: calls.append((name, arguments)) or "result",
)
list(

View file

@ -46,6 +46,68 @@ def test_dir_model_format_gguf_only(tmp_path):
assert models_route._dir_model_format(d) == "gguf"
def test_dir_model_format_mmproj_only_is_not_gguf(tmp_path):
# A lone vision adapter has nothing servable: the variant selector drops mmproj.
d = tmp_path / "model"
_touch(d / "mmproj-F16.gguf")
assert models_route._dir_model_format(d) is None
def test_dir_model_format_mmproj_beside_weights_is_still_gguf(tmp_path):
d = tmp_path / "model"
_touch(d / "mmproj-F16.gguf")
_touch(d / "model-Q4_K_M.gguf")
assert models_route._dir_model_format(d) == "gguf"
def test_dir_model_format_recursive_sees_split_quant_subdirs(tmp_path):
# HF cache snapshots keep split quants in per-quant subdirs. A flat glob reports
# no GGUF there, which would hide every sharded repo from the GGUF pickers.
d = tmp_path / "snapshot"
_touch(d / "UD-Q4_K_XL" / "model-00001-of-00002.gguf")
assert models_route._dir_model_format(d) is None
assert models_route._dir_model_format(d, recursive = True) == "gguf"
def test_dir_model_format_recursive_ignores_mmproj_only_subdirs(tmp_path):
d = tmp_path / "snapshot"
_touch(d / "mmproj" / "mmproj-F16.gguf")
assert models_route._dir_model_format(d, recursive = True) is None
def test_scan_models_dir_mmproj_only_folder_is_not_gguf(tmp_path):
# Same rule as _dir_model_format, applied by the parallel ./models scanner.
_touch(tmp_path / "vision" / "mmproj-F16.gguf")
_touch(tmp_path / "real" / "model-Q4_K_M.gguf")
formats = {m.display_name: m.model_format for m in models_route._scan_models_dir(tmp_path)}
assert formats["vision"] is None
assert formats["real"] == "gguf"
def test_scan_models_dir_skips_standalone_mmproj_file(tmp_path):
# A loose mmproj-*.gguf is a vision adapter with no weights to serve, so it must
# not be offered as a model the way a loose primary GGUF is.
_touch(tmp_path / "mmproj-F16.gguf")
_touch(tmp_path / "model-Q4_K_M.gguf")
names = {m.display_name for m in models_route._scan_models_dir(tmp_path)}
assert names == {"model-Q4_K_M"}
def test_scan_lmstudio_dir_skips_standalone_mmproj_file(tmp_path):
_touch(tmp_path / "mmproj-F16.gguf")
_touch(tmp_path / "model-Q4_K_M.gguf")
names = {m.display_name for m in models_route._scan_lmstudio_dir(tmp_path)}
assert names == {"model-Q4_K_M"}
def test_scan_lmstudio_dir_skips_mmproj_under_publisher(tmp_path):
# LM Studio's publisher/model.gguf layout classifies on a separate branch.
_touch(tmp_path / "Publisher" / "mmproj-F16.gguf")
_touch(tmp_path / "Publisher" / "model-Q4_K_M.gguf")
names = {m.display_name for m in models_route._scan_lmstudio_dir(tmp_path)}
assert names == {"model-Q4_K_M"}
def test_dir_model_format_gguf_with_config_is_still_gguf(tmp_path):
# A config.json alongside the .gguf must not flip it to non-GGUF.
d = tmp_path / "model"

View file

@ -1108,7 +1108,7 @@ def test_build_index_covers_legacy_default_lmstudio_and_custom_roots(monkeypatch
monkeypatch.setattr(
models_route,
"_scan_hf_cache",
lambda d: scanned.append(("hf", str(Path(d).resolve()))) or [],
lambda d, **_: scanned.append(("hf", str(Path(d).resolve()))) or [],
)
monkeypatch.setattr(
models_route,
@ -1337,6 +1337,56 @@ def test_hf_cache_entry_loads_from_local_snapshot_path(tmp_path):
# ── review round 5: concurrent-swap, repo-id identity, /v1/models id, gate, 503 ──
def _revision_pair(root, complete: bool):
"""Two revisions of one cache repo; the newer one is optionally half-downloaded."""
snaps = root / "models--org--Repo" / "snapshots"
old, new = snaps / "rev-old", snaps / "rev-new"
for path in (old, new):
path.mkdir(parents = True)
(old / "model-Q8_0.gguf").write_bytes(b"GGUF stub")
name = "model-Q4_K_M.gguf" if complete else "model-Q4_K_M-00001-of-00003.gguf"
(new / name).write_bytes(b"GGUF stub")
return old, new
def test_sibling_revision_resolves_to_its_own_weights(tmp_path):
# /v1/models advertises only the snapshot dir name, so a durable pin holds one
# revision hash. A newer snapshot must not strand it, and the old revision must
# resolve to ITS OWN directory rather than be redirected onto the newest.
old, new = _revision_pair(tmp_path, complete = True)
found = dict(resolver._sibling_revision_entries(str(new), "org/Repo"))
assert "rev-old" in found
assert found["rev-old"].load_path == str(old)
def test_incomplete_sibling_revision_is_not_indexed(tmp_path):
# A half-downloaded revision cannot load, so naming it must not resolve to it.
old, _new = _revision_pair(tmp_path, complete = False)
# Point the scan at the complete one; the partial sibling is the candidate here.
found = dict(resolver._sibling_revision_entries(str(old), "org/Repo"))
assert "rev-new" not in found
def test_sibling_revisions_ignore_a_scan_folder_named_snapshots(tmp_path):
# A user scan folder called "snapshots" holds unrelated models, not revisions of
# one repo; treating them as revisions would silently serve model-a as model-b.
snaps = tmp_path / "snapshots"
for name in ("model-a", "model-b"):
(snaps / name).mkdir(parents = True)
(snaps / name / "model-Q4_K_M.gguf").write_bytes(b"GGUF stub")
found = dict(resolver._sibling_revision_entries(str(snaps / "model-a"), "model-a"))
assert found == {}
def test_sibling_revisions_skip_plain_repo_ids():
assert dict(resolver._sibling_revision_entries("org/Repo-GGUF", "org/Repo-GGUF")) == {}
def test_already_loaded_by_repo_id_is_not_reswapped(monkeypatch):
# A model loaded normally has model_identifier == repo id, but the resolver
# returns the concrete load path. A request for that repo must count as already

File diff suppressed because it is too large Load diff

View file

@ -120,10 +120,7 @@ class TestParser:
# Only the wrapping newline is trimmed; code-argument indentation survives.
text = (
"<function=python><parameter=code>\n"
" indented = 1\n"
" more\n"
"</parameter></function>"
"<function=python><parameter=code>\n indented = 1\n more\n</parameter></function>"
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
@ -157,10 +154,7 @@ class TestParser:
def test_xml_param_preserves_leading_indentation(self):
# Only the wrapping newline is trimmed, so code-argument indentation survives (str.strip() destroyed it).
text = (
"<function=python><parameter=code>\n"
" indented = 1\n"
" more\n"
"</parameter></function>"
"<function=python><parameter=code>\n indented = 1\n more\n</parameter></function>"
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
@ -310,20 +304,18 @@ class TestParser:
tag has not arrived yet, so the strip regex has to accept
end-of-string as a terminator. Regression for the Gemini
high-severity flag on this PR."""
text = (
"<think>I should call web_search[ARGS]" '{"query":"weather"} next to find the answer.'
)
text = '<think>I should call web_search[ARGS]{"query":"weather"} next to find the answer.'
result = parse_tool_calls_from_text(text)
# Inside an unclosed think block no calls are yielded.
assert result == []
def test_rehearsal_inside_unclosed_bracket_think_is_ignored(self):
text = "[THINK]planning to use python[ARGS]" '{"code":"print(1)"} but not yet.'
text = '[THINK]planning to use python[ARGS]{"code":"print(1)"} but not yet.'
result = parse_tool_calls_from_text(text)
assert result == []
def test_rehearsal_after_closed_think_still_parsed(self):
text = "<think>planning</think>" 'python[ARGS]{"code":"print(1)"}'
text = '<think>planning</think>python[ARGS]{"code":"print(1)"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
@ -365,7 +357,7 @@ class TestParser:
def test_mistral_bracket_nested_json(self):
# Brace-balance scan handles nested objects and braces inside string literals.
text = "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}'
text = '[TOOL_CALLS]web_search{"query":"a {nested} brace","opts":{"limit":5}}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
import json as _json
@ -376,11 +368,7 @@ class TestParser:
def test_mistral_bracket_with_prose(self):
# Bracket-tag surrounded by prose is still recognised.
text = (
"Sure, I will look that up.\n"
'[TOOL_CALLS]web_search{"query":"weather"}\n'
"Calling now."
)
text = 'Sure, I will look that up.\n[TOOL_CALLS]web_search{"query":"weather"}\nCalling now.'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
@ -408,7 +396,7 @@ class TestParser:
assert "print(1)" in result[0]["function"]["arguments"]
def test_rehearsal_with_prose(self):
text = "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}'
text = 'I should call the python tool. Like this: python[ARGS]{"code":"x = 1"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
@ -489,16 +477,14 @@ class TestParser:
assert result[0]["function"]["name"] == "web_search"
def test_think_block_stripped_before_bracket_tag(self):
text = (
"<think>Let me search for that.</think>\n" '[TOOL_CALLS]web_search{"query":"weather"}'
)
text = '<think>Let me search for that.</think>\n[TOOL_CALLS]web_search{"query":"weather"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
def test_uppercase_think_tag_stripped(self):
# Some templates use [THINK]...[/THINK] instead of <think>.
text = "[THINK]planning my next call[/THINK]" '[TOOL_CALLS]python{"code":"print(1)"}'
text = '[THINK]planning my next call[/THINK][TOOL_CALLS]python{"code":"print(1)"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
@ -544,8 +530,7 @@ class TestParser:
def test_xml_wins_over_bracket(self):
# When a model emits both forms in one message, the XML form is canonical and wins.
text = (
'<tool_call>{"name":"primary","arguments":{}}</tool_call>'
'[TOOL_CALLS]secondary{"k":"v"}'
'<tool_call>{"name":"primary","arguments":{}}</tool_call>[TOOL_CALLS]secondary{"k":"v"}'
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
@ -728,7 +713,7 @@ class TestParserMultiFormat:
def test_llama3_python_tag_dot_call_multi_arg(self):
import json
text = "<|python_tag|>get_weather.call(" 'location="Tokyo", units="celsius", days=5)'
text = '<|python_tag|>get_weather.call(location="Tokyo", units="celsius", days=5)'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
args = json.loads(result[0]["function"]["arguments"])
@ -1330,12 +1315,7 @@ class TestParserDeepSeek:
def test_v3_1_strict_rejects_unclosed_envelope(self):
# Envelope truncated mid-stream (no <tool▁calls▁end>): healed by
# default, rejected with Auto-Heal off.
text = (
"<tool▁calls▁begin>"
"<tool▁call▁begin>get_time"
"<tool▁sep>"
'{"city": "Tokyo"}'
)
text = '<tool▁calls▁begin><tool▁call▁begin>get_time<tool▁sep>{"city": "Tokyo"}'
assert len(parse_tool_calls_from_text(text)) == 1
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
@ -1765,9 +1745,9 @@ class TestParserCrossFormatRouting:
for label, text, expected_name in cases:
result = parse_tool_calls_from_text(text)
assert len(result) == 1, f"{label}: parser missed the call"
assert result[0]["function"]["name"] == expected_name, (
f"{label}: got {result[0]['function']['name']!r}, " f"expected {expected_name!r}"
)
assert (
result[0]["function"]["name"] == expected_name
), f"{label}: got {result[0]['function']['name']!r}, expected {expected_name!r}"
def test_all_new_markers_in_tool_xml_signals(self):
# The safetensors / MLX streaming buffer must wake on every supported emission marker --
@ -2538,6 +2518,9 @@ class TestLoopBasic:
tools = [{"type": "function", "function": {"name": "render_html"}}],
execute_tool = exec_fn,
confirm_tool_calls = True,
# Unset defaults to "auto", which only gates render_html when it
# reaches the network, so this static canvas would not prompt.
permission_mode = "ask",
session_id = "sess",
max_tool_iterations = 3,
)
@ -3402,10 +3385,7 @@ class TestLoopRePrompt:
loop, exec_fn = _make_loop(
turns = [
["Let me search for that."],
[
'<tool_call>{"name":"web_search","arguments":'
'{"query":"sky color"}}</tool_call>'
],
['<tool_call>{"name":"web_search","arguments":{"query":"sky color"}}</tool_call>'],
["The sky is blue."],
],
exec_results = ["Blue (Rayleigh scattering)"],
@ -3513,7 +3493,7 @@ class TestLoopCanonicalHealKey:
def test_python_bare_string_heals_to_code(self):
loop, exec_fn = _make_loop(
turns = [
['<tool_call>{"name":"python","arguments":"print(1)"}' "</tool_call>"],
['<tool_call>{"name":"python","arguments":"print(1)"}</tool_call>'],
["done"],
],
exec_results = ["1\n"],
@ -3526,7 +3506,7 @@ class TestLoopCanonicalHealKey:
def test_terminal_bare_string_heals_to_command(self):
loop, exec_fn = _make_loop(
turns = [
['<tool_call>{"name":"terminal","arguments":"ls -la"}' "</tool_call>"],
['<tool_call>{"name":"terminal","arguments":"ls -la"}</tool_call>'],
["done"],
],
exec_results = ["..."],
@ -3537,7 +3517,7 @@ class TestLoopCanonicalHealKey:
def test_unknown_tool_bare_string_heals_to_query(self):
loop, exec_fn = _make_loop(
turns = [
['<tool_call>{"name":"web_search","arguments":"hello"}' "</tool_call>"],
['<tool_call>{"name":"web_search","arguments":"hello"}</tool_call>'],
["ok"],
],
exec_results = ["..."],
@ -3927,6 +3907,8 @@ class TestGuardrails:
turns = [['<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>']],
exec_results = ["OK"],
confirm_tool_calls = True,
# Unset defaults to "auto", which would not prompt this safe call.
permission_mode = "ask",
session_id = "sess",
max_tool_iterations = 1,
)
@ -3957,6 +3939,9 @@ class TestGuardrails:
loop, exec_fn = _make_loop(
turns = [["plain answer"]],
confirm_tool_calls = True,
# "ask" gates every call so autoinject waits; the companion test
# below covers "auto", where the safe retrieval never gates.
permission_mode = "ask",
rag_scope = {"thread_id": "t1"},
)
events = _collect_events(loop)
@ -4313,6 +4298,8 @@ class TestPlanWithoutActionReprompt:
["SHOULD NOT APPEAR"],
],
confirm_tool_calls = True,
# Only "ask" gates the always-safe web_search, so the deny path runs.
permission_mode = "ask",
session_id = "sess",
nudge_tool_calls = True,
)
@ -4367,20 +4354,18 @@ class TestRoutesPythonTagStrip:
def test_python_tag_multiline_with_less_than(self):
# Combined: multi-line code AND literal ``<`` in code.
text = (
'<|python_tag|>python.call(code="for i in range(10):\n'
" if i < 5:\n"
' print(i)")'
'<|python_tag|>python.call(code="for i in range(10):\n if i < 5:\n print(i)")'
)
assert self._strip(text) == ""
def test_python_tag_stops_at_eom_sentinel(self):
# Strip stops at the next Llama-3 ``<|`` sentinel so any
# trailing assistant content survives.
text = '<|python_tag|>python.call(code="multi\nline")' "<|eom_id|>final answer text"
text = '<|python_tag|>python.call(code="multi\nline")<|eom_id|>final answer text'
assert self._strip(text) == "<|eom_id|>final answer text"
def test_python_tag_stops_at_eot_sentinel(self):
text = '<|python_tag|>brave_search.call(query="x")' "<|eot_id|>after"
text = '<|python_tag|>brave_search.call(query="x")<|eot_id|>after'
assert self._strip(text) == "<|eot_id|>after"
def test_python_tag_json_form_multiline_stripped(self):
@ -4410,7 +4395,7 @@ class TestParserRobustness:
# too. Was extracting name only and silently dropping the args.
import json
text = "<tool_call>\n" '{"name": "search", "parameters": {"q": "ramen"}}\n' "</tool_call>"
text = '<tool_call>\n{"name": "search", "parameters": {"q": "ramen"}}\n</tool_call>'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "search"
@ -4421,7 +4406,7 @@ class TestParserRobustness:
# ``<function name="..."><param name="...">v</param></function>``.
import json
text = '<function name="get_weather">' '<param name="city">Tokyo</param>' "</function>"
text = '<function name="get_weather"><param name="city">Tokyo</param></function>'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "get_weather"

View file

@ -219,7 +219,7 @@ class TestUploadDenylist:
)
def test_plain_post_json_not_blocked(self):
_ok("import requests\n" 'requests.post("https://api.weather.gov/lookup", json={"k": "v"})')
_ok('import requests\nrequests.post("https://api.weather.gov/lookup", json={"k": "v"})')
class TestSandboxEnvIsolation:
@ -693,6 +693,51 @@ class TestBashBlocklistPosition:
def test_while_do_blocked(self):
assert "curl" in self._find()("while true; do curl --version; break; done")
# ---- `.` is the POSIX synonym for the blocked `source` builtin ----
def test_dot_source_blocked(self):
assert "." in self._find()(". ./script.sh")
assert "." in self._find()("cat x && . ./payload")
def test_dot_in_argument_position_allowed(self):
assert self._find()("find . -type f") == set()
assert self._find()("ls .") == set()
assert self._find()("cd .") == set()
# ---- ANSI-C quoting must not hide a blocked command name ----
def test_ansi_c_quoted_command_blocked(self):
assert "ssh" in self._find()("$'ssh' user@host")
assert "source" in self._find()("$'source' ./payload")
def test_ansi_c_data_with_newline_is_not_a_command(self):
# $'...' expands to a single word, so a newline inside it is data for
# printf, not a separator that starts a second command.
payload = "printf '%s' $'hello\\n" + "rm" + " -rf x\\n'"
assert self._find()(payload) == set()
def test_command_position_glob_matches_blocked_name(self):
# Bash expands the pattern to the blocked name after this scan runs.
assert "rm" in self._find()("/bin/r[m] -rf /tmp/victim")
assert "rm" in self._find()("/bin/r? -rf /tmp/victim")
def test_glob_without_literal_character_allowed(self):
# A bracket expression in argument position is not a command word.
assert self._find()("echo '[a]'") == set()
def test_attached_exec_flag_value_blocked(self):
# fd accepts the command attached to the flag, so the value is what runs.
assert "rm" in self._find()("fd victim . --exec=rm")
assert "rm" in self._find()("fd victim . --exec-batch=rm")
def test_short_flag_neighbour_not_read_as_command(self):
# Only the long spellings carry an attached command; -x belongs to too
# many other utilities to read its neighbour as one.
assert self._find()("grep -x rm file.txt") == set()
def test_alias_body_scanned_as_command(self):
# `alias zap='rm -rf'` stores a command bash runs when zap is invoked.
assert "rm" in self._find()("alias zap='rm -rf'")
assert self._find()("alias ll='ls -la'") == set()
class TestHfUploadImportGate:
"""Upload-method blocking requires an HF import in scope, so paramiko /
@ -737,15 +782,11 @@ class TestHfUploadImportGate:
def test_hf_bare_name_upload_folder_safe_allowed(self):
_ok(
"from huggingface_hub import upload_folder;"
" upload_folder(folder_path='x', repo_id='r')"
"from huggingface_hub import upload_folder; upload_folder(folder_path='x', repo_id='r')"
)
def test_hf_bare_name_create_commit_safe_allowed(self):
_ok(
"from huggingface_hub import create_commit;"
" create_commit(operations=[], repo_id='r')"
)
_ok("from huggingface_hub import create_commit; create_commit(operations=[], repo_id='r')")
def test_bare_name_upload_file_without_hf_import_allowed(self):
# No HF import -- local helper named upload_file passes.

View file

@ -19,6 +19,7 @@ from __future__ import annotations
import asyncio
import inspect
import socket
import sys
import threading
import time
@ -528,6 +529,358 @@ def test_runtime_recovery_is_single_flight(monkeypatch):
release.set()
def test_single_flight_claim_is_released_when_the_reload_cannot_start(monkeypatch):
# Only the reload thread's finally clears the claim, so if starting it raises the
# claim must not latch: nothing else resets it, and _respawn_if_dead then refuses
# forever, for every later model.
b = _recovery_backend()
class _NoThread:
def __init__(self, *args, **kwargs):
pass
def start(self):
raise RuntimeError("can't start new thread")
monkeypatch.setattr(llama_cpp_module.threading, "Thread", _NoThread)
assert b._maybe_recover_from_mtp_crash(RuntimeError()) is False
assert b._mtp_runtime_fallback_in_progress is False
def test_load_kwargs_are_read_once_before_the_claim(monkeypatch):
# Gate and snapshot must share one read: reading twice lets an unload null
# _last_load_kwargs in between, so dict(None) raises after the claim and strands
# the flag with no thread alive to clear it.
b = _recovery_backend()
class _CountingKwargs: # data descriptor, so it wins over the instance dict
def __init__(self, value):
self.value = value
self.reads = 0
def __get__(self, obj, owner):
if obj is None:
return self
self.reads += 1
return self.value
def __set__(self, obj, value):
self.value = value
counter = _CountingKwargs({"model_identifier": "owner/repo"})
monkeypatch.setattr(type(b), "_last_load_kwargs", counter, raising = False)
class _UnstartedThread: # keep the reload off-thread so only sync reads count
def __init__(self, *args, **kwargs):
pass
def start(self):
pass
monkeypatch.setattr(llama_cpp_module.threading, "Thread", _UnstartedThread)
assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
assert counter.reads == 1, f"read {counter.reads} times; an unload can race the claim"
def test_respawn_defers_to_an_inflight_mtp_reload(monkeypatch):
# "Already recovering" must not read as "not an MTP crash": respawning replays the
# crashing MTP kwargs and aborts the in-flight no-MTP reload on its "newer load" check.
b = _recovery_backend()
b._mtp_runtime_fallback_in_progress = True
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == []
# Once that reload finishes, an ordinary respawn works again.
b._mtp_runtime_fallback_in_progress = False
b._process.returncode = -9 # only the respawn path logs it
assert b._respawn_if_dead() is True
assert [kw.get("speculative_type") for kw in loads] == ["auto"]
def test_respawn_does_not_wait_out_the_grace_on_a_replacement(monkeypatch):
# Callers losing the same child queue on _respawn_lock and wake holding the healthy
# REPLACEMENT. Unable to tell it from their own child, each burns the reap grace, and
# that sleep is held under the lock, so N callers cost N grace periods.
class _LiveProcess(_FakeProcess):
returncode = None
def __init__(self):
self.polls = 0
def poll(self): # never reapable, so the grace loop runs to its deadline
self.polls += 1
return None
workers = 4
b = _recovery_backend()
b._healthy = True
b._process.returncode = -9 # only the respawn path logs it
live = _LiveProcess()
loads: list[dict] = []
guard = threading.Lock()
all_in_flight = threading.Event()
# Subclass this instance, not the class: a descriptor on LlamaCppBackend would
# redirect _process for every other live backend, including atexit-registered ones.
state = {"proc": b._process, "readers": set()}
class _Tracked(type(b)):
@property
def _process(self):
"""Reports when every worker has taken its pre-lock look at the child."""
with guard:
state["readers"].add(threading.get_ident())
everyone = len(state["readers"]) >= workers
if everyone:
all_in_flight.set()
return state["proc"]
@_process.setter
def _process(self, value):
state["proc"] = value
b.__class__ = _Tracked
def _load(**kwargs):
# A real load_model takes seconds, so every caller that lost this child is in
# flight before the replacement appears; waiting reproduces that ordering. The
# timeout keeps the pre-fix build, where losers cannot read until the lock is
# free, from hanging instead of failing.
all_in_flight.wait(timeout = 2)
with guard:
loads.append(kwargs)
b._process = live
b._healthy = True # the real load_model marks the new server healthy
return True
monkeypatch.setattr(b, "load_model", _load)
results: list[bool] = []
def _respawn():
outcome = b._respawn_if_dead()
with guard:
results.append(outcome)
threads = [threading.Thread(target = _respawn) for _ in range(workers)]
started = time.monotonic()
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout = 30)
elapsed = time.monotonic() - started
assert results == [True] * workers, results
assert len(loads) == 1, f"{len(loads)} reloads, expected one"
# The grace loop is the only poll() of a live process, so any count means a queued
# caller charged the wait to a server that never failed.
assert live.polls == 0, "queued caller waited out the grace on a healthy server"
assert elapsed < llama_cpp_module._RESPAWN_REAP_GRACE_S * (workers - 1)
class _DyingChild(_FakeProcess):
"""Alive for the first polls, then reapable: what a terminate() looks like."""
def __init__(
self,
code = -15,
alive_polls = 2,
on_death = None,
):
self.polls = 0
self.returncode = None
self._code = code
self._alive_polls = alive_polls
self._on_death = on_death
def poll(self):
self.polls += 1
if self.polls <= self._alive_polls:
return None
if self.returncode is None:
self.returncode = self._code
if self._on_death is not None:
self._on_death()
return self._code
def test_respawn_does_not_resurrect_a_deliberate_unload(monkeypatch):
# unload_model() sets _cancel_event before killing, so a request that loses the
# connection can watch that deliberate exit through the grace loop and call it a
# crash, with _last_load_kwargs still populated (unload clears it after the kill).
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild()
b._cancel_event.set()
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == [], "resurrected a model the user unloaded"
def test_respawn_rechecks_the_cancel_flag_after_the_grace_wait(monkeypatch):
# The unload can also begin while we are already sleeping in the grace loop.
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(on_death = b._cancel_event.set)
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == [], "checked the cancel flag only before the wait"
def test_respawn_does_not_revert_a_newer_load(monkeypatch):
# A model switch landing while we wait must win; replaying the old kwargs would
# swap the user's new model back out.
b = _recovery_backend()
b._healthy = True
replacement = _DyingChild(alive_polls = 10**6)
b._process = _DyingChild(on_death = lambda: setattr(b, "_process", replacement))
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
b._respawn_if_dead()
assert loads == [], "replayed stale kwargs over a newer load"
assert b._process is replacement
def test_respawn_still_recovers_an_ordinary_crash(monkeypatch):
# Guard rail: none of the above may disable the recovery this path exists for.
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(code = -9)
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is True
assert len(loads) == 1
class _NeverReapable(_FakeProcess):
"""A child that stays unreapable, so only the port can tell alive from dead."""
returncode = None
def poll(self):
return None
def test_a_transient_error_against_a_live_server_costs_nothing(monkeypatch):
# The reap grace must not be charged to a server that never died: the sleep is
# held under _respawn_lock, so a full grace per caller serialises into N seconds
# of added latency on an install that is working fine.
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(16)
try:
b = _recovery_backend()
b._healthy = True
b._process = _NeverReapable()
b._port = listener.getsockname()[1]
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
started = time.monotonic()
assert b._respawn_if_dead() is True
elapsed = time.monotonic() - started
assert loads == [], "a live server must not be reloaded"
assert (
elapsed < llama_cpp_module._RESPAWN_REAP_GRACE_S / 2
), f"waited {elapsed:.2f}s on a server that is still accepting"
finally:
listener.close()
def test_a_closed_port_still_waits_for_the_child_to_be_reapable(monkeypatch):
# The other half: no listener means the server really is gone, so the grace
# still runs and the reap-race fix is preserved.
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
probe.bind(("127.0.0.1", 0))
dead_port = probe.getsockname()[1]
probe.close()
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(code = -9)
b._port = dead_port
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is True
assert len(loads) == 1
def test_socket_fast_path_honours_a_pending_unload(monkeypatch):
# unload_model() sets _cancel_event before it kills, so the child is still
# accepting when the probe runs. Reporting it healthy aims the retry at a server
# that is deliberately going away.
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(8)
try:
b = _recovery_backend()
b._healthy = True
b._process = _NeverReapable()
b._port = listener.getsockname()[1]
b._cancel_event.set()
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == []
finally:
listener.close()
def test_an_unload_landing_during_the_reload_is_undone(monkeypatch):
# The cancel check cannot live under _serial_load_lock alone: unload_model never
# takes that lock, so it can land entirely between the check and load_model and
# the captured kwargs then restart a model the user stopped. load_model clears
# _cancel_event on the way in, so _unload_epoch is the surviving evidence.
b = _recovery_backend()
b._healthy = True
b._process = _FakeProcess()
b._process.returncode = -9
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
unloads: list[int] = []
real_unload = b.unload_model
monkeypatch.setattr(b, "unload_model", lambda: unloads.append(1) or real_unload())
# The warning marks the window: after the snapshot, before the reload.
real_warning = llama_cpp_module.logger.warning
fired: list[int] = []
def racing_warning(*args, **kwargs):
if not fired:
fired.append(1)
real_unload()
return real_warning(*args, **kwargs)
monkeypatch.setattr(llama_cpp_module.logger, "warning", racing_warning)
assert b._respawn_if_dead() is False
assert unloads, "the racing unload was not honoured"
def test_socket_probe_is_false_without_a_port():
# Unloaded backends have no port; the probe must not raise, and the caller
# then falls back to the poll-based grace.
b = _recovery_backend()
b._port = None
assert b._server_socket_is_open() is False
def test_runtime_recovery_rechecks_cancel_before_reload():
# recover() must re-check the cancel flag after the death poll (load_model
# clears it), so a reload scheduled just before /unload can't resurrect it.

View file

@ -94,6 +94,9 @@ def _drive(
execute_tool = exec_fn,
session_id = _SESSION,
confirm_tool_calls = True,
# The confirm-gate mechanics (allow/deny/reissue/dedup) need every call to
# prompt; unset defaults to "auto", which only gates high-risk calls.
permission_mode = "ask",
)
events = []
for ev in gen:

View file

@ -54,6 +54,8 @@ logger = structlog.get_logger(__name__)
DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp"
_INSTALL_TIMEOUT_SECONDS = 1800 # 30 min ceiling for download + build/validate
# install_llama_prebuilt.py EXIT_NO_SPACE: out of disk, retrying will not help.
_EXIT_NO_SPACE = 4
# Background job state. Single in-flight update at a time, guarded by _job_lock.
_JOB_IDLE = _flow.JOB_IDLE
@ -496,6 +498,16 @@ def _run_llama_phase(
+ (" Reload your model to use it." if model_was_active else "")
),
}
except _flow.InstallerExit as exc:
# Raw "installer exited 4: <log tail>" says nothing actionable in the UI.
if exc.returncode == _EXIT_NO_SPACE:
logger.warning("llama update: out of disk space")
raise RuntimeError(
"Not enough disk space to install llama.cpp. Free up space or point "
"UNSLOTH_STUDIO_HOME/TMPDIR at a larger volume, then retry."
) from exc
logger.warning("llama update: failed", error = str(exc))
raise
except Exception as exc:
logger.warning("llama update: failed", error = str(exc))
raise

View file

@ -0,0 +1,9 @@
<!-- Source: https://github.com/NousResearch/hermes-agent/blob/main/acp_registry/icon.svg -->
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 16 16" width="16" height="16" fill="none">
<path d="M8 1.5v13" stroke="currentColor" stroke-width="1.5" stroke-linecap="round"/>
<path d="M8 3.25c-2.35-1.4-4.7-.95-6.25.35 1.85-.2 3.8.2 5.55 1.55" stroke="currentColor" stroke-width="1.1" stroke-linecap="round" stroke-linejoin="round"/>
<path d="M8 3.25c2.35-1.4 4.7-.95 6.25.35-1.85-.2-3.8.2-5.55 1.55" stroke="currentColor" stroke-width="1.1" stroke-linecap="round" stroke-linejoin="round"/>
<path d="M8 13.25c-2.3-1-3.05-2.65-1.35-4.15-2 .8-2.35 2.95-.35 4" stroke="currentColor" stroke-width="1.1" stroke-linecap="round" stroke-linejoin="round"/>
<path d="M8 13.25c2.3-1 3.05-2.65 1.35-4.15 2 .8 2.35 2.95.35 4" stroke="currentColor" stroke-width="1.1" stroke-linecap="round" stroke-linejoin="round"/>
<circle cx="8" cy="1.8" r="1.1" fill="currentColor"/>
</svg>

After

Width:  |  Height:  |  Size: 976 B

View file

@ -0,0 +1,18 @@
<!-- Source: https://github.com/openclaw/openclaw/blob/main/apps/linux/src-tauri/icons/icon.svg -->
<svg viewBox="0 0 120 120" fill="none" xmlns="http://www.w3.org/2000/svg">
<defs>
<linearGradient id="lobster-gradient" x1="0%" y1="0%" x2="100%" y2="100%">
<stop offset="0%" stop-color="#ff4d4d"/>
<stop offset="100%" stop-color="#991b1b"/>
</linearGradient>
</defs>
<path d="M60 10 C30 10 15 35 15 55 C15 75 30 95 45 100 L45 110 L55 110 L55 100 C55 100 60 102 65 100 L65 110 L75 110 L75 100 C90 95 105 75 105 55 C105 35 90 10 60 10Z" fill="url(#lobster-gradient)"/>
<path d="M20 45 C5 40 0 50 5 60 C10 70 20 65 25 55 C28 48 25 45 20 45Z" fill="url(#lobster-gradient)"/>
<path d="M100 45 C115 40 120 50 115 60 C110 70 100 65 95 55 C92 48 95 45 100 45Z" fill="url(#lobster-gradient)"/>
<path d="M45 15 Q35 5 30 8" stroke="#ff4d4d" stroke-width="3" stroke-linecap="round"/>
<path d="M75 15 Q85 5 90 8" stroke="#ff4d4d" stroke-width="3" stroke-linecap="round"/>
<circle cx="45" cy="35" r="6" fill="#050810"/>
<circle cx="75" cy="35" r="6" fill="#050810"/>
<circle cx="46" cy="34" r="2.5" fill="#00e5cc"/>
<circle cx="76" cy="34" r="2.5" fill="#00e5cc"/>
</svg>

After

Width:  |  Height:  |  Size: 1.2 KiB

View file

@ -0,0 +1,19 @@
<!-- Source: https://github.com/anomalyco/opencode/blob/dev/packages/console/app/src/asset/brand/opencode-logo-dark-square.svg -->
<svg width="300" height="300" viewBox="0 0 300 300" fill="none" xmlns="http://www.w3.org/2000/svg">
<g transform="translate(30, 0)">
<g clip-path="url(#clip0)">
<mask id="mask0" style="mask-type:luminance" maskUnits="userSpaceOnUse" x="0" y="0" width="240" height="300">
<path d="M240 0H0V300H240V0Z" fill="white"/>
</mask>
<g mask="url(#mask0)">
<path d="M180 240H60V120H180V240Z" fill="#4B4646"/>
<path d="M180 60H60V240H180V60ZM240 300H0V0H240V300Z" fill="#F1ECEC"/>
</g>
</g>
</g>
<defs>
<clipPath id="clip0">
<rect width="240" height="300" fill="white"/>
</clipPath>
</defs>
</svg>

After

Width:  |  Height:  |  Size: 796 B

View file

@ -0,0 +1,19 @@
<!-- Source: https://github.com/anomalyco/opencode/blob/dev/packages/console/app/src/asset/brand/opencode-logo-light-square.svg -->
<svg width="300" height="300" viewBox="0 0 300 300" fill="none" xmlns="http://www.w3.org/2000/svg">
<g transform="translate(30, 0)">
<g clip-path="url(#clip0)">
<mask id="mask0" style="mask-type:luminance" maskUnits="userSpaceOnUse" x="0" y="0" width="240" height="300">
<path d="M240 0H0V300H240V0Z" fill="white"/>
</mask>
<g mask="url(#mask0)">
<path d="M180 240H60V120H180V240Z" fill="#CFCECD"/>
<path d="M180 60H60V240H180V60ZM240 300H0V0H240V300Z" fill="#211E1E"/>
</g>
</g>
</g>
<defs>
<clipPath id="clip0">
<rect width="240" height="300" fill="white"/>
</clipPath>
</defs>
</svg>

After

Width:  |  Height:  |  Size: 797 B

View file

@ -0,0 +1,21 @@
<!-- Source: https://pi.dev/favicon.svg (official Pi press-kit badge) -->
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 800 800">
<rect width="800" height="800" rx="120" fill="#09090b"/>
<path fill="#fff" fill-rule="evenodd" d="
M165.29 165.29
H517.36
V400
H400
V517.36
H282.65
V634.72
H165.29
Z
M282.65 282.65
V400
H400
V282.65
Z
"/>
<path fill="#fff" d="M517.36 400 H634.72 V634.72 H517.36 Z"/>
</svg>

After

Width:  |  Height:  |  Size: 475 B

View file

@ -362,10 +362,12 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
}
}, [isReasoningStreaming]);
// Reset dismissed flag on new stream.
// Reset per-round open state. manualOpen is sticky and regenerate reuses this
// instance, so a hand-opened block would stay pinned open and never collapse.
useEffect(() => {
if (isReasoningStreaming) {
setDismissedWhileStreaming(false);
setManualOpen(false);
}
}, [isReasoningStreaming]);

View file

@ -40,10 +40,9 @@ interface ApiProviderLogoProps {
title?: string;
}
/**
* Renders a registry provider's logo when its asset exists under
* `public/provider-logos/`. OpenAI's is inverted in dark mode for contrast.
*/
const DARK_INVERT_LOGOS = new Set(["openai", "ollama", "openrouter"]);
/** Provider logo from `public/provider-logos/`; monochrome ones invert in dark mode. */
export function ApiProviderLogo({ providerType, className, title }: ApiProviderLogoProps) {
const src = apiProviderLogoSrc(providerType);
if (!src && isCustomProviderType(providerType)) {
@ -63,7 +62,7 @@ export function ApiProviderLogo({ providerType, className, title }: ApiProviderL
aria-hidden
className={cn(
"shrink-0 object-contain",
providerType === "openai" && "dark:invert",
providerType && DARK_INVERT_LOGOS.has(providerType) && "dark:invert",
className,
)}
/>

View file

@ -3762,12 +3762,15 @@ export function createOpenAIStreamAdapter(
// Permission level for local tool calls is sent for every local
// chat, not only when a tool pill is on: a process policy
// (unsloth run --enable-tools) can open the tool loop with no pill,
// and the backend must still see the selected gate. ask/auto request
// the confirm gate ("auto" only pauses calls flagged unsafe); off
// and full never prompt, full also drops the sandbox.
// and the backend must still see the selected gate. "auto" OMITS
// confirm_tool_calls: an explicit true would make the backend treat
// every auto request as needing a stream and defeat the safe-only
// no-stream exception. "ask" sends true; off/full send false (full
// also drops the sandbox).
permission_mode: permissionMode,
confirm_tool_calls:
permissionMode === "ask" || permissionMode === "auto",
...(permissionMode === "auto"
? {}
: { confirm_tool_calls: permissionMode === "ask" }),
bypass_permissions: bypassPermissions,
...(supportsTools &&
(toolsEnabled ||

View file

@ -348,6 +348,9 @@ export interface LocalModelInfo {
// Backend-detected weights format ("gguf" when known), so the UI can
// classify scanned folders whose name lacks a -GGUF suffix.
model_format?: string | null;
// Set when a cached snapshot holds an incomplete download, so consumers can skip
// weights that cannot load yet.
partial?: boolean;
updated_at?: number | null;
}

View file

@ -11,9 +11,11 @@ export {
fetchGgufStagedMetadata,
getCachedModelPath,
getInferenceStatus,
listCachedGguf,
listChatAttachments,
listGgufVariants,
listLocalModels,
listModels,
listRecommendedFolders,
listScanFolders,
loadModel,
@ -28,7 +30,11 @@ export {
type LocalModelInfo,
type ScanFolderInfo,
} from "./api/chat-api";
export type { GgufVariantDetail } from "./types/api";
export type {
BackendModelDetails,
GgufVariantDetail,
InferenceStatusResponse,
} from "./types/api";
export {
ChatSettingsPanel,
ParamSlider,

View file

@ -52,7 +52,8 @@ export const PERMISSION_MODE_OPTIONS: readonly {
{
value: "auto",
label: "Approve for me",
description: "Only ask for actions detected as potentially unsafe",
description:
"Run tool calls, but ask before high-risk actions like credential access, privilege escalation, or destructive commands",
icon: ShieldCheck,
},
{
@ -76,6 +77,8 @@ export const FULL_ACCESS_WARNING =
export function permissionModeOption(mode: PermissionMode) {
return (
PERMISSION_MODE_OPTIONS.find((option) => option.value === mode) ??
// Unknown values fall back to the default ("Approve for me"), not row 0 ("Ask").
PERMISSION_MODE_OPTIONS.find((option) => option.value === "auto") ??
PERMISSION_MODE_OPTIONS[0]
);
}

View file

@ -51,8 +51,8 @@ export const CHAT_PERMISSION_MODE_KEY = "unsloth_chat_permission_mode";
/**
* Permission level for local tool calls:
* - "ask": always ask before every tool call runs.
* - "auto" ("Approve for me"): only ask for calls the backend detects as
* potentially unsafe; read-only calls run immediately. Sandbox stays on.
* - "auto" ("Approve for me", the default): only ask for calls the backend
* detects as high risk; ordinary dev commands run immediately. Sandbox stays on.
* - "off": never ask; tool calls run automatically inside the sandbox
* (the original default before permission levels existed).
* - "full" ("Full access"): no confirmations and the python/terminal sandbox

View file

@ -115,7 +115,7 @@ export interface GgufVariantDetail {
download_size_bytes?: number;
downloaded?: boolean;
update_available?: boolean;
/** True while an in-progress (.incomplete) blob exists for this variant. */
/** An interrupted download: some shards are missing, so it cannot load yet. */
partial?: boolean;
}
@ -171,7 +171,10 @@ export interface LoadModelResponse {
max_context_length?: number | null;
native_context_length?: number | null;
supports_reasoning?: boolean;
reasoning_style?: "enable_thinking" | "reasoning_effort" | "enable_thinking_effort";
reasoning_style?:
| "enable_thinking"
| "reasoning_effort"
| "enable_thinking_effort";
reasoning_effort_levels?: string[];
reasoning_always_on?: boolean;
supports_preserve_thinking?: boolean;
@ -222,7 +225,10 @@ export interface InferenceStatusResponse {
} | null;
requires_trust_remote_code?: boolean;
supports_reasoning?: boolean;
reasoning_style?: "enable_thinking" | "reasoning_effort" | "enable_thinking_effort";
reasoning_style?:
| "enable_thinking"
| "reasoning_effort"
| "enable_thinking_effort";
reasoning_effort_levels?: string[];
reasoning_always_on?: boolean;
supports_preserve_thinking?: boolean;
@ -391,7 +397,7 @@ export interface OpenAIChatCompletionsRequest {
| "xhigh"
| null;
preserve_thinking?: boolean | null;
thinking?: {type: "disabled" | "enabled";} | null;
thinking?: { type: "disabled" | "enabled" } | null;
enable_tools?: boolean | null;
enabled_tools?: string[];
/** Local models + enable_tools only. */

View file

@ -35,7 +35,7 @@
{
"column_type": "llm-structured",
"name": "llm_structured_1",
"drop": false,
"drop": true,
"model_alias": "provider_column",
"prompt": "Given ONLY this chunk: {{ chunk_text }} generate one answerable question, answer, and exact supporting quote from chunk. If not answerable, skip.",
"with_trace": "none",
@ -43,11 +43,7 @@
"output_format": {
"type": "object",
"additionalProperties": false,
"required": [
"question",
"answer",
"evidence_quote"
],
"required": ["question", "answer", "evidence_quote"],
"properties": {
"question": {
"type": "string"
@ -60,16 +56,41 @@
}
}
}
},
{
"column_type": "expression",
"name": "instruction",
"drop": false,
"expr": "{{ llm_structured_1.question }}",
"dtype": "str"
},
{
"column_type": "expression",
"name": "output",
"drop": false,
"expr": "{{ llm_structured_1.answer }}",
"dtype": "str"
},
{
"column_type": "expression",
"name": "input",
"drop": false,
"expr": "Evidence quote: {{ llm_structured_1.evidence_quote }}\n\nSource context: {{ chunk_text }}",
"dtype": "str"
}
],
"processors": []
"processors": [
{
"processor_type": "drop_columns",
"name": "drop_seed_columns",
"column_names": ["chunk_text", "source_file"]
}
]
},
"run": {
"rows": 5,
"preview": true,
"output_formats": [
"jsonl"
]
"output_formats": ["jsonl"]
},
"ui": {
"nodes": [
@ -102,7 +123,7 @@
"width": 400,
"node_type": "markdown_note",
"name": "note_3",
"markdown": "- LLM prompt: `{{ chunk_text }}`\n- Expression block: combine/format values using `{{ chunk_text }}`\n- Processor templates: use `{{ chunk_text }}` during transforms\n\nTip:\n- Start with medium chunk size + small overlap.\n- Increase overlap only if answers lose context between chunks.",
"markdown": "The structured LLM block generates a question, answer, and evidence quote from `{{ chunk_text }}`.\n\nExpression blocks then project the result into a training-ready Alpaca row:\n\n- `instruction`: generated question\n- `input`: evidence quote and source context\n- `output`: generated answer\n\nThe source chunk, source-file field, and nested structured intermediate are dropped only after these fields are created.",
"note_color": "#F3E8FF",
"note_opacity": "35"
},
@ -129,6 +150,24 @@
"x": 960,
"y": 1077,
"width": 400
},
{
"id": "instruction",
"x": 1440,
"y": 895,
"width": 400
},
{
"id": "output",
"x": 1440,
"y": 1077,
"width": 400
},
{
"id": "input",
"x": 1440,
"y": 1259,
"width": 400
}
],
"edges": [
@ -147,11 +186,39 @@
"target_handle": "data-in-top"
},
{
"from": "llm_structured_1",
"to": "seed",
"from": "seed",
"to": "llm_structured_1",
"type": "canvas",
"source_handle": "data-out-left",
"target_handle": "data-in-right"
"source_handle": "data-out",
"target_handle": "data-in"
},
{
"from": "llm_structured_1",
"to": "instruction",
"type": "canvas",
"source_handle": "data-out",
"target_handle": "data-in"
},
{
"from": "llm_structured_1",
"to": "output",
"type": "canvas",
"source_handle": "data-out",
"target_handle": "data-in"
},
{
"from": "llm_structured_1",
"to": "input",
"type": "canvas",
"source_handle": "data-out",
"target_handle": "data-in"
},
{
"from": "seed",
"to": "input",
"type": "canvas",
"source_handle": "data-out",
"target_handle": "data-in"
}
],
"layout_direction": "LR",
@ -164,4 +231,4 @@
"unstructured_chunk_size": "1200",
"unstructured_chunk_overlap": "200"
}
}
}

View file

@ -404,6 +404,10 @@ export function importRecipePayload(
uiSeedSourceTypeRaw === "unstructured"
? uiSeedSourceTypeRaw
: undefined;
const payloadSeedSourceIsUnstructured =
isRecord(recipe.seed_config) &&
isRecord(recipe.seed_config.source) &&
recipe.seed_config.source.seed_type === "unstructured";
const uiSeedColumns = Array.isArray(ui?.seed_columns)
? ui.seed_columns
.map((value) => (typeof value === "string" ? value.trim() : ""))
@ -478,7 +482,17 @@ export function importRecipePayload(
nextId += 1;
const seedConfig = parseSeedConfig(recipe.seed_config, id, {
preferredSourceType: uiSeedSourceType,
seed_columns: uiSeedColumns,
drop:
payloadSeedSourceIsUnstructured && payloadSeedDropColumns.length > 0,
// Payload-only unstructured recipes have no preview metadata, but their
// generated rows always expose these fields. Keep the imported drop
// processor usable until a real preview replaces this fallback.
seed_columns:
(uiSeedColumns?.length ?? 0) > 0
? uiSeedColumns
: uiSeedSourceType === "unstructured" || payloadSeedSourceIsUnstructured
? ["chunk_text", "source_file"]
: uiSeedColumns,
seed_drop_columns:
uiSeedDropColumns && uiSeedDropColumns.length > 0
? uiSeedDropColumns

View file

@ -193,6 +193,7 @@ export function parseSeedConfig(
id: string,
options?: {
preferredSourceType?: SeedSourceType;
drop?: boolean;
seed_columns?: string[];
seed_drop_columns?: string[];
seed_preview_rows?: Record<string, unknown>[];
@ -229,6 +230,7 @@ export function parseSeedConfig(
...makeDefaultSeedConfig(id),
...parsed, // payload-only fields override ui defaults
seed_source_type: sourceType,
...(options?.drop !== undefined ? { drop: options.drop } : {}),
...(options?.seed_columns ? { seed_columns: options.seed_columns } : {}),
...(options?.seed_drop_columns
? { seed_drop_columns: options.seed_drop_columns }

View file

@ -164,17 +164,24 @@ export function buildSeedDropProcessor(
): Record<string, unknown> | null {
const seedSourceType = config.seed_source_type ?? "hf";
const loadedCols = (config.seed_columns ?? []).map((c) => c.trim()).filter(Boolean);
const selectedDropColumns = (config.seed_drop_columns ?? [])
.map((c) => c.trim())
.filter(Boolean);
let cols: string[] = [];
if (seedSourceType === "unstructured") {
if (!config.drop) {
return null;
}
cols = loadedCols;
cols =
selectedDropColumns.length > 0
? loadedCols.length > 0
? selectedDropColumns.filter((col) => loadedCols.includes(col))
: selectedDropColumns
: loadedCols.length > 0
? loadedCols
: ["chunk_text", "source_file"];
} else {
const selectedDropColumns = (config.seed_drop_columns ?? [])
.map((c) => c.trim())
.filter(Boolean);
if (selectedDropColumns.length === 0) {
return null;
}

View file

@ -141,8 +141,9 @@ const AGENT_LABELS: Record<string, string> = {
};
const j = (s: string): string => JSON.stringify(s);
const shSingle = (s: string): string => s.replace(/'/g, "'\\''");
const psSingle = (s: string): string => s.replace(/'/g, "''");
// Inner escaping for a single-quoted argument (POSIX '\'' , PowerShell '').
export const shSingle = (s: string): string => s.replace(/'/g, "'\\''");
export const psSingle = (s: string): string => s.replace(/'/g, "''");
const toolsJson = TOOLS.map(j).join(", ");
function bodyExtraLines(variant: Variant, indent: string): string[] {

View file

@ -12,6 +12,7 @@ import { type TranslationKey, useT } from "@/i18n";
import { cn } from "@/lib/utils";
import { MicIcon } from "@/lib/mic-icon";
import {
BotIcon,
Cancel01Icon,
CloudIcon,
CpuIcon,
@ -40,6 +41,7 @@ import {
useSettingsDialogStore,
} from "./stores/settings-dialog-store";
import { AboutTab } from "./tabs/about-tab";
import { AgentsTab } from "./tabs/agents-tab";
import { ApiKeysTab } from "./tabs/api-keys-tab";
import { AppearanceTab } from "./tabs/appearance-tab";
import { ChatTab } from "./tabs/chat-tab";
@ -71,13 +73,11 @@ const TABS: TabDef[] = [
id: "resources",
labelKey: "settings.tabs.resources",
icon: CpuIcon,
badgeKey: "common.new",
},
{
id: "chat",
labelKey: "settings.tabs.chat",
icon: Message01Icon,
badgeKey: "common.new",
},
{
id: "api-keys",
@ -89,6 +89,12 @@ const TABS: TabDef[] = [
labelKey: "settings.tabs.connections",
icon: CloudIcon,
},
{
id: "agents",
labelKey: "settings.tabs.agents",
icon: BotIcon,
badgeKey: "common.new",
},
{
id: "voice",
labelKey: "settings.tabs.voice",
@ -124,6 +130,8 @@ function renderTab(tab: SettingsTab) {
return <DataTab />;
case "api-keys":
return <ApiKeysTab />;
case "agents":
return <AgentsTab />;
case "about":
return <AboutTab />;
}
@ -222,6 +230,7 @@ export function SettingsDialog() {
connections: null,
data: null,
"api-keys": null,
agents: null,
about: null,
});

View file

@ -103,6 +103,21 @@ export const SETTINGS_SEARCH_INDEX: Record<SettingsTab, TranslationKey[]> = {
"settings.apiKeys.description",
"settings.apiKeys.accessTokens",
],
agents: [
// Every key needs a rendered data-settings-label, or a hit has nothing to scroll to.
"settings.agents.title",
"settings.agents.description",
"settings.agents.intro",
"settings.agents.agent",
"settings.agents.model",
"settings.agents.quantization",
// subagent.title is deliberately absent: its label only mounts for the agents
// that support subagents, so a hit would have nothing to scroll to otherwise.
"settings.agents.options.title",
"settings.agents.remote.title",
"settings.agents.passthrough.title",
"settings.agents.dryRun.title",
],
connections: [],
voice: [
"settings.voice.dictation.sectionTitle",

View file

@ -13,6 +13,7 @@ export type SettingsTab =
| "connections"
| "data"
| "api-keys"
| "agents"
| "about";
export type SettingsScrollTarget = "about-updates";
@ -69,6 +70,7 @@ function loadInitialTab(): SettingsTab {
"connections",
"data",
"api-keys",
"agents",
"about",
];
return valid.includes(stored as SettingsTab)

File diff suppressed because it is too large Load diff

View file

@ -101,6 +101,7 @@ export const en = {
connections: "Connections",
data: "Data",
apiKeys: "API",
agents: "Agents",
about: "About",
},
voice: {
@ -162,7 +163,8 @@ export const en = {
},
dictionary: {
sectionTitle: "Dictation dictionary",
sectionDescription: "Set how dictation spells specific words or phrases",
sectionDescription:
"Set how dictation spells specific words or phrases",
manageLabel: "Custom spellings",
manage: "Manage",
backToVoice: "Back to Voice",
@ -466,7 +468,8 @@ export const en = {
"Unsupported file type. Use .woff2, .woff, .ttf, or .otf.",
errorTooLarge: "Font file is too large (max 1.5 MB).",
errorLimit: "You can import up to 3 fonts.",
errorStorageFull: "Not enough local storage for this font. Remove an imported font first.",
errorStorageFull:
"Not enough local storage for this font. Remove an imported font first.",
errorFailed: "Could not load this font file.",
},
uiFontSize: {
@ -580,6 +583,100 @@ export const en = {
unknown: "Unknown",
},
},
agents: {
title: "Agents",
description:
"Connect coding agents like Claude Code and Codex to a model running locally in Unsloth with unsloth start.",
intro:
"connects Claude Code, Codex, Hermes, OpenClaw, OpenCode, Pi and other agents to a model served locally by Unsloth, fully offline on your own hardware. It runs an OpenAI-compatible server for the agent and never touches your agent's config files.",
readDocs: "Read the docs",
copy: "Copy",
copied: "Copied",
commandBuilder: "Command builder",
agent: "Coding agent",
model: "Model",
searchModels: "Search GGUF models...",
noModels: "No matching GGUF models.",
showingModels:
"Showing {shown} of {total} matches. Keep typing to narrow the list.",
quantization: "Quantization",
loadingQuantizations: "Loading quantizations...",
noQuantizations: "No separate quantization",
recommended: "Recommended",
downloaded: "Downloaded",
quantizationLoadError:
"Couldn't load all quantizations. The command will use the available model value.",
generatedCommand: "Generated command",
docs: "Docs",
agentDocs: "Open {agent} setup docs",
copyGeneratedCommand: "Copy generated command",
modelNote:
"Codex requires a GGUF model served by llama-server. Other agents can also use transformer-backed models; remove --model to use the model already loaded in Unsloth Studio.",
subagent: {
title: "Use a local model as a subagent",
description:
"Keep {agent} on its current model and delegate selected tasks to this local Unsloth model.",
setupCommand: "Setup command",
copySetupCommand: "Copy subagent setup command",
usagePrompt: "Then in {agent}, type:",
copyUsagePrompt: "Copy subagent usage prompt",
defaultPrompt: "Spawn a local agent to implement this function.",
opencodePrompt: "@unsloth find the cause of this test failure",
},
quickstart: {
title: "Build a command",
description:
"Launch an agent against the model currently loaded in Studio. Load a model first, then swap claude for any supported agent below.",
noneDetected: "No supported agent CLIs were found on your PATH.",
installed: "Installed",
},
supportedAgents: {
title: "Supported agents",
description: "Each agent launches with its own command:",
requiresGguf: "Needs a GGUF model",
},
models: {
title: "Choosing a model",
description:
"Pass --model to pick a model and quantization, and --context-length to set the window. Use a quantization suffix, or an explicit --gguf-variant flag.",
suffixLabel: "With a quantization suffix",
variantLabel: "With an explicit variant flag",
},
options: {
title: "Common options",
description:
"Unsloth flags are parsed first; anything it doesn't recognize is passed straight through to the agent.",
model:
"Select a model. Without --model, unsloth start uses the model currently loaded in Studio and errors if none is loaded.",
contextLength:
"Set the requested context length (alias: --max-seq-length).",
ggufVariant: "Choose the GGUF quantization variant.",
loadIn4bit: "Toggle 4-bit loading for Hugging Face models.",
tensorParallel: "Toggle tensor-parallel across multiple GPUs.",
serve: "Enable or disable the automatic local server.",
launch: "Launch the agent, or just print the command and environment.",
persist: "Keep Unsloth-managed agent storage between runs.",
asSubagent:
"Keep the parent on its current model and register Unsloth as a local subagent (Claude Code, Codex, OpenCode, and Pi).",
apiKey: "Provide your Unsloth API key (or set UNSLOTH_API_KEY).",
yolo: "Skip approval prompts. Use only in trusted environments.",
},
remote: {
title: "Connect to a remote Studio",
description:
"Point unsloth start at a Studio running elsewhere by setting these before launching (or pass --api-key directly):",
},
passthrough: {
title: "Passing agent arguments",
description:
"Arguments after the Unsloth flags are forwarded to the agent itself, so native commands like resume still work:",
},
dryRun: {
title: "Preview without launching",
description:
"Add --no-launch to print the environment and command instead of launching the agent. If --model is set, the model may still be resolved and loaded.",
},
},
chat: {
title: "Chat",
description: "Customize how chat behaves on this device.",

View file

@ -63,6 +63,7 @@ EXIT_SUCCESS = 0
EXIT_FALLBACK = 2
EXIT_ERROR = 1
EXIT_BUSY = 3
EXIT_NO_SPACE = 4
# DiskPart-prompt suppression. RunAsInvoker does NOT stop amd-smi's runtime
# elevation (its manifest is asInvoker), so this is just harmless belt-and-
@ -3674,7 +3675,8 @@ def hydrate_source_tree(
break
except Exception as exc:
last_exc = exc
if index == len(source_urls) - 1:
# A full disk fails every mirror; stop so a later 404 cannot mask it.
if _environment_fatal_reason(exc) or index == len(source_urls) - 1:
raise
log(f"source tree download failed from {source_url}: {exc}")
if not downloaded:
@ -6000,6 +6002,14 @@ def validate_prebuilt_attempts(
)
raise ExistingInstallSatisfied(attempt, tried_fallback)
# Advisory: a few GB free usually fits, and rejecting here would also skip
# the source-build fallback.
if index == 0:
low_disk = _low_disk_warning(install_dir)
if low_disk is not None:
log(low_disk)
_log_disk_space_help()
staging_dir = create_install_staging_dir(install_dir)
quantized_path = work_dir / f"stories260K-q4-{index}.gguf"
if quantized_path.exists():
@ -6028,7 +6038,9 @@ def validate_prebuilt_attempts(
attempt_error = PrebuiltFallback(
f"candidate attempt failed before activation for {attempt.name}: {exc}"
)
if index == len(attempt_list) - 1:
if _environment_fatal_reason(exc) or index == len(attempt_list) - 1:
if attempt_error is exc:
raise
raise attempt_error from exc
log(
"selected CUDA bundle failed before activation; trying next prebuilt fallback "
@ -6149,6 +6161,132 @@ def diffusion_visual_server_backfill_needed(
return True
def _causal_chain(exc: BaseException) -> Iterable[BaseException]:
seen: set[int] = set()
current: BaseException | None = exc
while current is not None and id(current) not in seen:
seen.add(id(current))
yield current
# `raise X from None` sets __suppress_context__: the earlier exception is
# unrelated, so following __context__ anyway would misreport the cause.
if current.__cause__ is not None:
current = current.__cause__
elif current.__suppress_context__:
current = None
else:
current = current.__context__
# ERROR_HANDLE_DISK_FULL / ERROR_DISK_FULL. CPython's PC/errmap.h maps 112 to
# ENOSPC but has no case for 39, which arrives as EINVAL, so check winerror too.
_WINDOWS_DISK_FULL = (39, 112)
# A quota (NFS/XFS/container) leaves blocks this user cannot have, so the bigger
# source build is just as doomed; named apart from ENOSPC so df does not mislead.
# Guarded: the MSVC CRT has no EDQUOT, so on Windows CPython aliases it to the
# Winsock WSAEDQUOT (10069), which no file write raises.
_DISK_FULL_ERRNOS = {errno.ENOSPC: "no space left on device"}
if hasattr(errno, "EDQUOT"):
_DISK_FULL_ERRNOS[errno.EDQUOT] = "disk quota exceeded"
def _winerror_of(exc: OSError) -> Any:
"""exc.winerror, defensively. Not getattr(exc, ..., None): urllib's HTTPError
is an OSError that proxies unknown attributes to a wrapped file object and
raises KeyError (not AttributeError) on 3.9, which getattr will not swallow.
A 404 from a mirror must not crash the classifier."""
try:
return exc.winerror
except Exception:
return None
def _out_of_space_reason(exc: BaseException) -> str | None:
"""Why `exc` means the install cannot fit, or None if it means something else."""
if isinstance(exc, OSError):
reason = _DISK_FULL_ERRNOS.get(exc.errno)
if reason is not None:
return reason
if _winerror_of(exc) in _WINDOWS_DISK_FULL:
return "no space left on device"
# shutil.copytree stringifies each per-file OSError and raises Error(errors)
# outside the except block, so errno and the chain are gone and only text
# survives. OSError.__str__ returns early on winerror, so Windows reads
# "[WinError 112]" and never "[Errno 28]": match both, brackets included so
# WinError 112 does not match WinError 1120.
if isinstance(exc, shutil.Error):
text = str(exc)
for code, reason in _DISK_FULL_ERRNOS.items():
if f"[Errno {code}]" in text:
return reason
if any(f"[WinError {code}]" in text for code in _WINDOWS_DISK_FULL):
return "no space left on device"
return None
def _environment_fatal_reason(exc: BaseException) -> str | None:
for cause in _causal_chain(exc):
reason = _out_of_space_reason(cause)
if reason is not None:
return reason
return None
def _log_disk_space_help() -> None:
log(
"free up space or point TMPDIR and UNSLOTH_STUDIO_HOME at a larger "
"volume (e.g. /workspace), then re-run"
)
@contextmanager
def scratch_dir(prefix: str) -> Iterator[Path]:
"""Temp dir whose cleanup never raises: an rmtree failure on the way out would
replace the in-flight exception and lose EXIT_NO_SPACE. Not
TemporaryDirectory(ignore_cleanup_errors = True), which is 3.10+ (setup.sh
still runs this helper under the host python, and we support 3.9)."""
path = Path(tempfile.mkdtemp(prefix = prefix))
try:
yield path
finally:
shutil.rmtree(path, ignore_errors = True)
def _first_existing_ancestor(path: Path) -> Path:
current = path
while current != current.parent and not current.exists():
current = current.parent
return current
def _low_disk_warning(install_dir: Path, *, advised_gb: float = 5.0) -> str | None:
"""Advisory only, never fatal. A prebuilt install peaks well under 1 GB (the
largest published bundle is 0.77 GB, macOS is 0.01 GB), so a fixed threshold
cannot decide whether this host has room -- a real ENOSPC decides that. The
number here is the headroom a source-build fallback would want."""
advised = int(advised_gb * (1024**3))
targets = {
"build/download scratch (TMPDIR)": Path(tempfile.gettempdir()),
"llama.cpp install dir": _first_existing_ancestor(install_dir),
}
for label, path in targets.items():
try:
free = shutil.disk_usage(path).free
except OSError:
continue
if free < advised:
return (
f"low disk space for llama.cpp: {label} at {path} has "
f"{free / (1024**3):.1f} GB free (~{advised_gb:.0f} GB recommended)"
)
return None
def _fail_no_space(reason: str) -> None:
log(reason)
_log_disk_space_help()
raise SystemExit(EXIT_NO_SPACE)
def install_prebuilt(
install_dir: Path,
llama_tag: str,
@ -6217,8 +6355,7 @@ def install_prebuilt(
# recorded so the updater re-asserts it (#7213).
sync_marker_force_cpu(install_dir, persist_force_cpu)
return
with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp:
work_dir = Path(tmp)
with scratch_dir("unsloth-llama-prebuilt-") as work_dir:
probe_path = work_dir / "stories260K.gguf"
download_validation_model(probe_path, validation_model_cache_path(install_dir))
release_count = len(release_plans)
@ -6263,6 +6400,8 @@ def install_prebuilt(
except ExistingInstallSatisfied:
return
except PrebuiltFallback as exc:
if _environment_fatal_reason(exc):
raise
if release_index == release_count - 1:
raise
log(
@ -6296,6 +6435,11 @@ def install_prebuilt(
log(f"prebuilt busy reason: {exc}")
raise SystemExit(EXIT_BUSY) from exc
except PrebuiltFallback as exc:
fatal = _environment_fatal_reason(exc)
if fatal:
log(f"prebuilt install failed: {fatal}")
_log_disk_space_help()
raise SystemExit(EXIT_NO_SPACE) from exc
log("prebuilt install path failed; falling back to source build")
log(f"prebuilt fallback reason: {exc}")
report = collect_system_report(host, choice, install_dir)
@ -6466,6 +6610,10 @@ def main() -> int:
install_kind = args.install_kind,
)
except PrebuiltFallback as exc:
# A full disk is not a bad build: the CPU source rebuild needs more space.
fatal = _environment_fatal_reason(exc)
if fatal:
_fail_no_space(f"install validation failed: {fatal}")
print(str(exc), file = sys.stderr)
raise SystemExit(EXIT_FALLBACK) from exc
return EXIT_SUCCESS
@ -6595,9 +6743,15 @@ if __name__ == "__main__":
# Expected when the published repo (e.g. ggml-org/llama.cpp) has no
# prebuilt manifest. Exit quietly with EXIT_FALLBACK so the caller
# falls back to source build without a noisy "fatal helper error".
fatal = _environment_fatal_reason(exc)
if fatal:
_fail_no_space(f"prebuilt install failed: {fatal}")
log(textwrap.shorten(str(exc), width = 400, placeholder = "..."))
raise SystemExit(EXIT_FALLBACK)
except Exception as exc:
fatal = _environment_fatal_reason(exc)
if fatal:
_fail_no_space(f"prebuilt install failed: {fatal}")
message = textwrap.shorten(str(exc), width = 400, placeholder = "...")
log(f"fatal helper error: {message}")
raise SystemExit(EXIT_ERROR)

View file

@ -3763,6 +3763,18 @@ if ($LocalLlamaCppLinked) {
}
substep "Close Unsloth or other llama.cpp users and retry" "Yellow"
exit 3
} elseif ($prebuiltExit -eq 4) {
step "llama.cpp" "not enough disk space to install llama.cpp" "Yellow"
Write-LlamaFailureLog -Output $prebuiltOutput
substep "Free up disk or move UNSLOTH_STUDIO_HOME/TEMP to a larger volume, then re-run" "Yellow"
$PreservedLlamaServerFound = $false
foreach ($_cand in @(
(Join-Path $LlamaCppDir "llama-server.exe"),
(Join-Path $LlamaCppDir "build\bin\llama-server.exe"),
(Join-Path $LlamaCppDir "build\bin\Release\llama-server.exe"))) {
if (Test-Path -LiteralPath $_cand) { $PreservedLlamaServerFound = $true; break }
}
if (-not $PreservedLlamaServerFound) { $script:LlamaCppDegraded = $true }
} else {
step "llama.cpp" "prebuilt install failed (continuing)" "Yellow"
Write-LlamaFailureLog -Output $prebuiltOutput

View file

@ -67,6 +67,15 @@ fi
step() { printf " ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; }
substep() { printf " %-15s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; }
# ── Helper: can the controlling terminal actually be opened for reading? ──
# `test -r` only checks permission bits, which look fine in containers and
# systemd units where open() then fails with ENXIO. Probe with a real open.
# Mirrors install.sh's _can_read_tty; defined here too because setup.sh runs
# as its own process (install.sh invokes it, it does not source it).
_can_read_tty() {
( : </dev/tty ) >/dev/null 2>&1
}
_is_verbose() {
[ "${UNSLOTH_VERBOSE:-0}" = "1" ]
}
@ -1224,6 +1233,7 @@ LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp"
LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
_NEED_LLAMA_SOURCE_BUILD=false
_LLAMA_CPP_DEGRADED=false
_LLAMA_CPP_NO_SPACE=false
_LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}"
_HOST_SYSTEM="$(uname -s 2>/dev/null || true)"
@ -1451,6 +1461,13 @@ else
fi
substep "close Unsloth or other llama.cpp users and retry"
exit 3
elif [ "$_PREBUILT_STATUS" -eq 4 ]; then
step "llama.cpp" "not enough disk space to install llama.cpp" "$C_WARN"
print_llama_error_log "$_PREBUILT_LOG"
rm -f "$_PREBUILT_LOG"
substep "free up disk or move UNSLOTH_STUDIO_HOME/TMPDIR to a larger volume, then re-run"
_LLAMA_CPP_NO_SPACE=true
_has_local_llama_server "$LLAMA_CPP_DIR" || _LLAMA_CPP_DEGRADED=true
else
step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
print_llama_error_log "$_PREBUILT_LOG"
@ -1502,25 +1519,46 @@ if [ "$_NEED_LLAMA_SOURCE_BUILD" = true ] && grep -qi microsoft /proc/version 2>
step "gguf deps" "installed"
elif command -v sudo >/dev/null 2>&1; then
step "gguf deps" "sudo required for: $_STILL_MISSING" "$C_WARN"
printf " %-15s" ""
printf "accept? [Y/n] "
if [ -r /dev/tty ]; then
read -r REPLY </dev/tty || REPLY="y"
if _can_read_tty; then
printf " %-15s" ""
printf "accept? [Y/n] "
# The device opened, so a failed read is EOF, not consent: decline.
read -r REPLY </dev/tty || REPLY="n"
case "$REPLY" in
[nN]*)
substep "skipped -- run manually:"
substep "sudo apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
;;
*)
# Degrade like the no-sudo branch below rather than letting
# set -e abort setup on a bare apt error: missing GGUF build
# deps are recoverable, not fatal.
if sudo apt-get update -y </dev/null &&
sudo apt-get install -y $_STILL_MISSING </dev/null; then
step "gguf deps" "installed"
else
step "gguf deps" "install failed -- run manually:" "$C_WARN"
substep "sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
fi
;;
esac
else
REPLY="y"
fi
case "$REPLY" in
[nN]*)
substep "skipped -- run manually:"
substep "sudo apt-get install -y $_STILL_MISSING"
# Nobody can answer a prompt or type a password here, so -n makes
# sudo refuse rather than prompt into a closed stdin, and -k ignores
# any cached timestamp so only a real NOPASSWD rule gets through.
# Same treatment as install.sh's _smart_apt_install. This is the WSL
# GGUF-export case noted above, where sudo does want a password.
if sudo -n -k apt-get update -y </dev/null &&
sudo -n -k apt-get install -y $_STILL_MISSING </dev/null; then
step "gguf deps" "installed (non-interactive sudo)"
else
step "gguf deps" "needs sudo, no terminal -- run manually:" "$C_WARN"
substep "sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
;;
*)
sudo apt-get update -y
sudo apt-get install -y $_STILL_MISSING
step "gguf deps" "installed"
;;
esac
fi
fi
else
step "gguf deps" "missing (no sudo) -- install manually:" "$C_WARN"
substep "apt-get install -y $_STILL_MISSING"
@ -1946,7 +1984,14 @@ else
--validate-install "$_BUILD_TMP"
)
[ -n "$_SMOKE_KIND" ] && _SMOKE_CMD+=(--install-kind "$_SMOKE_KIND")
if ! run_quiet_no_exit "validate source llama.cpp" "${_SMOKE_CMD[@]}"; then
_SMOKE_RC=0
run_quiet_no_exit "validate source llama.cpp" "${_SMOKE_CMD[@]}" || _SMOKE_RC=$?
# Exit 4 is a full disk, not a bad build: the CPU rebuild needs even
# more space, so keep what we already have.
if [ "$_SMOKE_RC" -eq 4 ]; then
substep "not enough disk space to validate the $_FB_LABEL build; keeping it" "$C_WARN"
_LLAMA_CPP_NO_SPACE=true
elif [ "$_SMOKE_RC" -ne 0 ]; then
substep "$_FB_LABEL source build failed smoke test; retrying CPU build..." "$C_WARN"
_TRY_METAL_CPU_FALLBACK=false
rm -rf "$_BUILD_TMP/build"
@ -2003,8 +2048,10 @@ fi # end _SKIP_GGUF_BUILD check
# An arm64 Linux GPU host source-builds for the GPU above. If that produced no
# binary, install the fork's arm64 CPU prebuilt (app-<tag>-linux-arm64-cpu.tar.gz)
# instead of leaving the host without llama.cpp. --cpu-fallback drops the GPU
# attributes so the CPU bundle is selected rather than re-attempting CUDA.
# attributes so the CPU bundle is selected rather than re-attempting CUDA. Skipped
# on a full disk: the retry fails the same way and buries the hint.
if [ "$_LLAMA_CPP_DEGRADED" = true ] \
&& [ "$_LLAMA_CPP_NO_SPACE" != true ] \
&& [ "$_HOST_SYSTEM" = "Linux" ] \
&& { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; }; then
substep "GPU source build unavailable; trying arm64 CPU prebuilt..."

View file

@ -22,6 +22,20 @@ def apply() -> None:
if getattr(torch.cuda, "_unsloth_consolidated_spoof", False):
return
# Settle bitsandbytes against the real torch first. Its __init__ does
# `if torch.cuda.is_available(): from .backends.cuda import ops`, and that
# module reads torch._C._cuda_getCurrentRawStream at import. On a CPU-only
# wheel that attribute is absent, so a bitsandbytes imported AFTER this
# spoof raises AttributeError (or OSError hunting libhipblas for the ROCm
# spoof) rather than ImportError, which slips past the `except ImportError`
# guards its importers use. Importing it here, while is_available() is
# still False, caches the CPU path in sys.modules for everything that
# follows.
try:
import bitsandbytes # noqa: F401
except Exception:
pass
# Device probes (cheap, value-returning)
torch.cuda.is_available = lambda: True
torch.cuda.device_count = lambda: 1

View file

@ -0,0 +1,190 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
# ruff: noqa
"""GRPO smoke test for the ``fast_inference=True`` vLLM rollout path.
Exercises the vLLM LoRA activation path (`WorkerLoRAManager`) that regressed on
vLLM >= 0.25.0 (unsloth#7283): the stacked `WeightsMapper` collapsed q/k/v and
gate/up LoRA weights onto one key, crashing adapter activation with
`IndexError`. All seven attention and MLP projections are LoRA targets so both
the fused `qkv_proj` and `gate_up_proj` families are covered.
Kept deliberately tiny so it finishes in well under a minute: a 0.6B model,
`enforce_eager`, no torch.compile, three short training steps, and short
prompts/completions. Seeded, so the asserted metrics are reproducible.
Run directly (`python tests/fast_inference/test_fast_inference.py`) or via
pytest; it skips automatically when no CUDA device is present.
"""
import math
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).parents[2]
sys.path.insert(0, str(REPO_ROOT))
import pytest
import torch
from tests.utils import header_footer_context
MODEL_NAME = "unsloth/Qwen3-0.6B"
MAX_SEQ_LENGTH = 256
LORA_RANK = 8
NUM_GENERATIONS = 2
MAX_PROMPT_LENGTH = 64
MAX_COMPLETION_LENGTH = 16
# >1 so the updated LoRA adapter is re-synced into vLLM on every step, not just
# loaded once; that repeat sync is the path that regressed.
MAX_STEPS = 3
GPU_MEMORY_UTILIZATION = 0.3
COMPILATION_CONFIG = 0
# Pins torch's global RNG (via the Trainer's set_seed), which the colocated vLLM
# sampler draws from, so the rollout and every metric below is reproducible.
SEED = 42
# Loose sanity bounds, not fitted values: they catch divergence and degenerate
# rollouts while staying valid across GPUs, models and vLLM versions.
MAX_CHARS_PER_TOKEN = 20
MAX_GRAD_NORM = 1e3
MAX_KL = 1.0
# All attention + MLP projections, so both fused vLLM LoRA families (qkv_proj and
# gate_up_proj) are exercised -- the >= 0.25.0 collision hit both.
TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
SYSTEM_PROMPT = "Respond concisely."
QUESTIONS = ["What is the capital of France?", "What is 2 + 2?"]
PROMPTS = [
[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": q}]
for q in QUESTIONS
]
def length_reward_func(completions, **kwargs) -> list[float]:
"""Reward longer completions. The fractional tie-break keeps rewards distinct
even if the model samples equal-length completions, so GRPO advantages are
never all-zero and the step stays meaningful on any vLLM/GPU combination."""
n = len(completions)
return [float(len(c[0]["content"])) + i / (n + 1) for i, c in enumerate(completions)]
def _metric(metrics, *names):
"""First present key; TRL spells some metrics differently across versions."""
for name in names:
if name in metrics:
return metrics[name]
return None
@pytest.mark.skipif(not torch.cuda.is_available(), reason = "fast_inference needs a CUDA GPU + vLLM")
def test_fast_inference():
# Import here, not at module load: importing unsloth probes for an
# accelerator and errors on CPU-only machines, so deferring keeps pytest
# collection and the skip path import-free. Unsloth must precede TRL.
from unsloth import FastLanguageModel
from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer
with header_footer_context("Load model (fast_inference=True)"):
model, tokenizer = FastLanguageModel.from_pretrained(
model_name = MODEL_NAME,
max_seq_length = MAX_SEQ_LENGTH,
load_in_4bit = False,
fast_inference = True,
max_lora_rank = LORA_RANK,
gpu_memory_utilization = GPU_MEMORY_UTILIZATION,
enforce_eager = True, # skip CUDA graph capture for fast startup
compilation_config = COMPILATION_CONFIG,
)
assert hasattr(model, "vllm_engine"), "fast_inference=True did not attach a vLLM engine"
model = FastLanguageModel.get_peft_model(
model,
r = LORA_RANK,
target_modules = TARGET_MODULES,
lora_alpha = LORA_RANK,
use_gradient_checkpointing = False,
random_state = SEED,
)
dataset = Dataset.from_dict({"prompt": PROMPTS})
with header_footer_context("GRPO config and trainer"):
training_args = GRPOConfig(
learning_rate = 5e-6,
per_device_train_batch_size = NUM_GENERATIONS,
gradient_accumulation_steps = 1,
num_generations = NUM_GENERATIONS,
max_prompt_length = MAX_PROMPT_LENGTH,
max_completion_length = MAX_COMPLETION_LENGTH,
max_steps = MAX_STEPS,
logging_steps = 1,
report_to = "none",
seed = SEED,
)
trainer = GRPOTrainer(
model = model,
processing_class = tokenizer,
reward_funcs = [length_reward_func],
args = training_args,
train_dataset = dataset,
)
# The trainer must actually route rollouts through vLLM, otherwise it would
# fall back to HF generation and never exercise WorkerLoRAManager.
assert trainer.args.use_vllm, "GRPO is not configured to use vLLM"
assert getattr(trainer, "llm", None) is not None, "GRPO did not bind a vLLM engine"
with header_footer_context("GRPO train (vLLM LoRA rollout)"):
trainer_stats = trainer.train()
assert trainer_stats is not None, "trainer.train() returned None"
assert trainer_stats.global_step == MAX_STEPS, "GRPO ran the wrong number of steps"
assert math.isfinite(trainer_stats.training_loss), "training loss is not finite"
# Without these, a rollout that silently produced nothing, or an update that
# diverged to NaN, would still pass the wiring assertions above.
steps = [log for log in trainer.state.log_history if "loss" in log]
assert len(steps) == MAX_STEPS, f"expected {MAX_STEPS} logged steps, got {len(steps)}"
# Every reward is a completion's character count, so this bounds reward and
# its spread without hard-coding model-specific values.
max_reward = MAX_COMPLETION_LENGTH * MAX_CHARS_PER_TOKEN
for i, step in enumerate(steps, start = 1):
loss = step["loss"]
grad_norm = step.get("grad_norm")
reward = step.get("reward")
zero_std = step.get("frac_reward_zero_std")
kl = step.get("kl")
# Key names differ across the supported TRL range, so accept either.
length = _metric(step, "completion_length", "completions/mean_length")
reward_std = _metric(step, "reward_std", "rewards/std")
assert math.isfinite(loss), f"step {i}: loss not finite ({loss})"
assert grad_norm is not None, f"step {i}: no grad_norm logged"
assert math.isfinite(grad_norm), f"step {i}: grad_norm not finite ({grad_norm})"
# Sign check only: a step can legitimately be near zero (0.004 observed),
# so any tighter lower bound would be flaky.
assert 0.0 < grad_norm < MAX_GRAD_NORM, f"step {i}: grad_norm {grad_norm}"
assert length is not None, f"step {i}: no completion length logged"
assert 0.0 < length <= MAX_COMPLETION_LENGTH, f"step {i}: empty rollout ({length})"
assert reward is not None, f"step {i}: no reward logged"
assert 0.0 < reward <= max_reward, f"step {i}: reward {reward} out of range"
assert reward_std is not None, f"step {i}: no reward_std logged"
assert 0.0 < reward_std <= max_reward, f"step {i}: no reward spread ({reward_std})"
assert zero_std in (None, 0.0), f"step {i}: {zero_std} of groups had no spread"
assert kl is None or math.isfinite(kl), f"step {i}: kl not finite ({kl})"
assert kl is None or abs(kl) < MAX_KL, f"step {i}: kl diverged ({kl})"
print("fast_inference GRPO rollout completed:", trainer_stats)
if __name__ == "__main__":
if torch.cuda.is_available():
test_fast_inference()
else:
print("Skipping fast_inference test: needs a CUDA GPU + vLLM")

View file

@ -89,6 +89,184 @@ assert_contains "mentions apt-get" "$_smart" 'sudo apt-get'
assert_contains "mentions official repos" "$_smart" "official repositories"
assert_contains "rejects tarball worry" "$_smart" "not a third-party tarball"
# ── No-TTY sudo escalation (#7307 Problem 7) ────────────────────────
# The old code assumed consent when /dev/tty was unreadable, then ran sudo with
# stdin closed, so a password-requiring host died on a raw sudo error. Drive the
# real function with /dev/tty rewritten to a fixture, the same trick used for
# /etc/os-release above, so every TTY state is reachable hermetically.
echo "=== _smart_apt_install no-TTY escalation ==="
# Closest portable stand-in for the /dev/tty inside containers and systemd
# units: the mode bits satisfy `test -r`, but open() fails with ENXIO. Callers
# must verify the shape before relying on it.
make_unopenable() {
python3 -c 'import socket,sys; socket.socket(socket.AF_UNIX).bind(sys.argv[1])' \
"$1" 2>/dev/null
}
# $1 tty: "tty" | "notty" | "unopenable"
# $2 sudo: "nopasswd" | "needspasswd" | "aptneedspasswd" | "cached" | "absent"
run_smart() {
_tty_mode="$1"; _sudo_mode="$2"
_d=$(mktemp -d -p "$_TMP_ROOT")
case "$_tty_mode" in
tty) printf 'y\n' > "$_d/tty" ;;
# Opens fine but reads EOF straight away (drained/half-closed
# terminal): openable is not the same as answerable.
eof) : > "$_d/tty" ;;
unopenable) make_unopenable "$_d/tty" ;;
esac
_f=$(mktemp -p "$_TMP_ROOT")
sed -n -e '/^_can_read_tty()/,/^}/p' \
-e '/^_smart_apt_install()/,/^}/p' "$INSTALL_SH" \
| sed -e "s#/dev/tty#$_d/tty#g" > "$_f"
(
TAURI_MODE=false
_apt_distro_description() { echo "TestOS 1.0 (debian-like)"; }
_is_pkg_installed() { return 1; } # nothing ever installs
apt-get() { return 1; } # unprivileged attempt fails
command() {
if [ "$1" = -v ] && [ "$2" = sudo ]; then
[ "$_sudo_mode" != absent ]; return $?
fi
builtin command "$@"
}
# Models real sudo: -n refuses (exit 1, nothing runs) when a password
# would be needed. -k ignores any cached timestamp for this invocation
# (sudo(8)), so only a real NOPASSWD rule counts as passwordless.
sudo() {
_noninteractive=false
_ignore_cache=false
while :; do
case "$1" in
-n) _noninteractive=true; shift ;;
-k) _ignore_cache=true; shift ;;
*) break ;;
esac
done
if [ "$_noninteractive" = true ]; then
case "$_sudo_mode" in
nopasswd) ;;
# A valid timestamp from an earlier, unrelated sudo. Without
# -k this looks passwordless; with -k it must not.
cached) [ "$_ignore_cache" = true ] && return 1 ;;
# Authorized for everything, NOPASSWD only on trivial
# commands: `sudo -l` says yes while execution still needs
# a password. Authorization is not the question to ask.
aptneedspasswd)
case " $* " in
*" apt-get "*) return 1 ;;
esac
;;
*) return 1 ;;
esac
fi
# Sudoers refuses the command outright, with or without -n.
if [ "$_sudo_mode" = denied ]; then
echo "sudo: user is not allowed to execute that" >&2
return 1
fi
echo "SUDO_RAN: $*"
}
# shellcheck disable=SC1090
. "$_f"
_smart_apt_install cmake 2>&1
echo "EXIT:$?"
) || true
}
_out=$(run_smart notty needspasswd)
assert_contains "no tty + password sudo: says it cannot run unattended" \
"$_out" "cannot be done unattended"
assert_contains "no tty + password sudo: gives the manual command" \
"$_out" "sudo apt-get update -y && sudo apt-get install -y cmake"
assert_contains "no tty + password sudo: names the distro" \
"$_out" "TestOS 1.0 (debian-like)"
case "$_out" in
*SUDO_RAN*) echo " FAIL: no tty + password sudo must not run apt-get as root"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: no tty + password sudo runs nothing as root"; PASS=$((PASS + 1)) ;;
esac
case "$_out" in
*"Accept? [Y/n]"*) echo " FAIL: must not print an unanswerable prompt"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: no dangling Accept? prompt without a tty"; PASS=$((PASS + 1)) ;;
esac
# Passwordless sudo is the one case where unattended escalation is legitimate.
_out=$(run_smart notty nopasswd)
assert_contains "no tty + passwordless sudo: still installs" "$_out" "SUDO_RAN: apt-get install -y cmake"
assert_contains "no tty + passwordless sudo: says why it proceeded" \
"$_out" "passwordless sudo"
# A readable tty must behave exactly as before: prompt, then honour the answer.
_out=$(run_smart tty needspasswd)
assert_contains "tty present: still prompts" "$_out" "Accept? [Y/n]"
assert_contains "tty present: accepts and installs" "$_out" "SUDO_RAN: apt-get install -y cmake"
# Consent given at a real tty, but the elevated apt-get fails anyway (sudoers
# denial, wrong password, apt error). The interactive branch must say what to
# run by hand, like the headless branch does, not die on the bare sudo error.
_out=$(run_smart tty denied)
assert_contains "tty + denied sudo: gives the manual command" \
"$_out" "sudo apt-get update -y && sudo apt-get install -y cmake"
# No sudo at all keeps its own message.
_out=$(run_smart notty absent)
assert_contains "no sudo binary: unchanged message" "$_out" "sudo is not available on this system"
# A /dev/tty that passes `test -r` but cannot be opened counts as no tty.
# Only assert where the platform can actually produce that shape.
_probe=$(mktemp -d -p "$_TMP_ROOT")
if make_unopenable "$_probe/tty" && [ -r "$_probe/tty" ] && ! ( : <"$_probe/tty" ) 2>/dev/null; then
_out=$(run_smart unopenable needspasswd)
assert_contains "unopenable tty: treated as no tty" "$_out" "cannot be done unattended"
case "$_out" in
*"Accept? [Y/n]"*) echo " FAIL: unopenable tty must not print a prompt"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: unopenable tty prints no prompt"; PASS=$((PASS + 1)) ;;
esac
else
echo " SKIP: this platform cannot fake a readable-but-unopenable /dev/tty"
fi
# A tty that opens but yields EOF must decline: a failed read is nobody
# answering, and calling that "yes" escalates through the branch that does
# have a terminal.
_out=$(run_smart eof needspasswd)
assert_contains "eof tty: declines instead of escalating" \
"$_out" "Please install these packages first"
case "$_out" in
*SUDO_RAN*) echo " FAIL: eof tty must not escalate"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: eof tty runs nothing as root"; PASS=$((PASS + 1)) ;;
esac
# A cached timestamp from an earlier, unrelated sudo must not count as
# passwordless: nobody answered this run's prompt and the apt-get rule still
# carries PASSWD. Asserts the -k is present and effective.
_out=$(run_smart notty cached)
assert_contains "cached credentials: says it cannot run unattended" \
"$_out" "cannot be done unattended"
case "$_out" in
*SUDO_RAN*) echo " FAIL: a cached timestamp must not authorise unattended install"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: cached credentials run nothing as root"; PASS=$((PASS + 1)) ;;
esac
# The failure message must not blame a password when apt itself failed: sudo
# passes the command's own exit status through when the command runs.
assert_contains "failure message does not blame a password exclusively" \
"$_out" "or apt-get itself"
# Authorized for apt-get but not NOPASSWD on it. Both `sudo -n true` and
# `sudo -n -l -- apt-get ...` read this as unattended, since list mode answers
# authorization, not authentication. Only running it with -n is truthful.
_out=$(run_smart notty aptneedspasswd)
assert_contains "apt-get needs a password: says it cannot run unattended" \
"$_out" "cannot be done unattended"
case "$_out" in
*SUDO_RAN*) echo " FAIL: apt-get needing a password must not run as root"; FAIL=$((FAIL + 1)) ;;
*) echo " PASS: apt-get needing a password runs nothing as root"; PASS=$((PASS + 1)) ;;
esac
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ]

View file

@ -0,0 +1,391 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Out-of-disk handling in the llama.cpp prebuilt installer: ENOSPC classification
through exception chains, EXIT_NO_SPACE, and the advisory low-disk warning. Offline."""
from __future__ import annotations
import errno
import importlib.util
import shutil
import sys
import urllib.error
from pathlib import Path
import pytest
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
MODULE_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
SPEC = importlib.util.spec_from_file_location("studio_install_llama_prebuilt", MODULE_PATH)
assert SPEC is not None and SPEC.loader is not None
INSTALL_LLAMA_PREBUILT = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = INSTALL_LLAMA_PREBUILT
SPEC.loader.exec_module(INSTALL_LLAMA_PREBUILT)
M = INSTALL_LLAMA_PREBUILT
PrebuiltFallback = M.PrebuiltFallback
AssetChoice = M.AssetChoice
ApprovedReleaseChecksums = M.ApprovedReleaseChecksums
GB = 1024**3
def linux_host() -> "M.HostInfo":
return M.HostInfo(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
)
def choice(name: str, tag: str = "release-2") -> "M.AssetChoice":
return AssetChoice(
repo = "unslothai/llama.cpp",
tag = tag,
name = name,
url = f"https://example.com/{name}",
source_label = "published",
install_kind = "linux-cpu",
)
def checksums(release_tag: str, llama_tag: str) -> "M.ApprovedReleaseChecksums":
return ApprovedReleaseChecksums(
repo = "unslothai/llama.cpp",
release_tag = release_tag,
upstream_tag = llama_tag,
source_commit = None,
artifacts = {},
)
def plan(llama_tag: str, release_tag: str, attempts) -> "M.InstallReleasePlan":
return M.InstallReleasePlan(
requested_tag = "latest",
llama_tag = llama_tag,
release_tag = release_tag,
attempts = attempts,
approved_checksums = checksums(release_tag, llama_tag),
)
def fake_disk_usage(free_bytes: int):
def _usage(path):
return shutil._ntuple_diskusage(100 * GB, 100 * GB - free_bytes, free_bytes)
return _usage
def install_harness(monkeypatch: pytest.MonkeyPatch, plans, *, free_bytes: int) -> list[str]:
"""Wire install_prebuilt down to a fake per-candidate validation. Returns the
list of candidate names the run actually reached."""
monkeypatch.setattr(M, "detect_host", lambda: linux_host())
monkeypatch.setattr(
M,
"resolve_simple_install_release_plans",
lambda llama_tag, host, published_repo, published_release_tag: ("latest", plans),
)
monkeypatch.setattr(
M, "download_validation_model", lambda probe_path, cache_path: probe_path.write_bytes(b"p")
)
monkeypatch.setattr(M.shutil, "disk_usage", fake_disk_usage(free_bytes))
monkeypatch.setattr(M, "existing_install_matches_plan", lambda *args, **kwargs: False)
monkeypatch.setattr(M, "existing_install_matches_choice", lambda *args, **kwargs: False)
monkeypatch.setattr(M, "activate_install_tree", lambda *args, **kwargs: None)
monkeypatch.setattr(M, "ensure_converter_scripts", lambda *args, **kwargs: None)
monkeypatch.setattr(M, "ensure_diffusion_visual_server", lambda *args, **kwargs: None)
monkeypatch.setattr(M, "collect_system_report", lambda *args, **kwargs: "report")
reached: list[str] = []
monkeypatch.setattr(
M, "validate_prebuilt_choice", lambda attempt, *a, **k: reached.append(attempt.name)
)
return reached
# ── the low-disk check is advisory, never fatal ──
def test_low_disk_warning_reports_the_starved_volume(tmp_path, monkeypatch):
monkeypatch.setattr(M.shutil, "disk_usage", fake_disk_usage(1 * GB))
reason = M._low_disk_warning(tmp_path / "llama.cpp")
assert reason is not None and "low disk space for llama.cpp" in reason
def test_low_disk_warning_silent_when_roomy(tmp_path, monkeypatch):
monkeypatch.setattr(M.shutil, "disk_usage", fake_disk_usage(50 * GB))
assert M._low_disk_warning(tmp_path / "llama.cpp") is None
def test_low_disk_warning_ignores_unstatable_paths(tmp_path, monkeypatch):
def _boom(path):
raise OSError(errno.EACCES, "permission denied")
monkeypatch.setattr(M.shutil, "disk_usage", _boom)
assert M._low_disk_warning(tmp_path / "llama.cpp") is None
def test_low_disk_does_not_block_an_install_that_fits(tmp_path, monkeypatch, capsys):
"""A 15 MB CPU bundle installs fine on a host with 3 GB free; the fixed
threshold must warn rather than reject it (and skip the source fallback)."""
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
only = plan("b10079", "release-2", [choice("app-b10079-linux-x64-cpu.tar.gz")])
reached = install_harness(monkeypatch, [only], free_bytes = 3 * GB)
M.install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
assert reached == ["app-b10079-linux-x64-cpu.tar.gz"]
captured = capsys.readouterr()
assert "low disk space for llama.cpp" in captured.out + captured.err
# ── ENOSPC classification ──
def test_classifies_direct_and_chained_enospc():
assert M._environment_fatal_reason(OSError(errno.ENOSPC, "No space left on device"))
for wrap in ("cause", "context"):
try:
try:
raise OSError(errno.ENOSPC, "No space left on device")
except OSError as inner:
if wrap == "cause":
raise PrebuiltFallback("download failed") from inner
raise PrebuiltFallback("download failed")
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer), wrap
def test_ignores_unrelated_errors_and_cycles():
assert M._environment_fatal_reason(OSError(errno.EACCES, "denied")) is None
first, second = PrebuiltFallback("a"), PrebuiltFallback("b")
first.__cause__, second.__cause__ = second, first
assert M._environment_fatal_reason(first) is None
def test_suppressed_context_is_not_treated_as_disk_full():
"""`raise ... from None` means the earlier ENOSPC is unrelated."""
try:
try:
raise OSError(errno.ENOSPC, "No space left on device")
except OSError:
raise PrebuiltFallback("checksum mismatch") from None
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer) is None
def test_windows_disk_full_winerrors_are_classified():
"""CPython maps ERROR_DISK_FULL (112) to ENOSPC but has no case for
ERROR_HANDLE_DISK_FULL (39), which arrives as EINVAL."""
for winerror, code in ((112, errno.ENOSPC), (39, errno.EINVAL)):
exc = OSError(code, "The disk is full")
exc.winerror = winerror
assert M._environment_fatal_reason(exc), winerror
other = OSError(errno.EACCES, "sharing violation")
other.winerror = 32
assert M._environment_fatal_reason(other) is None
def test_http_errors_in_the_chain_do_not_crash_the_classifier():
"""HTTPError is an OSError that proxies unknown attributes to a wrapped file
and raises KeyError, not AttributeError, on 3.9."""
err = urllib.error.HTTPError("https://example.com/a", 404, "Not Found", {}, None)
assert M._environment_fatal_reason(err) is None
try:
try:
raise err
except urllib.error.HTTPError as inner:
raise PrebuiltFallback("mirror failed") from inner
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer) is None
@pytest.mark.skipif(not hasattr(errno, "EDQUOT"), reason = "EDQUOT is POSIX only")
def test_quota_exhaustion_counts_as_out_of_space():
"""A quota'd home has free blocks this user cannot have, so the larger source
build is just as doomed. Reported as a quota so df does not mislead."""
assert M._environment_fatal_reason(OSError(errno.EDQUOT, "Disk quota exceeded")) == (
"disk quota exceeded"
)
try:
try:
raise OSError(errno.EDQUOT, "Disk quota exceeded")
except OSError as inner:
raise PrebuiltFallback("bundle download failed") from inner
except PrebuiltFallback as outer:
assert M._environment_fatal_reason(outer) == "disk quota exceeded"
def test_a_bare_oserror_never_matches():
"""errno is None on a bare OSError, so it must not collide with a code."""
assert M._environment_fatal_reason(OSError()) is None
assert M._environment_fatal_reason(shutil.Error("copy failed")) is None
def test_flattened_markers_are_not_matched_as_prefixes():
"""Bare "WinError 112" would also match WinError 1120; the brackets pin it."""
assert (
M._environment_fatal_reason(
shutil.Error("[('a', 'b', '[WinError 1120] a serial write completed')]")
)
is None
)
assert (
M._environment_fatal_reason(
shutil.Error(f"[('a', 'b', '[Errno {errno.ENOSPC}0] not a real code')]")
)
is None
)
def test_windows_flattened_disk_full_text_is_classified():
"""copytree stringifies the per-file OSError, and on Windows str(OSError)
prints [WinError 112] and never [Errno 28] (confirmed on a real NTFS volume)."""
flattened = (
"[('D:\\\\a\\\\src\\\\big.bin', 'T:\\\\dst\\\\big.bin', "
"'[WinError 112] There is not enough space on the disk')]"
)
assert M._environment_fatal_reason(shutil.Error(flattened))
assert M._environment_fatal_reason(
shutil.Error("[('a', 'b', '[WinError 39] The disk is full')]")
)
assert (
M._environment_fatal_reason(shutil.Error("[('a', 'b', '[WinError 32] sharing violation')]"))
is None
)
def test_validate_install_mode_exits_no_space(tmp_path, monkeypatch):
"""setup.sh reacts to a failed staged validation by deleting the finished GPU
build and starting a CPU rebuild, which needs more of the space that ran out."""
def boom(*args, **kwargs):
try:
raise OSError(errno.ENOSPC, "No space left on device")
except OSError as inner:
raise PrebuiltFallback("validation model unavailable") from inner
monkeypatch.setattr(M, "validate_existing_install", boom)
monkeypatch.setattr(
sys, "argv", ["install_llama_prebuilt.py", "--validate-install", str(tmp_path)]
)
with pytest.raises(SystemExit) as caught:
M.main()
assert caught.value.code == M.EXIT_NO_SPACE
def test_validate_install_mode_still_falls_back_on_ordinary_failure(tmp_path, monkeypatch):
monkeypatch.setattr(
M,
"validate_existing_install",
lambda *a, **k: (_ for _ in ()).throw(PrebuiltFallback("llama-server crashed")),
)
monkeypatch.setattr(
sys, "argv", ["install_llama_prebuilt.py", "--validate-install", str(tmp_path)]
)
with pytest.raises(SystemExit) as caught:
M.main()
assert caught.value.code == M.EXIT_FALLBACK
def test_classifies_enospc_hidden_in_a_shutil_error(tmp_path):
"""copytree stringifies the per-file OSError, so errno and the chain are gone."""
src = tmp_path / "src" / "sub"
src.mkdir(parents = True)
(src / "f").write_text("x", encoding = "utf-8")
def boom(*args, **kwargs):
raise OSError(errno.ENOSPC, "No space left on device")
with pytest.raises(shutil.Error) as caught:
shutil.copytree(tmp_path / "src", tmp_path / "dst", copy_function = boom)
assert caught.value.errno is None
assert M._environment_fatal_reason(caught.value)
def test_source_tree_enospc_is_not_masked_by_a_later_mirror_error(tmp_path, monkeypatch):
"""A full disk fails every mirror, so the first ENOSPC must win over a 404."""
calls: list[str] = []
def fake_download(
url,
path,
*,
expected_sha256 = None,
label = None,
):
calls.append(url)
if len(calls) == 1:
raise OSError(errno.ENOSPC, "No space left on device")
raise urllib.error.HTTPError(url, 404, "Not Found", {}, None)
monkeypatch.setattr(M, "download_file_verified", fake_download)
with pytest.raises(PrebuiltFallback) as caught:
M.hydrate_source_tree(
"deadbeef",
tmp_path / "install",
tmp_path,
source_repo = "unslothai/llama.cpp",
expected_sha256 = None,
exact_source = True,
asset_url = "https://example.com/llama.cpp-source.tar.gz",
)
assert len(calls) == 1, f"stopped after the first ENOSPC, tried: {calls}"
assert M._environment_fatal_reason(caught.value)
# ── exit codes ──
def test_enospc_exits_no_space_without_trying_older_releases(tmp_path, monkeypatch):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
newer = plan("b9002", "release-2", [choice("app-b9002-linux-x64-cpu.tar.gz")])
older = plan("b9001", "release-1", [choice("app-b9001-linux-x64-cpu.tar.gz", "release-1")])
reached = install_harness(monkeypatch, [newer, older], free_bytes = 50 * GB)
def enospc(attempt, *args, **kwargs):
reached.append(attempt.name)
raise OSError(errno.ENOSPC, "No space left on device")
monkeypatch.setattr(M, "validate_prebuilt_choice", enospc)
with pytest.raises(SystemExit) as caught:
M.install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
assert caught.value.code == M.EXIT_NO_SPACE
assert reached == ["app-b9002-linux-x64-cpu.tar.gz"]
def test_ordinary_failure_still_exits_fallback(tmp_path, monkeypatch):
install_dir = tmp_path / "llama.cpp"
install_dir.mkdir()
only = plan("b9002", "release-2", [choice("app-b9002-linux-x64-cpu.tar.gz")])
install_harness(monkeypatch, [only], free_bytes = 50 * GB)
monkeypatch.setattr(
M,
"validate_prebuilt_choice",
lambda *a, **k: (_ for _ in ()).throw(PrebuiltFallback("checksum mismatch")),
)
with pytest.raises(SystemExit) as caught:
M.install_prebuilt(install_dir, "latest", "unslothai/llama.cpp", "")
assert caught.value.code == M.EXIT_FALLBACK

View file

@ -12,6 +12,7 @@ at import) resolves from a clean process.
from __future__ import annotations
import json
import os
import subprocess
import sys
from pathlib import Path
@ -43,6 +44,15 @@ _ARCHES = {
_CHILD = """
import json, sys
sys.path.insert(0, {tests!r})
# Import bitsandbytes under the real torch first. unsloth_zoo pulls it in, and it
# picks a compute backend at import: once the spoof reports an AMD GPU, it loads
# its ROCm/CUDA ops, which a CPU-only torch cannot satisfy (no libhipblas, no
# torch._C._cuda_getCurrentRawStream) and the child dies before printing RESULT.
# Nothing here tests bitsandbytes, so let it see the honest hardware.
try:
import bitsandbytes # noqa: F401
except Exception:
pass
import _zoo_rocm_spoof as spoof
arches = {arches!r}
spoof.apply(arches[0])
@ -60,7 +70,11 @@ print("RESULT " + json.dumps({{"device_type": device_type, "targets": targets}})
@pytest.fixture(scope = "module")
def routed():
code = _CHILD.format(tests = str(_TESTS_DIR), arches = list(_ARCHES))
proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True)
# get_device_type() returns "mlx" before it ever looks at torch on Darwin arm64
# with mlx installed, so the spoof would be ignored. Force the GPU path to keep
# the assertion live there instead of skipping it.
env = {**os.environ, "UNSLOTH_FORCE_GPU_PATH": "1"}
proc = subprocess.run([sys.executable, "-c", code], capture_output = True, text = True, env = env)
line = next((l for l in proc.stdout.splitlines() if l.startswith("RESULT ")), None)
assert line, f"child produced no result.\nstdout:\n{proc.stdout}\nstderr:\n{proc.stderr}"
return json.loads(line[len("RESULT ") :])

View file

@ -16,6 +16,7 @@ import os
import sys
from pathlib import Path
from playwright.sync_api import TimeoutError as PWTimeout
from playwright.sync_api import sync_playwright
sys.path.insert(0, str(Path(__file__).resolve().parent))
@ -46,6 +47,18 @@ def near(
return a is not None and b is not None and abs(a - b) <= tol
_VP = 'document.querySelector("[data-radix-select-viewport]")'
SCROLL_TOP_JS = f"() => {_VP}.scrollTop"
SCROLLABLE_JS = f"() => {{ const vp = {_VP}; return !!vp && vp.scrollHeight > vp.clientHeight; }}"
VIEWPORT_STATE_JS = f"""
() => {{
const vp = {_VP};
return vp
? {{ scrollHeight: vp.scrollHeight, clientHeight: vp.clientHeight, top: vp.scrollTop }}
: null;
}}
"""
MEASURE_JS = """
() => {
const fs = (el) => (el ? parseFloat(getComputedStyle(el).fontSize) : null);
@ -83,15 +96,22 @@ def set_input(page, label, value):
def open_appearance(page):
page.keyboard.press("Control+,")
page.wait_for_timeout(700)
if page.get_by_role("dialog").count() == 0:
page.keyboard.press("Meta+,")
page.wait_for_timeout(700)
if page.get_by_role("dialog").count() == 0:
fail("settings dialog did not open")
page.get_by_role("dialog").get_by_role("button").filter(has_text = "Appearance").first.click()
page.wait_for_timeout(600)
# The shortcut can fire before the app has wired its key handler, so press
# each chord once behind a fixed sleep and a slow boot loses the dialog.
# Alternate them on a bounded retry, waiting on the dialog itself.
dialog = page.get_by_role("dialog")
for attempt in range(10):
page.keyboard.press("Meta+," if attempt % 2 else "Control+,")
try:
dialog.first.wait_for(state = "visible", timeout = 2_000)
break
except PWTimeout:
continue
if dialog.count() == 0:
fail("settings dialog did not open after 10 attempts")
dialog.get_by_role("button").filter(has_text = "Appearance").first.click()
# Wait for the control the caller is about to drive, not a fixed interval.
page.locator("input[aria-label='UI font size']").wait_for(state = "visible", timeout = 15_000)
def main():
@ -155,39 +175,46 @@ def main():
page.wait_for_timeout(400)
step("overflowing select scrolls its Radix viewport")
page.get_by_role("dialog").get_by_role("button").filter(has_text = "Voice").first.click()
page.wait_for_timeout(600)
voice = page.get_by_role("dialog").get_by_role("button").filter(has_text = "Voice").first
voice.click()
page.set_viewport_size({"width": 1440, "height": 480})
page.locator("[aria-label='Dictation language']").click()
page.wait_for_timeout(700)
state = page.evaluate(
"""
() => {
const vp = document.querySelector("[data-radix-select-viewport]");
return vp
? { scrollable: vp.scrollHeight > vp.clientHeight, top: vp.scrollTop }
: null;
}
"""
)
if not state or not state["scrollable"]:
fail(f"select viewport not scrollable: {state}")
for _ in range(6):
trigger = page.locator("[aria-label='Dictation language']")
trigger.wait_for(state = "visible")
trigger.click()
viewport = page.locator("[data-radix-select-viewport]")
viewport.wait_for(state = "visible")
# Wait for the overflow itself rather than a fixed sleep: the list is
# populated asynchronously, so measuring too early reads it as short.
try:
page.wait_for_function(SCROLLABLE_JS, timeout = 10_000)
except PWTimeout:
fail(f"select viewport not scrollable: {page.evaluate(VIEWPORT_STATE_JS)}")
# Radix moves focus into the listbox after the content opens, so a fixed
# burst of presses can land on the trigger and scroll nothing. Press until
# it moves instead; a real regression still fails, just after more tries.
kb_top = 0
for _ in range(40):
page.keyboard.press("ArrowDown")
page.wait_for_timeout(100)
kb_top = page.evaluate(
"() => document.querySelector('[data-radix-select-viewport]').scrollTop"
)
kb_top = page.evaluate(SCROLL_TOP_JS)
if kb_top > 0:
break
page.wait_for_timeout(50)
if not kb_top > 0:
fail(f"keyboard did not scroll the select viewport: {kb_top}")
vp_box = page.locator("[data-radix-select-viewport]").bounding_box()
fail(f"keyboard did not scroll the select viewport after 40 presses: {kb_top}")
vp_box = viewport.bounding_box()
page.mouse.move(vp_box["x"] + vp_box["width"] / 2, vp_box["y"] + 40)
page.mouse.wheel(0, -400)
page.wait_for_timeout(300)
wheel_top = page.evaluate(
"() => document.querySelector('[data-radix-select-viewport]').scrollTop"
)
if not wheel_top < kb_top:
try:
page.wait_for_function(
"top => document.querySelector('[data-radix-select-viewport]').scrollTop < top",
arg = kb_top,
timeout = 10_000,
)
except PWTimeout:
wheel_top = page.evaluate(SCROLL_TOP_JS)
fail(f"wheel did not scroll the select viewport: {kb_top} -> {wheel_top}")
page.keyboard.press("Escape")
page.set_viewport_size({"width": 1440, "height": 900})

View file

@ -85,6 +85,21 @@ def test_reasoning_keeps_streaming_height_cap_through_automatic_collapse():
assert "streaming={isReasoningStreaming || retainStreamingHeight}" in src
def test_reasoning_clears_manual_open_on_a_new_stream():
"""A hand-opened block must not stay pinned open when the stream restarts.
isOpen is `(streaming && !dismissed) || manualOpen` and manualOpen is only
settable while idle, so the new-stream reset has to clear it too.
"""
src = REASONING_TSX.read_text()
marker = "setDismissedWhileStreaming(false)"
start = src.find(marker)
assert start != -1, "new-stream reset effect is missing"
effect = src[src.rfind("useEffect(() => {", 0, start) : src.find("});", start)]
assert "setManualOpen(false)" in effect
def test_response_details_metadata_is_persisted_without_backend_schema_change():
src = ADAPTER_TS.read_text()
assert "interface ResponseDetailsMetadata" in src

View file

@ -0,0 +1,244 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Contracts and opt-in runtime coverage for the PDF grounded QA recipe."""
from __future__ import annotations
import copy
import importlib.util
import json
import os
import re
import sys
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import pytest
REPO = Path(__file__).resolve().parents[2]
RECIPE_PATH = (
REPO / "studio/frontend/src/features/data-recipes/learning-recipes/pdf-grounded-qa.json"
)
TRAINING_ACTIONS_PATH = REPO / "studio/frontend/src/features/training/hooks/use-training-actions.ts"
SEED_BUILDER_PATH = (
REPO / "studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts"
)
RECIPE_IMPORTER_PATH = REPO / "studio/frontend/src/features/recipe-studio/utils/import/importer.ts"
SEED_PARSER_PATH = (
REPO / "studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts"
)
FORMAT_DETECTION_PATH = REPO / "studio/backend/utils/datasets/format_detection.py"
def _load_payload() -> dict:
return json.loads(RECIPE_PATH.read_text(encoding = "utf-8"))
def _render_expression(template: str, row: dict) -> str:
def replace(match: re.Match[str]) -> str:
value = row
for part in match.group(1).strip().split("."):
value = value[part]
return str(value)
return re.sub(r"\{\{\s*([^}]+?)\s*\}\}", replace, template)
def test_pdf_qa_recipe_projects_and_cleans_training_columns():
recipe = _load_payload()["recipe"]
columns = {column["name"]: column for column in recipe["columns"]}
assert list(columns) == ["llm_structured_1", "instruction", "output", "input"]
assert columns["llm_structured_1"]["drop"] is True
assert columns["instruction"]["expr"] == "{{ llm_structured_1.question }}"
assert columns["output"]["expr"] == "{{ llm_structured_1.answer }}"
assert "llm_structured_1.evidence_quote" in columns["input"]["expr"]
assert "chunk_text" in columns["input"]["expr"]
assert recipe["processors"] == [
{
"processor_type": "drop_columns",
"name": "drop_seed_columns",
"column_names": ["chunk_text", "source_file"],
}
]
def test_pdf_qa_recipe_sample_row_is_qlora_ready():
recipe = _load_payload()["recipe"]
row = {
"chunk_text": "Paris is the capital of France.",
"source_file": "facts.pdf",
"llm_structured_1": {
"question": "What is the capital of France?",
"answer": "Paris.",
"evidence_quote": "Paris is the capital of France.",
},
}
for column in recipe["columns"]:
if column["column_type"] == "expression":
row[column["name"]] = _render_expression(column["expr"], row)
for column in recipe["columns"]:
if column.get("drop"):
row.pop(column["name"], None)
for processor in recipe["processors"]:
for name in processor["column_names"]:
row.pop(name, None)
assert row == {
"instruction": "What is the capital of France?",
"output": "Paris.",
"input": (
"Evidence quote: Paris is the capital of France.\n\n"
"Source context: Paris is the capital of France."
),
}
def test_pdf_qa_canvas_edges_cover_expression_dependencies():
payload = _load_payload()
recipe = payload["recipe"]
node_ids = {node["id"] for node in payload["ui"]["nodes"]}
edges = {(edge["from"], edge["to"]) for edge in payload["ui"]["edges"]}
assert all(source in node_ids and target in node_ids for source, target in edges)
assert ("seed", "llm_structured_1") in edges
assert ("llm_structured_1", "instruction") in edges
assert ("llm_structured_1", "output") in edges
assert ("llm_structured_1", "input") in edges
assert ("seed", "input") in edges
column_names = {column["name"] for column in recipe["columns"]}
assert {"instruction", "output"} <= column_names
def test_pdf_qa_fields_match_studio_alpaca_mapping():
source = TRAINING_ACTIONS_PATH.read_text(encoding = "utf-8")
assert 'alpaca: { user: "instruction", system: "input", assistant: "output" }' in source
assert 'if (fmt === "alpaca") return roles.has("instruction") && roles.has("output");' in source
def test_pdf_qa_fields_are_detected_as_alpaca():
spec = importlib.util.spec_from_file_location("_pdf_qa_format_detection", FORMAT_DETECTION_PATH)
assert spec and spec.loader
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
detected = module.detect_dataset_format(
[{"instruction": "What is the capital?", "input": "source", "output": "Paris."}]
)
assert detected["format"] == "alpaca"
assert detected["needs_standardization"] is False
def test_unstructured_seed_drop_toggle_round_trip_contract():
builder = SEED_BUILDER_PATH.read_text(encoding = "utf-8")
importer = RECIPE_IMPORTER_PATH.read_text(encoding = "utf-8")
parser = SEED_PARSER_PATH.read_text(encoding = "utf-8")
assert 'if (seedSourceType === "unstructured")' in builder
assert "if (!config.drop)" in builder
assert "selectedDropColumns.length > 0" in builder
assert ': ["chunk_text", "source_file"];' in builder
assert "payloadSeedSourceIsUnstructured && payloadSeedDropColumns.length > 0" in importer
assert "payloadSeedSourceIsUnstructured" in importer
assert '? ["chunk_text", "source_file"]' in importer
assert "drop?: boolean;" in parser
assert "...(options?.drop !== undefined ? { drop: options.drop } : {})" in parser
class _MockOpenAIHandler(BaseHTTPRequestHandler):
requests: list[dict] = []
def log_message(self, format: str, *args) -> None:
return
def do_POST(self) -> None:
raw = self.rfile.read(int(self.headers.get("Content-Length", "0")))
self.requests.append(json.loads(raw or b"{}"))
structured = {
"question": "What is the capital of France?",
"answer": "Paris.",
"evidence_quote": "Paris is the capital of France.",
}
body = json.dumps(
{
"id": "chatcmpl-pdf-qa-test",
"object": "chat.completion",
"created": 0,
"model": "mock-model",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": f"```json\n{json.dumps(structured)}\n```",
},
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30,
},
}
).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def test_pdf_qa_recipe_runs_with_pinned_data_designer(tmp_path, monkeypatch):
if os.environ.get("UNSLOTH_PDF_QA_MANAGED_INTEGRATION") != "1":
pytest.skip("set UNSLOTH_PDF_QA_MANAGED_INTEGRATION=1 to run this integration")
backend = REPO / "studio/backend"
sys.path.insert(0, str(backend))
pytest.importorskip("data_designer")
pytest.importorskip("data_designer_unstructured_seed")
from core.data_recipe import service
source_path = tmp_path / "facts.txt"
source_path.write_text("Paris is the capital of France.", encoding = "utf-8")
monkeypatch.setattr(service, "recipe_datasets_root", lambda: tmp_path / "artifacts")
server = ThreadingHTTPServer(("127.0.0.1", 0), _MockOpenAIHandler)
thread = threading.Thread(target = server.serve_forever, daemon = True)
thread.start()
try:
recipe = copy.deepcopy(_load_payload()["recipe"])
recipe["seed_config"]["source"] = {
"seed_type": "unstructured",
"paths": [str(source_path)],
"chunk_size": 1200,
"chunk_overlap": 200,
}
recipe["model_providers"][0].update(
{
"endpoint": f"http://127.0.0.1:{server.server_port}/v1",
"api_key": "test-only",
}
)
recipe["model_configs"][0].update({"model": "mock-model", "skip_health_check": True})
dataset, _, _ = service.preview_recipe(recipe, 1)
finally:
server.shutdown()
server.server_close()
thread.join(timeout = 5)
assert dataset == [
{
"instruction": "What is the capital of France?",
"output": "Paris.",
"input": (
"Evidence quote: Paris is the capital of France.\n\n"
"Source context: Paris is the capital of France."
),
}
]
assert _MockOpenAIHandler.requests

View file

@ -0,0 +1,670 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""The compressed (FP8/NVFP4) export must free GPU weights before its llm-compressor
subprocess loads a second copy from disk, including for accelerate-dispatched multi-GPU
shards, which the old single-device-only ``.to("cpu")`` skipped and left resident.
Pulls the release/restore helpers out of unsloth/save.py via AST (importing the module
needs torch/transformers) and exercises them with fakes.
"""
from __future__ import annotations
import ast
import gc
import sys
import types
from pathlib import Path
import pytest
_SAVE_PY = Path(__file__).resolve().parent.parent / "unsloth" / "save.py"
_WANTED = {
"_accelerate_dispatch_root",
"_snapshot_dispatch_state",
"_drop_accelerator_tied_param_cache",
"_accelerate_move_guards",
"_split_tensor_path",
"_lookup_tensor",
"_share_tensor",
"_restore_dispatch_state",
"_offload_model_for_quantize_subprocess",
"_restore_model_after_quantize_subprocess",
}
_WANTED_ASSIGNS = {
"_DISPATCH_SNAPSHOT_ATTR",
"_ACCELERATE_MOVE_GUARDS",
} # module constants the helpers close over
class _FakeLogger:
def __init__(self):
self.warnings = []
def warning_once(self, msg):
self.warnings.append(msg)
def _load_helpers(fake_torch, fake_logger):
tree = ast.parse(_SAVE_PY.read_text(encoding = "utf-8"))
keep = [
node
for node in tree.body
if (isinstance(node, ast.FunctionDef) and node.name in _WANTED)
or (
isinstance(node, ast.Assign)
and any(isinstance(t, ast.Name) and t.id in _WANTED_ASSIGNS for t in node.targets)
)
]
n_fns = sum(1 for node in keep if isinstance(node, ast.FunctionDef))
assert n_fns == len(_WANTED), "release helpers missing from save.py"
namespace = {"torch": fake_torch, "logger": fake_logger}
exec( # noqa: S102 - loading trusted repo source
compile(ast.Module(body = keep, type_ignores = []), str(_SAVE_PY), "exec"),
namespace,
)
return namespace
def _fake_torch(cuda_available = True):
t = types.ModuleType("torch")
t.cuda = types.SimpleNamespace(is_available = lambda: cuda_available)
return t
class _FakeModel:
def __init__(
self,
device_map = None,
devices = ("cuda:0",),
quantized = False,
):
if device_map is not None:
self.hf_device_map = device_map
self._devices = [types.SimpleNamespace(device = d) for d in devices]
self.moved_to = []
self.is_loaded_in_4bit = quantized
def parameters(self):
return iter(self._devices)
def to(self, target):
self.moved_to.append(str(target))
return self
@pytest.fixture
def _fake_accelerate(monkeypatch):
calls = {"removed": [], "dispatched": [], "dispatch_kwargs": [], "hooks_added": []}
accel = types.ModuleType("accelerate")
def _dispatch(model, device_map, **kwargs):
calls["dispatched"].append((model, dict(device_map)))
calls["dispatch_kwargs"].append(kwargs)
accel.dispatch_model = _dispatch
hooks = types.ModuleType("accelerate.hooks")
hooks.remove_hook_from_submodules = lambda model: calls["removed"].append(model)
hooks.add_hook_to_module = lambda module, hook: calls["hooks_added"].append((module, hook))
accel.hooks = hooks
monkeypatch.setitem(sys.modules, "accelerate", accel)
monkeypatch.setitem(sys.modules, "accelerate.hooks", hooks)
return calls
def test_dispatched_multi_gpu_model_is_released_and_redispatched(_fake_accelerate):
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 0, "model.layers.1": 1}
model = _FakeModel(device_map = device_map, devices = ("cuda:0", "cuda:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model] # hooks removed before the move
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_dispatched_move_failure_redispatches_and_returns_none(_fake_accelerate):
# If .to("cpu") raises after the hooks came off, the model must be re-dispatched,
# not left hookless and half-moved.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.1": 1}
class _MoveFails(_FakeModel):
def to(self, target):
raise RuntimeError("host RAM cannot hold the sharded model")
model = _MoveFails(device_map = device_map, devices = ("cuda:0", "cuda:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token is None # offload aborted
assert _fake_accelerate["removed"] == [model] # hooks were removed...
assert _fake_accelerate["dispatched"] == [(model, device_map)] # ...then restored
def test_single_device_move_failure_restores_and_returns_none():
ns = _load_helpers(_fake_torch(), _FakeLogger())
class _MoveFails(_FakeModel):
def __init__(self):
super().__init__(devices = ("cuda:0",))
def to(self, target):
self.moved_to.append(str(target))
if target == "cpu":
raise RuntimeError("move failed")
return self
model = _MoveFails()
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token is None
# attempted the cpu move, then restored back to the original device
assert model.moved_to == ["cpu", "cuda:0"]
def test_cpu_spilled_map_still_releases_its_gpu_shards(_fake_accelerate):
# One module spilled to CPU, but the rest is the GPU memory the reload needs, and
# the spilled weights are already in host RAM, so the move is safe.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1, "model.layers.9": "cpu"}
model = _FakeModel(device_map = device_map)
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model]
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_disk_offloaded_map_is_left_alone(_fake_accelerate):
# disk/meta entries are not on the model, so moving would materialize the whole
# checkpoint into RAM.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": 0, "model.layers.9": "disk"})
assert ns["_offload_model_for_quantize_subprocess"](model) is None
assert model.moved_to == []
assert _fake_accelerate["removed"] == []
def test_all_cpu_map_is_left_alone(_fake_accelerate):
# Nothing on an accelerator: no GPU memory to reclaim, so do not churn the hooks.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": "cpu", "model.layers.0": "cpu"})
assert ns["_offload_model_for_quantize_subprocess"](model) is None
assert model.moved_to == []
assert _fake_accelerate["removed"] == []
def test_single_device_model_keeps_plain_move():
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(devices = ("cuda:0",))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert model.moved_to == ["cpu"]
assert token is not None and token[0] == "device"
ns["_restore_model_after_quantize_subprocess"](model, token)
assert model.moved_to[-1] == "cuda:0"
def test_quantized_model_is_released_when_the_stack_allows_it():
# Studio exports load 4-bit by DEFAULT, so skipping quantized models left a shard
# on every GPU. Release them too where the move is accepted.
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(devices = ("cuda:0",), quantized = True)
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token == ("device", "cuda:0")
assert model.moved_to == ["cpu"]
def test_quantized_model_that_refuses_to_move_is_left_usable():
# transformers rejects .to() for some bitsandbytes builds and raises before
# anything moves, so the old behaviour must hold: no token, nothing escaping.
ns = _load_helpers(_fake_torch(), _FakeLogger())
class _Refuses(_FakeModel):
def to(self, target):
raise ValueError("`.to` is not supported for 4-bit bitsandbytes models")
model = _Refuses(devices = ("cuda:0",), quantized = True)
assert ns["_offload_model_for_quantize_subprocess"](model) is None
def test_no_cuda_is_noop_and_restore_none_is_noop():
ns = _load_helpers(_fake_torch(cuda_available = False), _FakeLogger())
model = _FakeModel()
assert ns["_offload_model_for_quantize_subprocess"](model) is None
ns["_restore_model_after_quantize_subprocess"](model, None) # must not raise
assert model.moved_to == []
def test_restore_failure_warns_instead_of_raising(_fake_accelerate):
fake_logger = _FakeLogger()
ns = _load_helpers(_fake_torch(), fake_logger)
class _ExplodingModel(_FakeModel):
def to(self, target):
raise RuntimeError("device gone")
model = _ExplodingModel(devices = ("cuda:0",))
ns["_restore_model_after_quantize_subprocess"](model, ("device", "cuda:0"))
assert fake_logger.warnings # warned, did not raise
def test_lora_merge_budgets_per_device():
# A merged tensor W lives on the GPU of its source layer, so budget against W's
# own device, not GPU0, else a sharded model OOMs GPU1+ (#7053).
src = _SAVE_PY.read_text(encoding = "utf-8")
tree = ast.parse(src)
fn = next(
(
n
for n in ast.walk(tree)
if isinstance(n, ast.FunctionDef) and n.name == "unsloth_save_model"
),
None,
)
assert fn is not None, "unsloth_save_model not found"
body = ast.get_source_segment(src, fn)
# Budget keyed on W's device, not a hardcoded device 0 / unqualified alloc.
assert "torch.cuda.memory_allocated(W.device)" in body
assert "_device_vram_budget(W.device)" in body
assert "get_device_properties(0).total_memory * maximum_memory_usage" not in body
# ── the torchao ("portable" FP8/INT8) export shares the same release ──
def _fake_torch_xpu():
t = types.ModuleType("torch")
t.cuda = types.SimpleNamespace(is_available = lambda: False)
t.xpu = types.SimpleNamespace(is_available = lambda: True)
return t
def test_dispatched_xpu_model_is_released(_fake_accelerate):
# torchao runs on Intel GPUs too, so an XPU-dispatched shard must release exactly
# like a CUDA one.
ns = _load_helpers(_fake_torch_xpu(), _FakeLogger())
device_map = {"model.embed": "xpu:0", "model.layers.0": "xpu:1"}
model = _FakeModel(device_map = device_map, devices = ("xpu:0", "xpu:1"))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert _fake_accelerate["removed"] == [model]
assert model.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
ns["_restore_model_after_quantize_subprocess"](model, token)
assert _fake_accelerate["dispatched"] == [(model, device_map)]
def test_single_device_xpu_model_is_released():
ns = _load_helpers(_fake_torch_xpu(), _FakeLogger())
model = _FakeModel(devices = ("xpu:0",))
token = ns["_offload_model_for_quantize_subprocess"](model)
assert token == ("device", "xpu:0")
assert model.moved_to == ["cpu"]
def test_torchao_export_uses_the_shared_release():
"""The torchao path must not re-inline a single-device-only ``.to("cpu")``.
A plain move is invalid on a dispatched model, so single-device-only handling left
a multi-GPU shard resident while ``device_map="auto"`` loaded a second copy.
"""
src = _SAVE_PY.read_text(encoding = "utf-8")
torchao = src.split("def _unsloth_save_torchao(", 1)[1].split("\ndef ", 1)[0]
assert "_offload_model_for_quantize_subprocess(model)" in torchao
assert "_restore_model_after_quantize_subprocess(model" in torchao
# No hand-rolled single-device gate left behind.
assert "len(_devs) == 1" not in torchao
# ── regressions for the multi-GPU dispatch branch ──
class _Child:
"""Minimal stand-in for an nn.Module leaf, enough for the dispatch walk."""
def __init__(
self,
name = "inner",
device_map = None,
):
self._modules = {}
self.__dict__["_name"] = name
if device_map is not None:
self.hf_device_map = device_map
def named_modules(self):
yield "", self
for key, child in self._modules.items():
for sub_name, sub in child.named_modules():
yield (f"{key}.{sub_name}" if sub_name else key), sub
def get_submodule(self, target):
node = self
for part in target.split("."):
node = node._modules[part]
return node
def named_parameters(self, remove_duplicate = True):
return iter(())
def named_buffers(self, remove_duplicate = True):
return iter(())
class _PeftLikeWrapper(_Child):
"""Proxies unknown attributes to the wrapped model, like ``PeftModelForCausalLM``:
``hasattr(wrapper, "_hf_hook")`` is True while ``delattr`` fails, which is what made
the offload a silent no-op."""
def __init__(self, inner):
super().__init__(name = "wrapper")
self._modules["base_model"] = inner
self.moved_to = []
def __getattr__(self, item):
return getattr(self._modules["base_model"], item)
def to(self, target):
self.moved_to.append(str(target))
return self
def parameters(self):
return iter(self._modules["base_model"]._devices)
def test_dispatch_root_is_the_inner_model_for_a_peft_style_wrapper(_fake_accelerate):
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1}
inner = _Child(device_map = device_map)
inner._devices = [types.SimpleNamespace(device = "cuda:0")]
wrapper = _PeftLikeWrapper(inner)
assert ns["_accelerate_dispatch_root"](wrapper) is inner
token = ns["_offload_model_for_quantize_subprocess"](wrapper)
# hooks must come off the INNER module, not the proxying wrapper
assert _fake_accelerate["removed"] == [inner]
assert wrapper.moved_to == ["cpu"]
assert token == ("dispatch", device_map)
def test_dispatch_root_falls_back_to_the_model_it_was_given():
ns = _load_helpers(_fake_torch(), _FakeLogger())
model = _FakeModel(device_map = {"model.embed": 0})
assert ns["_accelerate_dispatch_root"](model) is model
def test_offload_failure_is_logged_not_swallowed():
# A bare `return None` is indistinguishable from "nothing to move".
fake_logger = _FakeLogger()
ns = _load_helpers(_fake_torch(), fake_logger)
class _Explodes(_FakeModel):
@property
def hf_device_map(self):
raise RuntimeError("boom")
assert ns["_offload_model_for_quantize_subprocess"](_Explodes()) is None
assert any("boom" in w for w in fake_logger.warnings)
def test_restore_without_a_snapshot_forwards_skip_keys(_fake_accelerate):
# dispatch_model() defaults skip_keys to None, which moves every forward kwarg to
# the executing device, wrong for tensors transformers marks device-invariant.
ns = _load_helpers(_fake_torch(), _FakeLogger())
device_map = {"model.embed": 0, "model.layers.0": 1}
model = _FakeModel(device_map = device_map, devices = ("cuda:0", "cuda:1"))
model._skip_keys_device_placement = ["past_key_values"]
ns["_restore_model_after_quantize_subprocess"](model, ("dispatch", device_map))
assert _fake_accelerate["dispatched"] == [(model, device_map)]
assert _fake_accelerate["dispatch_kwargs"] == [{"skip_keys": ["past_key_values"]}]
def test_snapshot_restores_a_forward_patched_after_the_dispatch(_fake_accelerate):
"""accelerate restores ``forward = _old_forward`` on removal, and ``_old_forward``
is the forward from when the hook was FIRST attached. unsloth patches forwards after
the dispatch, so a naive remove/re-add throws every fused kernel away for good."""
ns = _load_helpers(_fake_torch(), _FakeLogger())
root = _Child(device_map = {"model.embed": 0, "mlp": 1})
mlp = _Child(name = "mlp")
root._modules["mlp"] = mlp
stock_forward = lambda *a, **k: "stock" # noqa: E731
fused_forward = lambda *a, **k: "unsloth-fused" # noqa: E731
mlp._hf_hook = object()
mlp._old_forward = stock_forward # captured by accelerate at dispatch time
mlp.forward = fused_forward # installed by unsloth afterwards
snapshot = ns["_snapshot_dispatch_state"](root)
# what accelerate's removal does
del mlp.__dict__["_hf_hook"]
mlp.forward = mlp._old_forward
del mlp.__dict__["_old_forward"]
assert mlp.forward() == "stock"
ns["_restore_dispatch_state"](root, snapshot)
assert mlp.forward() == "unsloth-fused"
assert mlp.__dict__["_old_forward"] is stock_forward
def test_snapshot_reties_shared_parameters(_fake_accelerate):
"""A CPU round trip repoints every tensor, so replaying the hooks alone leaves tied
weights as independent copies: double VRAM, and updates to one never reach the other."""
import torch
root = _Child(device_map = {"embed": 0, "head": 0})
shared = torch.nn.Parameter(torch.zeros(4, 4))
for name in ("embed", "head"):
child = _Child(name = name)
child._parameters = {"weight": shared}
child._buffers = {}
root._modules[name] = child
def named(remove_duplicate = True):
seen, out = set(), []
for mod_name, mod in root._modules.items():
for attr, tensor in mod._parameters.items():
if remove_duplicate and id(tensor) in seen:
continue
seen.add(id(tensor))
out.append((f"{mod_name}.{attr}", tensor))
return iter(out)
root.named_parameters = named
ns = _load_helpers(_fake_torch(), _FakeLogger())
snapshot = ns_ties = ns["_snapshot_dispatch_state"](root)
assert ns_ties[3] == [["embed.weight", "head.weight"]]
# what the replay leaves behind before the retie step
root._modules["head"]._parameters["weight"] = torch.nn.Parameter(shared.detach().clone())
assert (
root._modules["embed"]._parameters["weight"].data_ptr()
!= root._modules["head"]._parameters["weight"].data_ptr()
)
ns["_restore_dispatch_state"](root, snapshot)
assert (
root._modules["embed"]._parameters["weight"].data_ptr()
== root._modules["head"]._parameters["weight"].data_ptr()
)
def test_meta_tensors_never_form_tie_groups(_fake_accelerate):
"""Offloaded parameters all sit on meta with storage pointer 0, so grouping by
pointer alone would collapse them into one fake tie and overwrite them all."""
import torch
root = _Child(device_map = {"a": 0, "b": "cpu", "c": "cpu"})
live = torch.nn.Parameter(torch.zeros(4, 4))
offloaded = [
torch.nn.Parameter(torch.empty(4, 4, device = "meta")),
torch.nn.Parameter(torch.empty(8, 2, device = "meta")),
]
def named(remove_duplicate = True):
return iter([("a.weight", live), ("b.weight", offloaded[0]), ("c.weight", offloaded[1])])
root.named_parameters = named
ns = _load_helpers(_fake_torch(), _FakeLogger())
_hooks, places, _attrs, ties, _grads = ns["_snapshot_dispatch_state"](root)
assert ties == [] # nothing is tied here
assert "b.weight" in places # still tracked for placement
def test_accelerate_move_guards_survive_the_replay(_fake_accelerate):
"""remove_hook_from_module also deletes the to/cuda/... guards dispatch_model
installs to stop a caller moving an offloaded model."""
ns = _load_helpers(_fake_torch(), _FakeLogger())
root = _Child(device_map = {"": 0})
guard = lambda *a, **k: "blocked" # noqa: E731
root._hf_hook = object()
root.to = guard
root.cuda = guard
snapshot = ns["_snapshot_dispatch_state"](root)
del root.__dict__["_hf_hook"], root.__dict__["to"], root.__dict__["cuda"]
ns["_restore_dispatch_state"](root, snapshot)
assert root.__dict__["to"] is guard
assert root.__dict__["cuda"] is guard
def test_gradients_survive_the_offload_round_trip():
"""init_hook rebuilds the Parameter and drops .grad, so the snapshot has to carry it."""
import torch
root = _Child(device_map = {"": 0})
weight = torch.nn.Parameter(torch.zeros(4, 4))
weight.grad = torch.full((4, 4), 3.0)
root._parameters = {"weight": weight}
root.named_parameters = lambda remove_duplicate = True: iter([("weight", weight)])
ns = _load_helpers(_fake_torch(), _FakeLogger())
snapshot = ns["_snapshot_dispatch_state"](root)
assert torch.equal(snapshot[4]["weight"], torch.full((4, 4), 3.0))
# What init_hook does: same name, fresh Parameter, no grad.
replacement = torch.nn.Parameter(torch.zeros(4, 4))
assert replacement.grad is None
root._parameters = {"weight": replacement}
ns["_restore_dispatch_state"](root, snapshot)
assert replacement.grad is not None, "the restore must put the gradient back"
assert torch.equal(replacement.grad, torch.full((4, 4), 3.0))
def test_the_other_torchao_path_also_clears_the_failed_copy():
"""Both torchao paths must drop the copy and the traceback pinning it before restoring."""
src = _SAVE_PY.read_text(encoding = "utf-8")
body = src.split("\ndef _unsloth_save_torchao(", 1)[1].split("\ndef ", 1)[0]
finally_block = body.split(" finally:", 1)[1]
assert "del quantized_model" in finally_block
assert "traceback.clear_frames" in finally_block
restore_at = finally_block.index("_restore_model_after_quantize_subprocess")
assert finally_block.index("del quantized_model") < restore_at
assert finally_block.index("traceback.clear_frames") < restore_at
def test_cpu_spill_rejection_is_retryable():
"""bitsandbytes rejects a CPU-spilled map with a ValueError that says nothing about
memory, so the single-device retry has to match it explicitly."""
import importlib.util
from pathlib import Path
export_py = (
Path(__file__).resolve().parent.parent
/ "studio"
/ "backend"
/ "core"
/ "export"
/ "export.py"
)
src = ast.parse(export_py.read_text(encoding = "utf-8"))
keep = [
n
for n in src.body
if isinstance(n, ast.FunctionDef) and n.name in {"_is_oom_error", "_is_cpu_spill_rejection"}
]
assert len(keep) == 2
namespace = {"torch": None}
exec( # noqa: S102 - loading trusted repo source
compile(ast.Module(body = keep, type_ignores = []), str(export_py), "exec"), namespace
)
bnb = ValueError(
"Some modules are dispatched on the CPU or the disk. Make sure you have enough "
"GPU RAM to fit the quantized model."
)
assert not namespace["_is_oom_error"](bnb)
assert namespace["_is_cpu_spill_rejection"](bnb)
assert namespace["_is_oom_error"](RuntimeError("CUDA out of memory. Tried to allocate 1 GiB"))
assert not namespace["_is_cpu_spill_rejection"](RuntimeError("some other failure"))
def test_torchao_releases_the_quantized_copy_in_finally():
"""If save_pretrained raises, the quantized copy must still be dropped before the
original is restored, or both are resident at once."""
src = _SAVE_PY.read_text(encoding = "utf-8")
body = src.split("def _unsloth_save_torchao_with_given_config(", 1)[1].split("\ndef ", 1)[0]
finally_block = body.split(" finally:", 1)[1]
assert "del quantized_model" in finally_block
assert "_restore_model_after_quantize_subprocess(model, model_restore)" in finally_block
# and the restore must come after the copy is dropped
assert finally_block.index("del quantized_model") < finally_block.index(
"_restore_model_after_quantize_subprocess"
)
# dropping the local is not enough: the live traceback still holds the frames
assert "traceback.clear_frames" in finally_block
assert finally_block.index("traceback.clear_frames") < finally_block.index(
"_restore_model_after_quantize_subprocess"
)
def test_a_live_traceback_pins_the_failed_copy_until_its_frames_are_cleared():
"""Why the clear_frames call above is load-bearing, on plain objects."""
import sys
import traceback
import weakref
class _Copy:
pass
def _build_and_fail(sink):
copy = _Copy() # noqa: F841 -- the point is that the frame retains it
sink.append(weakref.ref(copy))
raise RuntimeError("save_pretrained failed")
def _run(clear_frames):
# try/finally with the exception still in flight, exactly as in save.py
sink = []
alive = None
try:
try:
_build_and_fail(sink)
finally:
if clear_frames:
exc = sys.exc_info()[1]
if exc is not None:
traceback.clear_frames(exc.__traceback__)
gc.collect()
alive = sink[0]() is not None
except RuntimeError:
pass
return alive
assert _run(clear_frames = False), "expected the traceback to pin the copy"
assert not _run(clear_frames = True), "clear_frames must release it"

View file

@ -48,6 +48,7 @@ import functools
from transformers.models.llama.modeling_llama import logger
from .kernels import fast_dequantize, QUANT_STATE, get_lora_parameters_bias
import subprocess
import traceback
import psutil
import re
from transformers.models.llama.modeling_llama import logger
@ -1122,7 +1123,19 @@ def unsloth_save_model(
torch_dtype
)
max_vram = int(torch.cuda.get_device_properties(0).total_memory * maximum_memory_usage)
# A merged tensor lives on the GPU of its source layer, so budget against W's own
# device, not GPU0, else a sharded model OOMs GPU1+ while only GPU0 is checked.
_max_vram_by_device = {}
def _device_vram_budget(dev):
if dev.type != "cuda":
return None
idx = dev.index if dev.index is not None else torch.cuda.current_device()
if idx not in _max_vram_by_device:
_max_vram_by_device[idx] = int(
torch.cuda.get_device_properties(idx).total_memory * maximum_memory_usage
)
return _max_vram_by_device[idx]
print("Unsloth: Saving model... This might take 5 minutes ...")
@ -1138,8 +1151,15 @@ def unsloth_save_model(
if bias is not None:
state_dict[f"model.layers.{j}.{item}.bias"] = bias
if (torch.cuda.memory_allocated() + W.nbytes) < max_vram:
# Save to GPU memory
_dev_budget = _device_vram_budget(W.device)
if (
_dev_budget is not None
and (torch.cuda.memory_allocated(W.device) + W.nbytes) < _dev_budget
):
# Fits on W's own GPU
state_dict[name] = W
elif W.device.type != "cuda":
# Already off-GPU: keeping it costs no VRAM
state_dict[name] = W
# [TODO] Saving to RAM seems to leak memory???
# elif (max_ram - W.nbytes) > 0:
@ -4524,30 +4544,61 @@ def _unsloth_save_torchao_with_given_config(
else:
kwargs = {"dtype": torch.bfloat16}
# Reload with quantization applied
quantized_model = auto_model.from_pretrained(
save_directory,
device_map = "auto",
quantization_config = quantization_config,
**kwargs,
)
# Else the original stays resident on every GPU while device_map="auto" below
# loads a second copy.
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
torchao_save_directory = save_directory + "-torchao"
# TorchAO does not support safe_serialization right now 0.14.0 seems broken!
safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0")
safe_serialization = False
if push_to_hub:
quantized_model.push_to_hub(
torchao_save_directory, safe_serialization = safe_serialization, token = token
# The original stays offloaded until the quantized copy is saved AND released,
# else both are resident at once and the restore OOMs.
try:
# Reload with quantization applied
quantized_model = auto_model.from_pretrained(
save_directory,
device_map = "auto",
quantization_config = quantization_config,
**kwargs,
)
tokenizer.push_to_hub(torchao_save_directory, token = token)
else:
quantized_model.save_pretrained(
torchao_save_directory, safe_serialization = safe_serialization
)
tokenizer.save_pretrained(torchao_save_directory, token = token)
torchao_save_directory = save_directory + "-torchao"
# TorchAO does not support safe_serialization right now 0.14.0 seems broken!
safe_serialization = Version(importlib_version("torchao")) > Version("0.14.0")
safe_serialization = False
if push_to_hub:
quantized_model.push_to_hub(
torchao_save_directory, safe_serialization = safe_serialization, token = token
)
tokenizer.push_to_hub(torchao_save_directory, token = token)
else:
quantized_model.save_pretrained(
torchao_save_directory, safe_serialization = safe_serialization
)
tokenizer.save_pretrained(torchao_save_directory, token = token)
finally:
# del here, not at the end of the try: if save_pretrained raises, the copy
# would otherwise still be resident while the original is restored.
quantized_model = None
del quantized_model
# A failed save leaves a live traceback whose frames still hold the copy, so
# dropping the local alone does not free its VRAM.
_exc = sys.exc_info()[1]
if _exc is not None:
traceback.clear_frames(_exc.__traceback__)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
_restore_model_after_quantize_subprocess(model, model_restore)
# Clean up the intermediate unquantized model
if os.path.exists(save_directory):
@ -4585,6 +4636,318 @@ def _print_compressed_hw_note(scheme, out_dir):
)
_DISPATCH_SNAPSHOT_ATTR = "_unsloth_dispatch_snapshot"
def _accelerate_move_guards():
"""The instance methods dispatch_model wraps to block moving an offloaded model."""
try:
from accelerate.hooks import _accelerate_added_attributes
return tuple(_accelerate_added_attributes)
except Exception:
return ("to", "cuda", "npu", "xpu", "mlu", "sdaa", "musa")
_ACCELERATE_MOVE_GUARDS = _accelerate_move_guards()
def _accelerate_dispatch_root(model):
"""The module that really owns the accelerate dispatch.
A PEFT wrapper only proxies ``_hf_hook``, so ``delattr`` fails and
``remove_hook_from_submodules`` raises before removing anything; ``hf_device_map``
keys are relative to the inner root too. Walks real children, never ``__getattr__``.
"""
node, seen = model, set()
while id(node) not in seen:
seen.add(id(node))
if "hf_device_map" in getattr(node, "__dict__", {}):
return node
children = getattr(node, "__dict__", {}).get("_modules") or {}
nxt = next(
(
children[a]
for a in ("base_model", "model")
if hasattr(children.get(a), "named_modules")
),
None,
)
if nxt is None:
return model
node = nxt
return model
def _snapshot_dispatch_state(root):
"""Hooks, tensor placements and instance forwards, so the dispatch can be replayed.
Re-deriving it with ``dispatch_model`` is not equivalent: PEFT reparents each
targeted ``Linear`` after transformers dispatched, so accelerate hooks modules that
never had any (measured: 395 -> 1379) and the logits shift enough to reorder top-5.
"""
hooks = [
(name, mod.__dict__["_hf_hook"])
for name, mod in root.named_modules()
if "_hf_hook" in mod.__dict__
]
# remove_duplicate=False: the default hides one half of every tied pair, exactly
# the half that needs re-tying below.
named = list(root.named_parameters(remove_duplicate = False)) + list(
root.named_buffers(remove_duplicate = False)
)
places = {name: tensor.device for name, tensor in named}
# Tied weights share one storage, but the CPU round trip repoints every tensor and
# tied_params_map is keyed on the old pointer, so replaying the hooks alone gives
# independent copies: double VRAM, and updates to one no longer reach the other.
# Skip meta tensors: offloaded parameters all sit on meta with pointer 0, which
# would collapse into one fake "tied" group of differently shaped tensors, and
# each side of a tie is already its own meta placeholder so nothing is lost.
groups = {}
for name, tensor in named:
if tensor.device.type == "meta":
continue
ptr = tensor.untyped_storage().data_ptr()
if ptr:
groups.setdefault(ptr, []).append(name)
ties = [names for names in groups.values() if len(names) > 1]
# Removing a hook restores `forward = _old_forward`, captured before unsloth patched
# the module, so a remove/re-add permanently drops every fused kernel installed after
# the dispatch (measured: apply_lora_mlp_swiglu on all 28 MLPs). It also deletes the
# `to`/`cuda`/... move guards, so record those too.
attrs = ("forward", "_old_forward") + tuple(_ACCELERATE_MOVE_GUARDS)
saved_attrs = {
name: {a: mod.__dict__[a] for a in attrs if a in mod.__dict__}
for name, mod in root.named_modules()
if any(a in mod.__dict__ for a in attrs)
}
# Re-adding a hook runs init_hook -> set_module_tensor_to_device, which builds a fresh
# Parameter and so drops .grad. Snapshot the gradients and reattach them on restore.
grads = {
name: getattr(tensor, "grad", None)
for name, tensor in root.named_parameters(remove_duplicate = False)
if getattr(tensor, "grad", None) is not None
}
return hooks, places, saved_attrs, ties, grads
def _drop_accelerator_tied_param_cache(snapshot) -> None:
"""Drop the GPU tensors accelerate caches in each hook's ``tied_params_map``.
Holding the hooks across the offload pins a GPU copy of the tied embedding (0.31 GB
of 1.24 GB here). The entries are keyed on the pre-move ``data_ptr`` so they are
stale anyway, and re-attaching repopulates them.
"""
for _name, hook in snapshot[0]:
cache = getattr(hook, "tied_params_map", None)
if not cache:
continue
for ptr in list(cache):
entry = cache[ptr]
for device in list(entry):
if str(device) != "cpu":
del entry[device]
if not entry:
del cache[ptr]
def _split_tensor_path(root, full_name):
"""``("model.embed_tokens.weight")`` -> ``(the module, "weight")``."""
mod_name, _, attr = full_name.rpartition(".")
try:
return (root.get_submodule(mod_name) if mod_name else root), attr
except AttributeError:
return None, attr
def _lookup_tensor(root, full_name):
mod, attr = _split_tensor_path(root, full_name)
if mod is None:
return None
for store in ("_parameters", "_buffers"):
found = (getattr(mod, store, None) or {}).get(attr)
if found is not None:
return found
return None
def _share_tensor(root, full_name, leader) -> None:
"""Point ``full_name`` back at ``leader``, restoring a tie."""
mod, attr = _split_tensor_path(root, full_name)
if mod is None:
return
for store in ("_parameters", "_buffers"):
target = getattr(mod, store, None)
if target is None or attr not in target:
continue
current = target[attr]
if current is None or current.device != leader.device or current.shape != leader.shape:
return # not actually the same tensor; leave it alone
target[attr] = leader
return
def _restore_dispatch_state(root, snapshot) -> None:
"""Replay ``_snapshot_dispatch_state``."""
from accelerate.hooks import add_hook_to_module
hooks, places, saved_attrs, ties, grads = snapshot
for name, hook in hooks:
add_hook_to_module(root.get_submodule(name) if name else root, hook)
# Re-adding a hook rewraps whatever `_old_forward` now holds, so put the exact
# callables back, `_old_forward` first.
for name, values in saved_attrs.items():
mod = root.get_submodule(name) if name else root
for attr in ("_old_forward", "forward", *_ACCELERATE_MOVE_GUARDS):
if attr in values:
mod.__dict__[attr] = values[attr]
# init_hook only re-places tensors the hooked module owns, so anything added after
# the dispatch (the LoRA adapters) is still on CPU.
for mod_name, mod in root.named_modules():
for attr in ("_parameters", "_buffers"):
store = getattr(mod, attr, None)
if not store:
continue
for tensor_name, tensor in list(store.items()):
if tensor is None:
continue
full = f"{mod_name}.{tensor_name}" if mod_name else tensor_name
want = places.get(full)
if want is None or tensor.device == want:
continue
if getattr(tensor, "quant_state", None) is not None:
# Only bitsandbytes' own .to() moves absmax/code/state2 with the data.
mod.to(want)
else:
tensor.data = tensor.data.to(want)
# Reattach the gradients init_hook discarded, on their weight's device.
for name, grad in grads.items():
tensor = _lookup_tensor(root, name)
if tensor is not None and tensor.grad is None and tensor.shape == grad.shape:
tensor.grad = grad.to(tensor.device)
# Re-tie last, once every tensor is back on its own device.
for names in ties:
leader = _lookup_tensor(root, names[0])
if leader is None:
continue
for follower in names[1:]:
_share_tensor(root, follower, leader)
# init_hook refilled tied_params_map with the pre-retie tensors, now unreferenced
# by the model but still pinned by the map.
if ties:
_drop_accelerator_tied_param_cache(snapshot)
def _offload_model_for_quantize_subprocess(model):
"""Best-effort: move the model's weights off the GPU before the quantized export
loads its own copy from disk, so the GPUs need not hold both at once. Returns an
opaque token for ``_restore_model_after_quantize_subprocess`` (None if nothing moved).
Two shapes are handled:
* single-device CUDA/XPU model -> ``.to("cpu")``, restored with ``.to(device)``;
* accelerate-dispatched model (a multi-GPU ``device_map`` shard, e.g. the Studio
multi-GPU export load) -> hooks removed and moved to CPU, restored by replaying
the dispatch. A plain ``.to("cpu")`` is invalid here, which is why the old
single-device-only move left every GPU holding a full copy. A map spilling to
CPU is still released, but disk/meta targets are left alone: accelerate keeps
those parameters off the model, so moving would materialize the whole checkpoint.
Quantized (bnb) models are attempted too rather than skipped: Studio exports load
4-bit by DEFAULT, so skipping them left a shard on every GPU. transformers refuses
``.to()`` for some bitsandbytes builds and that refusal raises before anything moves,
so the failure path restores the model and returns None, i.e. the old behaviour.
"""
try:
_has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
if not ((torch.cuda.is_available() or _has_xpu) and hasattr(model, "parameters")):
return None
device_map = getattr(model, "hf_device_map", None)
if device_map:
targets = {str(v).lower() for v in device_map.values()}
# A cpu spill is fine to move, it is already in host RAM. disk/meta is not:
# those parameters are off the model, so .to("cpu") would materialize the
# whole checkpoint into RAM.
if not all(t.isdigit() or t.startswith(("cuda", "xpu")) or t == "cpu" for t in targets):
return None
if not any(t.isdigit() or t.startswith(("cuda", "xpu")) for t in targets):
return None # nothing on an accelerator: no GPU memory to reclaim
from accelerate.hooks import remove_hook_from_submodules
# A PEFT wrapper only proxies the hooks; they live on the inner root.
root = _accelerate_dispatch_root(model)
try:
setattr(root, _DISPATCH_SNAPSHOT_ATTR, _snapshot_dispatch_state(root))
except Exception as snap_exc:
# Restore will fall back to re-deriving from the device_map.
logger.warning_once(
f"Unsloth: could not snapshot the accelerate dispatch "
f"({type(snap_exc).__name__}: {snap_exc}); re-dispatching on restore."
)
remove_hook_from_submodules(root)
try:
model.to("cpu")
except Exception:
# The move failed after the hooks came off; re-dispatch so the model is
# left usable rather than hookless and half-moved across CPU/GPUs.
_restore_model_after_quantize_subprocess(model, ("dispatch", dict(device_map)))
return None
snapshot = getattr(root, _DISPATCH_SNAPSHOT_ATTR, None)
if snapshot is not None:
_drop_accelerator_tied_param_cache(snapshot)
return ("dispatch", dict(device_map))
devices = {str(p.device) for p in model.parameters()}
if len(devices) == 1 and next(iter(devices)).startswith(("cuda", "xpu")):
device = next(model.parameters()).device
try:
model.to("cpu")
except Exception:
_restore_model_after_quantize_subprocess(model, ("device", device))
return None
return ("device", device)
except Exception as exc:
# A silent `return None` is indistinguishable from "nothing to move", which
# hides a real bug behind a merely slower export.
logger.warning_once(
f"Unsloth: could not free the model's accelerator memory before the quantized "
f"export ({type(exc).__name__}: {exc}); continuing with the model resident."
)
return None
return None
def _restore_model_after_quantize_subprocess(model, restore_token) -> None:
"""Undo ``_offload_model_for_quantize_subprocess``; warns instead of raising."""
if restore_token is None:
return
kind, value = restore_token
try:
if kind == "dispatch":
root = _accelerate_dispatch_root(model)
snapshot = root.__dict__.pop(_DISPATCH_SNAPSHOT_ATTR, None)
if snapshot is not None:
_restore_dispatch_state(root, snapshot)
else:
from accelerate import dispatch_model
# skip_keys matters: without it accelerate moves every forward kwarg
# to the executing device, wrong for device-invariant cache tensors.
dispatch_model(
root,
device_map = value,
skip_keys = getattr(root, "_skip_keys_device_placement", None),
)
else:
model.to(value) # restore the model to its original device
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device(s) after the "
"quantized export; it may remain on CPU."
)
def _unsloth_save_compressed_tensors(
model,
save_directory: Union[str, os.PathLike],
@ -4654,7 +5017,7 @@ def _unsloth_save_compressed_tensors(
# 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and
# quantize inside an isolated temp dir instead of writing ./<repo_id> into the cwd.
repo_id, work_tmp, calib_tmp, model_dev = None, None, None, None
repo_id, work_tmp, calib_tmp, model_restore = None, None, None, None
if push_to_hub:
repo_id = os.fspath(save_directory)
work_tmp = tempfile.mkdtemp(prefix = "unsloth-compressed-")
@ -4806,23 +5169,8 @@ def _unsloth_save_compressed_tensors(
cmd += ["--variant", variant]
# Free the in-memory model's CUDA memory before the subprocess loads its own copy from
# disk, so a single GPU need not hold both at once. Best-effort and restored in finally;
# skipped for quantized or multi-device models where moving is unsafe.
try:
if (
torch.cuda.is_available()
and hasattr(model, "parameters")
and not getattr(model, "is_loaded_in_4bit", False)
and not getattr(model, "is_loaded_in_8bit", False)
and not getattr(model, "is_quantized", False)
):
_devs = {str(p.device) for p in model.parameters()}
if len(_devs) == 1 and next(iter(_devs)).startswith("cuda"):
_dev = next(model.parameters()).device
model.to("cpu")
model_dev = _dev # set only after a successful move, so finally can restore
except Exception:
model_dev = None
# disk, so the GPUs need not hold both at once. Best-effort, restored in finally.
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
@ -4893,14 +5241,7 @@ def _unsloth_save_compressed_tensors(
_print_compressed_hw_note(scheme, result)
return result
finally:
if model_dev is not None:
try:
model.to(model_dev) # restore the model to its original device
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device after compressed "
"export; it may remain on CPU."
)
_restore_model_after_quantize_subprocess(model, model_restore)
if calib_tmp is not None and os.path.isdir(calib_tmp):
shutil.rmtree(calib_tmp, ignore_errors = True)
if work_tmp is not None:
@ -4959,7 +5300,7 @@ def _unsloth_save_torchao(
# Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected
# 16-bit export written to save_directory is not overwritten or deleted; the torchao output is
# the sibling "<save_directory>-<suffix>" (or the repo id on a hub push).
repo_id, work_tmp, model_dev = None, None, None
repo_id, work_tmp, model_restore = None, None, None
work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-")
if push_to_hub:
repo_id = os.fspath(save_directory)
@ -5037,25 +5378,12 @@ def _unsloth_save_torchao(
auto_model = AutoModelForCausalLM
auto_processor = AutoProcessor if is_vlm else AutoTokenizer
# 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk.
# Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit
# resident alongside the reloaded copy and OOM a device that fit the model once.
# 3) Free the in-memory model's accelerator memory before reloading a fresh copy from
# disk, else it sits resident alongside the copy and OOMs a device that fit the
# model once. Covers CUDA and XPU (torchao runs on Intel GPUs too) plus multi-GPU
# dispatched shards, which a plain .to("cpu") cannot move.
_has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available()
try:
if (
(torch.cuda.is_available() or _has_xpu)
and hasattr(model, "parameters")
and not getattr(model, "is_loaded_in_4bit", False)
and not getattr(model, "is_loaded_in_8bit", False)
and not getattr(model, "is_quantized", False)
):
_devs = {str(p.device) for p in model.parameters()}
if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")):
_dev = next(model.parameters()).device
model.to("cpu")
model_dev = _dev
except Exception:
model_dev = None
model_restore = _offload_model_for_quantize_subprocess(model)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
@ -5126,14 +5454,20 @@ def _unsloth_save_torchao(
)
return result
finally:
if model_dev is not None:
try:
model.to(model_dev)
except Exception:
logger.warning_once(
"Unsloth: could not restore the model to its original device after torchao "
"export; it may remain on CPU."
)
# A raise pins the copy in the local and the live traceback, so free both or the
# restore below OOMs.
quantized_model = None
del quantized_model
_exc = sys.exc_info()[1]
if _exc is not None:
traceback.clear_frames(_exc.__traceback__)
for _ in range(3):
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
if hasattr(torch, "xpu") and torch.xpu.is_available():
torch.xpu.empty_cache()
_restore_model_after_quantize_subprocess(model, model_restore)
if work_tmp is not None:
shutil.rmtree(work_tmp, ignore_errors = True)
for _ in range(3):

View file

@ -559,11 +559,18 @@ def _subagent_model_id(
)
if status.get("is_gguf"):
variant = status.get("gguf_variant")
return (
_display_model_spec(model_id, str(variant))
if variant and _is_hub_model_id(model_id)
else model_id
)
if variant and _is_hub_model_id(model_id):
return _display_model_spec(model_id, str(variant))
if variant:
# A path load is advertised as a bare basename with no ":variant" channel,
# so the quant cannot be recorded and a later reload picks for itself.
typer.echo(
f"Warning: {model_id} loaded from a path, so the subagent config cannot "
f"pin the {variant} quant; a reload may choose a different one. Load the "
"model by repository id to pin it.",
err = True,
)
return model_id
def _fail(message: str) -> NoReturn:
@ -572,9 +579,8 @@ def _fail(message: str) -> NoReturn:
def _reject_as_subagent(agent: str, args: list) -> None:
# Reject early; otherwise the flag reaches the agent binary and fails after
# Studio has already loaded the model.
if "--as-subagent" in args:
# Reject early, or the flag reaches the agent binary after Studio loaded the model.
if any(arg == "--as-subagent" or arg.startswith("--as-subagent=") for arg in args):
_fail(f"--as-subagent is not supported for {agent}.")
@ -1386,6 +1392,37 @@ def _is_hub_model_id(value: object) -> bool:
return True
def _is_model_path(value: str) -> bool:
"""Mirrors core.inference.model_ids._looks_like_path: a repo id is exactly
``org/model``; anything else with a separator, drive, prefix or .gguf is a path.
Deliberately not named _looks_like_path: that name is taken further down by the
WSLENV classifier, which only matches absolute paths and would shadow this one.
"""
if value.lower().endswith(".gguf"):
return True
if value.startswith(("/", "\\", "./", "../", ".\\", "..\\", "~")):
return True
if len(value) >= 2 and value[1] == ":":
return True
return value.count("/") >= 2 or "\\" in value
def _public_model_id(value: Optional[str]) -> Optional[str]:
"""The id Unsloth advertises for a model loaded by path.
/v1/models never echoes a host path: it reports the file or directory name
with any .gguf suffix stripped (core.inference.model_ids.public_model_id), so
a path we asked to load has to be matched by that name too.
"""
if not value or not _is_model_path(value):
return None
name = os.path.basename(value.replace("\\", "/").rstrip("/"))
if name.lower().endswith(".gguf"):
name = name[: -len(".gguf")]
return name or None
def _model_id_matches(
actual: object,
requested: object,
@ -1486,7 +1523,7 @@ def _resolve_model(
# casing) that /v1/models echoes but which may differ from the path we
# passed; match on the id the load reports so we don't silently fall
# through to models[0] and connect to a different loaded model.
wanted = {requested}
wanted = {requested, _public_model_id(requested)} - {None}
if isinstance(loaded, dict):
wanted |= {loaded.get("model"), loaded.get("display_name")} - {None}
models = _loaded_models(base, key)
@ -1954,7 +1991,15 @@ def _opencode_subagent_inline_config(path: Path, permission: dict) -> dict:
def merge_provider_filters(effective_config: dict) -> None:
enabled = effective_config.get("enabled_providers")
if isinstance(enabled, list):
inline["enabled_providers"] = list(dict.fromkeys([*enabled, _OPENCODE_PROVIDER]))
inherited_enabled = inline.get("enabled_providers")
if not isinstance(inherited_enabled, list):
inherited_enabled = []
providers = [
provider
for provider in [*inherited_enabled, *enabled]
if provider != _OPENCODE_PROVIDER
]
inline["enabled_providers"] = list(dict.fromkeys([*providers, _OPENCODE_PROVIDER]))
disabled = effective_config.get("disabled_providers")
if isinstance(disabled, list) and _OPENCODE_PROVIDER in disabled:
inline["disabled_providers"] = [
@ -2871,7 +2916,13 @@ def write_pi_config(base: str, key: str, model: dict, path: Path) -> None:
typer.echo(f"Updated {path}")
def write_pi_subagent_config(base: str, key: str, model: dict, path: Path) -> None:
def write_pi_subagent_config(
base: str,
key: str,
model: dict,
path: Path,
approve: bool = False,
) -> None:
"""Write private bootstrap data for the bundled Pi extension."""
window = model.get("context_length") or model.get("max_context_length")
window = int(window) if window else 32768
@ -2883,6 +2934,7 @@ def write_pi_subagent_config(base: str, key: str, model: dict, path: Path) -> No
"model": model["id"],
"contextWindow": window,
"maxTokens": min(window // 4, 8192),
"approve": approve,
},
)
@ -3452,7 +3504,13 @@ def pi(
extension = _agent_config_path(_PI_SUBAGENT_EXTENSION, ["pi"])
with _session_config("pi-subagent", launch, persist = persist) as config:
config_path = config / "subagent.json"
write_pi_subagent_config(base, key, subagent_model, config_path)
write_pi_subagent_config(
base,
key,
subagent_model,
config_path,
approve = yolo,
)
command = [
"pi",
"--extension",

View file

@ -5,7 +5,8 @@ import { fileURLToPath } from "node:url";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
const provider = "unsloth";
// Distinct from the normal `unsloth` provider: subagent mode preserves the user's Pi config.
const provider = "unsloth-studio-subagent";
const maxResultCharacters = 100_000;
const maxParallelAgents = 4;
const cancelGraceMilliseconds = 2_000;
@ -26,6 +27,7 @@ if (configPath) {
const model = typeof config.model === "string" ? config.model : "";
const baseUrl = typeof config.baseUrl === "string" ? config.baseUrl : "";
const apiKey = typeof config.apiKey === "string" ? config.apiKey : "";
const approve = config.approve === true;
const contextWindow = positiveInt(config.contextWindow, 32768);
const maxTokens = positiveInt(config.maxTokens, Math.min(Math.floor(contextWindow / 4), 8192));
let activeAgents = 0;
@ -168,6 +170,7 @@ async function runLocalAgent(
"json",
"--print",
"--no-session",
...(approve ? ["--approve"] : []),
"--provider",
provider,
"--model",

View file

@ -885,8 +885,9 @@ def test_subagent_model_id_warns_when_status_unavailable(monkeypatch, capsys):
@pytest.mark.parametrize("agent", ["openclaw", "hermes"])
def test_unsupported_agents_reject_as_subagent(agent):
result = CliRunner().invoke(start.start_app, [agent, "--as-subagent"])
@pytest.mark.parametrize("flag", ["--as-subagent", "--as-subagent=true", "--as-subagent=false"])
def test_unsupported_agents_reject_as_subagent(agent, flag):
result = CliRunner().invoke(start.start_app, [agent, flag])
assert result.exit_code == 1
assert f"--as-subagent is not supported for {agent}." in result.output
@ -1296,6 +1297,66 @@ def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch, cap
assert "please wait" not in output
def test_resolve_model_matches_snapshot_path_by_public_id(monkeypatch):
"""A GGUF loaded by snapshot path is advertised by its basename, not the path."""
snapshot = "/home/u/.cache/legacy/models--Org--Model/snapshots/abc123"
state = {"loaded": False}
def http_json(
method,
url,
token,
payload = None,
timeout = 30,
error = None,
):
if url.endswith("/v1/models"):
return {"data": [{"id": "abc123"}] if state["loaded"] else []}
if url.endswith("/api/inference/load"):
state["loaded"] = True
# The load echoes the path it was given, which /v1/models never lists.
return {"model": snapshot, "display_name": snapshot}
raise AssertionError(f"unexpected request: {method} {url}")
monkeypatch.setattr(start, "_http_json", http_json)
entry = start._resolve_model(BASE, "sk-test", snapshot, start.LoadOptions())
assert entry["id"] == "abc123"
def test_subagent_model_id_warns_when_a_path_load_cannot_pin_the_quant(capsys):
"""A path is advertised as a bare basename, so the quant cannot be recorded."""
model_id = start._subagent_model_id(BASE, "sk-test", {"id": "abc123"}, None, "UD-Q4_K_XL")
assert model_id == "abc123"
assert "cannot pin the UD-Q4_K_XL quant" in capsys.readouterr().err
def test_subagent_model_id_pins_the_quant_for_repo_ids(capsys):
model_id = start._subagent_model_id(
BASE, "sk-test", {"id": "unsloth/gemma-4-E4B-it-GGUF"}, None, "UD-Q4_K_XL"
)
assert model_id == "unsloth/gemma-4-E4B-it-GGUF:UD-Q4_K_XL"
assert capsys.readouterr().err == ""
def test_public_model_id_leaves_repo_ids_alone():
"""Only a path gets reduced; a repo id must not match some unrelated model.
Relative and multi-segment paths are covered too: _looks_like_path is defined
twice in this module (the WSLENV one wins), so this must use its own classifier.
"""
assert start._public_model_id("unsloth/gemma-4-E4B-it-GGUF") is None
assert start._public_model_id("org/model") is None
assert start._public_model_id("/srv/models/Qwen3-Q4_K_M.gguf") == "Qwen3-Q4_K_M"
assert start._public_model_id("/a/b/snapshots/rev1") == "rev1"
assert start._public_model_id("./models/foo") == "foo"
assert start._public_model_id("cache/snapshots/rev") == "rev"
assert start._public_model_id("a/b/c") == "c"
def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch):
# A cached-but-unloaded catalog entry (loaded == False) that only case-differs must
# not be treated as ready; the load endpoint must still be called so the requested
@ -3296,9 +3357,13 @@ def test_write_opencode_config_as_subagent_preserves_parent_model(tmp_path):
def test_opencode_subagent_inline_keeps_parent_provider_filters(monkeypatch, tmp_path):
config_path = tmp_path / "opencode.json"
inherited = {"theme": "tokyonight"}
inherited = {
"theme": "tokyonight",
"enabled_providers": ["anthropic"],
}
monkeypatch.setenv("OPENCODE_CONFIG_CONTENT", json.dumps(inherited))
monkeypatch.setattr(start, "_which_with_install_dirs", lambda _: "/usr/bin/opencode")
monkeypatch.setattr(start, "_wsl_windows_executable", lambda _: None)
captured = {}
def run(command, **kwargs):
@ -3324,7 +3389,11 @@ def test_opencode_subagent_inline_keeps_parent_provider_filters(monkeypatch, tmp
assert captured["env"]["OPENCODE_CONFIG"] == str(config_path)
assert inline == {
"theme": "tokyonight",
"enabled_providers": ["opencode-go", start._OPENCODE_PROVIDER],
"enabled_providers": [
"anthropic",
"opencode-go",
start._OPENCODE_PROVIDER,
],
"disabled_providers": ["ollama"],
"subagent_depth": 1,
"permission": permission,
@ -3721,21 +3790,26 @@ def test_connect_pi_no_launch(fake_studio, tmp_path):
assert not any(c[1].endswith("/api/inference/status") for c in fake_studio)
def test_connect_pi_as_subagent_preserves_cloud_parent(fake_studio, tmp_path):
@pytest.mark.parametrize("yolo", [False, True])
def test_connect_pi_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, yolo):
args = [
"pi",
"--as-subagent",
"--no-launch",
"--model",
MODEL["id"] + ":UD-Q4_K_XL",
]
if yolo:
args.insert(2, "--yolo")
result = CliRunner().invoke(
start.start_app,
[
"pi",
"--as-subagent",
"--no-launch",
"--model",
MODEL["id"] + ":UD-Q4_K_XL",
],
args,
)
assert result.exit_code == 0, result.output
command = _launch_command(result.output)
assert command[:2] == ["pi", "--extension"]
assert command[2].endswith("unsloth_cli/pi_subagent.ts")
assert ("--approve" in command) is yolo
assert "--provider" not in command
assert "--model" not in command
assert "PI_CODING_AGENT_DIR" not in result.output
@ -3750,6 +3824,7 @@ def test_connect_pi_as_subagent_preserves_cloud_parent(fake_studio, tmp_path):
"model": MODEL["id"] + ":UD-Q4_K_XL",
"contextWindow": 4096,
"maxTokens": 1024,
"approve": yolo,
}
assert "Ask Pi to spawn an Unsloth or local agent." in result.output