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:
commit
b7b37e165c
62 changed files with 9917 additions and 436 deletions
84
install.sh
84
install.sh
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
242
studio/backend/tests/test_export_multi_gpu_device_map.py
Normal file
242
studio/backend/tests/test_export_multi_gpu_device_map.py
Normal 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
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
9
studio/frontend/public/agent-logos/hermes.svg
Normal file
9
studio/frontend/public/agent-logos/hermes.svg
Normal 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 |
18
studio/frontend/public/agent-logos/openclaw.svg
Normal file
18
studio/frontend/public/agent-logos/openclaw.svg
Normal 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 |
19
studio/frontend/public/agent-logos/opencode-dark.svg
Normal file
19
studio/frontend/public/agent-logos/opencode-dark.svg
Normal 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 |
19
studio/frontend/public/agent-logos/opencode-light.svg
Normal file
19
studio/frontend/public/agent-logos/opencode-light.svg
Normal 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 |
21
studio/frontend/public/agent-logos/pi.svg
Normal file
21
studio/frontend/public/agent-logos/pi.svg
Normal 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 |
|
|
@ -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]);
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -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 ||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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. */
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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[] {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
1435
studio/frontend/src/features/settings/tabs/agents-tab.tsx
Normal file
1435
studio/frontend/src/features/settings/tabs/agents-tab.tsx
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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..."
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
190
tests/fast_inference/test_fast_inference.py
Normal file
190
tests/fast_inference/test_fast_inference.py
Normal 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")
|
||||
|
|
@ -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 ]
|
||||
|
|
|
|||
391
tests/studio/install/test_llama_prebuilt_no_space.py
Normal file
391
tests/studio/install/test_llama_prebuilt_no_space.py
Normal 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
|
||||
|
|
@ -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 ") :])
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
244
tests/studio/test_pdf_qa_recipe_contract.py
Normal file
244
tests/studio/test_pdf_qa_recipe_contract.py
Normal 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
|
||||
670
tests/test_compressed_export_gpu_release.py
Normal file
670
tests/test_compressed_export_gpu_release.py
Normal 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"
|
||||
490
unsloth/save.py
490
unsloth/save.py
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue