diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 3968f2e80a..ec437e0c32 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -30,6 +30,13 @@ on: - 'unsloth/**' - 'unsloth_cli/**' - 'tests/**' + # The root installers: tests/sh/*.sh and tests/studio/install/* assert + # against these two files, so a change here must run the suite that + # covers it. Without them an install-only edit (the shape most AMD/ROCm + # routing fixes take) skipped Backend CI entirely. + - 'install.sh' + - 'install.ps1' + - 'scripts/**' - 'pyproject.toml' - '.github/workflows/studio-backend-ci.yml' push: @@ -217,27 +224,32 @@ jobs: tests/studio/test_xpu_spoof_pipeline.py - name: Shell installer tests - # Subset that does not depend on a writable / pristine install.sh - # tree; test_install_host_defaults.sh checks install.ps1 layout - # which has drifted (separate followup). + # Auto-discovered rather than allowlisted. The old hardcoded list had + # silently fallen seven files behind tests/run_all.sh, including + # test_strixhalo_wsl_reroute.sh -- the only shell coverage of the ROCm + # WSL reroute -- so that suite never ran on a PR. Skips are explicit, + # each with a reason, and tests/studio/test_ci_shell_suite_coverage.py + # fails if this step stops discovering the directory or the skip list + # grows without one. + # + # Skipped: + # test_install_host_defaults.sh: asserts an install.ps1 layout that + # has drifted (separate followup). + # test_install_rollback_lifecycle.sh: already runs on both platforms + # in cross-platform-parity-ci.yml. run: | set -e - for s in \ - tests/sh/test_get_torch_index_url.sh \ - tests/sh/test_mac_intel_compat.sh \ - tests/sh/test_node_decision.sh \ - tests/sh/test_studio_home_node_dir.sh \ - tests/sh/test_system_node_readonly.sh \ - tests/sh/test_nvcc_meets_llama_minimum.sh \ - tests/sh/test_resolve_cuda_archs.sh \ - tests/sh/test_staged_validation_enabled.sh \ - tests/sh/test_tauri_install_exit_order.sh \ - tests/sh/test_torch_constraint.sh \ - tests/sh/test_torch_flavor.sh \ - tests/sh/test_with_llama_cpp_dir_flag.sh \ - tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do + skip="test_install_host_defaults.sh test_install_rollback_lifecycle.sh" + found=0 + for s in tests/sh/test_*.sh; do + case " $skip " in + *" $(basename "$s") "*) echo "skipping $s (see workflow comment)"; continue ;; + esac + found=$((found + 1)) echo "::group::$s" bash "$s" echo "::endgroup::" done + [ "$found" -gt 0 ] || { echo "::error::no shell tests discovered under tests/sh"; exit 1; } + echo "ran $found shell installer test files" diff --git a/install.ps1 b/install.ps1 index 8a962c0662..96078ae06d 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1790,12 +1790,14 @@ exit 0 # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI); targets only arches the ROCm prebuilts cover (gfx120X/110X/1151/1150/103X), unknown names fall back to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( - @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080) - @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060) + @{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: RX 9070 XT / 9070 GRE / 9070 / 9080) + @{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: RX 9060 XT / 9060) @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) - @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) - @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) + @{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + @{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + @{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) + @{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32) + @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33) @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family @@ -2087,6 +2089,7 @@ exit 0 $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) + "gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point) "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000) @@ -2099,15 +2102,18 @@ exit 0 $torchFloorMap = @{ "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" + "gfx1152" = "torch>=2.11.0,<2.12.0" } # Companion ranges track the torch ceiling so pip resolves a consistent trio on AMD's per-arch index (each published independently). Mirrors setup.ps1 / install_python_stack.py; bump all three together for 2.12.x. $torchvisionFloorMap = @{ "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" + "gfx1152" = "torchvision>=0.26.0,<0.27.0" } $torchaudioFloorMap = @{ "gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0" "gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0" + "gfx1152" = "torchaudio>=2.11.0,<2.12.0" } $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } if ($archFamily) { @@ -2135,7 +2141,7 @@ exit 0 $_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2) } # Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare. - $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf + $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $_pinLeaf if ($_pinGfx211 -or $_pinRocm211) { $ROCmIndexUrl = $TorchIndexUrl $ROCmTorchFloor = "torch>=2.11.0,<2.12.0" diff --git a/install.sh b/install.sh index 8808138386..b38d952705 100755 --- a/install.sh +++ b/install.sh @@ -624,6 +624,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/null 2>&1 +} + # ── Helper: install packages via apt, escalating to sudo only if needed ── # Usage: _smart_apt_install pkg1 pkg2 pkg3 ... _smart_apt_install() { @@ -664,24 +673,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/null; then + if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]' /proc/cpuinfo 2>/dev/null; then echo gfx1150 return 0 fi + if [ -n "$_gpu_evidence" ] && grep -qiE '860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then + echo gfx1152 + return 0 + fi if command -v lspci >/dev/null 2>&1; then # A non-AMD controller can enumerate first (Intel/ASPEED before an AMD # dGPU), so scan every display-class line and take the first AMD one @@ -2792,7 +2847,7 @@ if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \ # whole handoff (a user-set override re-exports unchanged). export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx" case "$_linux_inferred_gfx" in - gfx1201|gfx1200|gfx1151|gfx1150) + gfx1201|gfx1200|gfx1151|gfx1150|gfx1152) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" @@ -2855,7 +2910,7 @@ fi # rocm7.2 and the per-gfx indexes (Strix _grouped_mm fix) ship torch 2.11.0: raise the floor and pin companions; match the FINAL leaf only. # (cu*/cpu/custom leaves all use the default <2.12 trio above.) case "$_torch_index_leaf" in - rocm7.2|gfx120x-all|gfx1151|gfx1150) + rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" @@ -2953,7 +3008,7 @@ case "$_torch_index_leaf" in fi _strix_gfx="" case "$_runtime_gfx" in - gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;; + gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;; esac # Skip rocm7.13+ generic indexes: they already ship the fixes, so the # arch build (rocm7.13) would be a downgrade rather than a rescue. @@ -3037,12 +3092,14 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then # Kept in sync with install.ps1 nameArchTable; gfx1102 matched before gfx1100 ("RX 7700S"). case "$_gpu_disp_mkt" in - *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 - *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 + *9070*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 (Navi 48) + *9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 (Navi 44) *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) - *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) + *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1152" ;; # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) + *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) _gpu_disp_gfx="gfx1101" ;; # RDNA 3 (Navi 32) + *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point) *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _gpu_disp_gfx="gfx1030" ;; # RDNA 2 (Navi 21) *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _gpu_disp_gfx="gfx1032" ;; # RDNA 2 (Navi 23) @@ -3663,9 +3720,11 @@ echo "" # Interactive terminals prompt before starting; non-interactive environments 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 str: Unsloth Studio Login (Colab)

- Log in to Studio with the Cloudflare link above using these credentials. This cell - is visible only in your notebook session. + Log in as {username} with this password. This cell is visible only in + your notebook session.

- Username: {username}
Password: {password}

@@ -441,8 +440,29 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: return False -def _shareable_link_html(cloudflare_url: str) -> str: - """Branded card for the shareable Cloudflare link, styled like the show_link banner.""" +def _shareable_link_html( + cloudflare_url: str, + password: "str | None" = None, + username: "str | None" = None, +) -> str: + """Branded card for the shareable Cloudflare link, styled like the show_link banner. + + *password* renders under the link so the credential sits in the card with the button + it unlocks. The username is always the default admin, so it reads inline. + """ + login_block = "" + if password: + login_block = f""" +

+ Password +

+

{password}

+

+ Log in as {username} with this password. Shown only in your + notebook session, and never included in the shared link. +

""" return f"""
@@ -460,11 +480,12 @@ def _shareable_link_html(cloudflare_url: str) -> str: Open Unsloth Studio

- This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab. + This Cloudflare HTTPS link works from any device, so you can share it with anyone.

- 🔗 {cloudflare_url} -

+ 🔗 {cloudflare_url} +

{login_block}
""" @@ -555,28 +576,37 @@ def _show_and_embed( cloudflare_url = cloudflare_url, ) + # Fold the credentials into the link card rather than a second card below it. + credentials_shown = False if cloudflare_url: try: from IPython.display import HTML, display - display(HTML(_shareable_link_html(cloudflare_url))) + + username, password = colab_login if colab_login else (None, None) + display(HTML(_shareable_link_html(cloudflare_url, password, username))) + credentials_shown = bool(colab_login) except Exception as e: logger.info(f"Could not render Cloudflare link card ({e}).") - if colab_login: + if colab_login and not credentials_shown: try: _show_colab_login_credentials(*colab_login) except Exception as e: logger.info(f"Could not render Colab login card ({e}).") - try: - show_link( - port, - _url = url, - has_cloudflare_link = bool(cloudflare_url), - cloudflare_requested = cloudflare_requested, - ) - except Exception as e: - logger.info(f"Could not render Unsloth link card ({e}).") + # With a tunnel up the embed below is skipped, so the ready card would only restate + # the link card and print a proxy URL that 404s outside this tab. + skip_ready_card = _is_colab_runtime() and bool(cloudflare_url) + if not skip_ready_card: + try: + show_link( + port, + _url = url, + has_cloudflare_link = bool(cloudflare_url), + cloudflare_requested = cloudflare_requested, + ) + except Exception as e: + logger.info(f"Could not render Unsloth link card ({e}).") # On Colab with a working tunnel, skip the in-cell proxy embed (often blank). if _is_colab_runtime() and cloudflare_url: diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 1b24c46e65..e364ea4f3a 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -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.""" diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 147174451e..397cba8842 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -33,6 +33,7 @@ from typing import ( List, Literal, Mapping, + MutableMapping, Optional, Union, ) @@ -306,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( @@ -2098,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 @@ -3182,7 +3189,7 @@ class LlamaCppBackend: @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: - """True only for AMD unified-memory APUs (gfx1150/gfx1151), where + """True only for AMD unified-memory APUs (gfx1150/gfx1151/gfx1152), where GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the selected GPUs, so a dGPU on a mixed host is not treated as unified-memory; @@ -3212,7 +3219,9 @@ class LlamaCppBackend: ) arch_by_id[pid] = _arch.split(":")[0].strip().lower() for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id): - if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}: + # gfx1152 is Krackan Point (Radeon 860M/840M), the third RDNA 3.5 + # APU: same shared GPU/system-RAM pool as Strix Point/Halo. + if arch_by_id.get(_i) in {"gfx1150", "gfx1151", "gfx1152"}: return True except Exception: return False @@ -3696,6 +3705,14 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # V cache types that llama.cpp can run WITHOUT flash attention. Only the V + # axis has the dependency: a quantized V cache (q8_0/q4_0/q4_1/q5_0/q5_1/ + # iq4_nl) aborts init with "V cache quantization requires flash_attn", while + # a quantized K cache runs fine without FA. So the flash-attn-off crash- + # recovery fallback must reset a quantized V cache to f16 before it can + # launch (and leaves K alone). These three are the only non-quantized types. + _NON_QUANTIZED_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # Main-model placement settings that Manual mode owns. They must not leak # from Studio's parent environment into llama-server and silently override # the command assembled from the current request. Draft-model placement is @@ -6149,6 +6166,21 @@ class LlamaCppBackend: cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode) ) + @staticmethod + def _canonical_long_flag(name: str) -> str: + """Return ``name`` with llama.cpp's long-option underscore normalization. + + llama.cpp runs ``std::replace(arg.begin(), arg.end(), '_', '-')`` on any + argv token that starts with ``--`` before looking it up, so a legal + pass-through spelling like ``--cache_type_v`` parses as + ``--cache-type-v``. Mirror that here so managed-flag matching sees the + same canonical name. Short flags (``-ctv``) never carry underscores and + keep their exact spelling; pass only the flag name (no attached value). + """ + if name.startswith("--"): + return name.replace("_", "-") + return name + @staticmethod def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: """Return cmd with flash attention forced off, or None when its effective @@ -6181,8 +6213,76 @@ class LlamaCppBackend: out[i + 1] = "off" elif explicit(i) is None: # bare flag (reads as on) -> explicit off out[i] = f"{tok}=off" + + # A quantized V cache requires flash attention in llama.cpp: the init + # aborts with "V cache quantization requires flash_attn". A quantized K + # cache has no such requirement and runs fine without FA, so it is left + # untouched -- resetting it would needlessly enlarge the K cache and can + # OOM a memory-constrained config. Studio launches with FA on, so a + # quantized --cache-type-v is legal at launch but would make THIS FA-off + # retry crash on init instead of recovering. Reset a quantized V cache -- + # main and draft (the draft context shares the global --flash-attn flag, + # so its V cache aborts too) -- to f16 (the llama.cpp default); + # non-quantized types -- f16/bf16/f32 -- run fine without FA and are left + # untouched. The value is rewritten in place so the list length is + # preserved for downstream slices, matching the flash-attn flip above. + _v_cache_flags = ( + "--cache-type-v", + "-ctv", + "--cache-type-v-draft", + "--spec-draft-type-v", + "-ctvd", + ) + _cache_reset = False + for i, tok in enumerate(out): + # llama.cpp rewrites '_' to '-' for any argv token starting with + # '--' before matching, so a legal pass-through spelling such as + # --cache_type_v parses as --cache-type-v and still enables a + # quantized V cache. Canonicalize the flag name the same way so the + # reset recognizes the underscore aliases too; short flags (-ctv) + # and the type value are left untouched. + name = LlamaCppBackend._canonical_long_flag(tok.partition("=")[0]) + if name not in _v_cache_flags: + continue + if "=" in tok: + flag, _, value = tok.partition("=") + if value.strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i] = f"{flag}=f16" + _cache_reset = True + elif i + 1 < len(out): + if out[i + 1].strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i + 1] = "f16" + _cache_reset = True + if _cache_reset: + logger.info( + "V cache dtype reset to f16 because flash attention was disabled " + "by the crash-recovery fallback (quantized V cache requires flash " + "attention in llama.cpp; the K cache is left untouched)." + ) return out + @staticmethod + def _drop_env_quantized_v_cache(env: MutableMapping[str, str]) -> bool: + """Drop an inherited quantized V-cache env var (main or draft) in place + before a flash-attn-off retry, returning True if anything was removed. + + The argv rewrite in ``_with_flash_attn_off`` only reaches flags on the + command line. Studio deliberately lets an env-only cache type reach the + child untouched (an asymmetric K/V env must survive), so a quantized V + cache set purely through ``LLAMA_ARG_CACHE_TYPE_V`` (or the draft + ``LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V``) would still abort the FA-off retry + with "V cache quantization requires flash_attn". Dropping it lets + llama.cpp fall back to the f16 default. Only V is dropped: a quantized K + cache runs fine without flash attention, so its env var is preserved. + """ + dropped = False + for var in ("LLAMA_ARG_CACHE_TYPE_V", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V"): + value = (env.get(var) or "").strip().lower() + if value and value not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + env.pop(var, None) + dropped = True + return dropped + @staticmethod def _strip_mmproj_args(cmd: list[str]) -> list[str]: """Return cmd without the '--mmproj ' pair (text-only retry). @@ -8339,6 +8439,13 @@ class LlamaCppBackend: _fa_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = _spawn_and_wait(_fa_cmd, label = "-noflash") @@ -8384,6 +8491,13 @@ class LlamaCppBackend: _probe_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = ( _spawn_and_wait(_fa_cmd, label = "-noflash-mtp") @@ -9200,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 @@ -9999,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: @@ -10055,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: @@ -10527,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. @@ -10536,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, @@ -10855,7 +11081,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 @@ -11115,7 +11340,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, ): @@ -12152,7 +12377,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, ): diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 0ef6dd46cf..b31db3faf6 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -49,6 +49,7 @@ from loggers import get_logger logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes +_DISABLE_DNS_PINNING_ENV = "UNSLOTH_STUDIO_DISABLE_DNS_PINNING" # Splits the UI source-map from the result; loops strip it (like __IMAGES__). RAG_SOURCES_SENTINEL = "\n__RAG_SOURCES__:" @@ -4194,13 +4195,18 @@ def _fetch_url_raw( budget_error = _fetch_budget_exceeded(deadline, cancel_event) if budget_error is not None: return budget_error, "", "" - # Pin to the validated IP (prevents DNS rebinding): rewrite URL to - # the IP, set the Host header. cp = urlparse(current_url) - # Bracket IPv6 addresses so the netloc is valid in a URL. - ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip - ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str - pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + validated_netloc = f"[{current_host}]" if ":" in current_host else current_host + if cp.port: + validated_netloc = f"{validated_netloc}:{cp.port}" + if os.environ.get(_DISABLE_DNS_PINNING_ENV) == "1": + # Enterprise proxies need the hostname in CONNECT for policy and TLS interception. + request_url = urlunparse(cp._replace(netloc = validated_netloc)) + else: + # Pin to the validated IP to prevent DNS rebinding. + ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip + ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str + request_url = urlunparse(cp._replace(netloc = ip_netloc)) opener = urllib.request.build_opener( _NoRedirect, @@ -4209,11 +4215,11 @@ def _fetch_url_raw( headers = { "User-Agent": ua, - "Host": current_host, + "Host": validated_netloc, } if extra_headers: headers.update(extra_headers) - req = urllib.request.Request(pinned_url, headers = headers) + req = urllib.request.Request(request_url, headers = headers) try: # Cap the socket timeout at the time left on the overall deadline # so a single slow hop cannot outlast the whole fetch budget. diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 03327d3320..baf6329dae 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -764,8 +764,8 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: - ``gcn_arch``: canonical arch string (e.g. ``"gfx1151"``) when a known attribute is present, else ``""``. - ``is_unified``: ``True`` for AMD APUs with a shared GPU/system-RAM pool - (gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower - ``set_per_process_memory_fraction`` cap to leave OS headroom. + (gfx1150 Strix Point, gfx1151 Strix Halo, gfx1152 Krackan Point) — these + need a lower ``set_per_process_memory_fraction`` cap to leave OS headroom. Classification priority: 1. ``props.is_integrated`` truthy (hipDeviceProp_t.integrated -- the @@ -778,6 +778,7 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: - gfx1151 Strix Halo / Gorgon Halo: ``Radeon 8065S`` (Ryzen AI Max+ 495), ``Radeon 8060S`` (Ryzen AI MAX+ 395), ``Radeon 8050S`` (cut-down SKU) + - gfx1152 Krackan Point: ``Radeon 860M``, ``Radeon 840M`` """ gcn_arch = "" for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"): @@ -797,9 +798,13 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: return gcn_arch, True if gcn_arch: - return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"} + # gfx1152 is Krackan Point, the third RDNA 3.5 APU: same shared + # GPU/system-RAM pool as Strix Point (gfx1150) and Strix Halo (gfx1151). + return gcn_arch, gcn_arch in {"gfx1150", "gfx1151", "gfx1152"} - # Arch attrs absent — fall back to device-name matching. + # Arch attrs absent — fall back to device-name matching. Only reached under + # _hw.IS_ROCM, so the NVIDIA GeForce 840M cannot collide with the Krackan + # markers here. dev_lower = (getattr(props, "name", "") or "").lower() is_unified = ( "890m" in dev_lower @@ -807,6 +812,8 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: or "8065s" in dev_lower or "8060s" in dev_lower or "8050s" in dev_lower + or "860m" in dev_lower + or "840m" in dev_lower ) return gcn_arch, is_unified @@ -2828,7 +2835,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # On ROCm, exhausting VRAM can hang the HIP driver instead of raising. # set_per_process_memory_fraction caps the allocator so PyTorch raises # OutOfMemoryError first (NVIDIA already has a graceful OOM path). - # Unified-memory APUs (gfx1150/gfx1151) share GPU+system RAM, so use 0.80 + # Unified-memory APUs (gfx1150/gfx1151/gfx1152) share GPU+system RAM, so use 0.80 # vs 0.90 for discrete. Classify via gcnArchName, else device-name markers. # Non-fatal: skipped if torch is not importable. if _hw.IS_ROCM: diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 392a4e0d02..ae65712146 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -47,7 +47,14 @@ _SAFE = re.compile(r"[^A-Za-z0-9._-]+") def _sanitize_filename(name: str) -> str: base = os.path.basename(name or "").strip() or "document" base = _SAFE.sub("_", base) - return base[:200] + if len(base) <= 200: + return base + # Trim the stem, not the extension: _save_upload gates on the extension, so + # a plain truncation would reject a long-named .txt as "unsupported". + stem, ext = os.path.splitext(base) + if not ext or len(ext) > 32: + return base[:200] + return stem[: 200 - len(ext)] + ext def _save_upload(file: UploadFile) -> tuple[str, str]: diff --git a/studio/backend/run.py b/studio/backend/run.py index d9569c46f6..f1fc8c6062 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -991,10 +991,88 @@ class _TeeStream: except Exception: pass + def close(self): + # We do NOT own the console stream (it is the terminal / Jupyter kernel + # stream we wrapped), so closing the tee must never take the server down. + # Flush the log copy, then forward close() to the wrapped stream + # best-effort: on Colab that stream is an ipykernel OutStream whose + # close() can raise (see _harden_console_close / ipython/ipykernel#867). + try: + self._log_fh.flush() + except Exception: + pass + try: + self._stream.close() + except Exception: + pass + def __getattr__(self, name): return getattr(self._stream, name) +_WATCH_FD_THREAD_ATTR = "watch_fd_thread" + + +def _is_missing_watch_fd_thread(exc): + """True only for ipython/ipykernel#867's missing-``watch_fd_thread`` error. + + ``AttributeError.name`` exists from Python 3.10; the message carries the + attribute name on every version (possibly with a "Did you mean" tail), so + check both and let every other AttributeError through. + """ + if getattr(exc, "name", None) == _WATCH_FD_THREAD_ATTR: + return True + return _WATCH_FD_THREAD_ATTR in str(exc) + + +def _harden_console_close(stream): + """Stop a displaced console stream's close() from aborting Studio startup. + + ``_setup_server_disk_logging`` replaces ``sys.stdout``/``sys.stderr`` with a + tee. That changes the object identity of the console stream, so a third-party + logging handler that captured the ORIGINAL stream (notably Colab's ``absl`` + logging handler, whose ``close()`` skips ``sys.stdout``/``sys.stderr`` but not + a stream that is no longer either) treats it as an ordinary stream and calls + ``close()`` on it during logging teardown -- ``uvicorn.Config()`` -> + ``logging.config.dictConfig()`` -> ``logging.shutdown()``. + + A Jupyter/Colab ``ipykernel`` ``OutStream`` created with ``watchfd=False`` + (the Colab default, and every in-process kernel) never gains a + ``watch_fd_thread``, yet the ``OutStream.close()`` shipped in the affected + ipykernel versions joins that thread unconditionally and raises + ``AttributeError: 'OutStream' object has no attribute 'watch_fd_thread'`` + (ipython/ipykernel#867). That AttributeError propagates out of + ``uvicorn.Config(...)`` and aborts startup ("Unsloth Studio failed to start"). + + Wrap the stream's ``close()`` in a transparent pass-through that swallows + ONLY that specific teardown AttributeError. A healthy close() (a real console + stream, or an OutStream with fd-watching on) runs to completion exactly as + before and any other error still propagates, so nothing changes off Colab. A + stream whose ``close`` cannot be reassigned keeps its original close(). + """ + try: + _orig_close = stream.close + except Exception: + return + + def _safe_close(*args, **kwargs): + try: + return _orig_close(*args, **kwargs) + except AttributeError as exc: + if not _is_missing_watch_fd_thread(exc): + # A real teardown failure; never hide it. + raise + # ipython/ipykernel#867: watchfd=False OutStream.close() joins a + # thread that was never created. Nothing to clean up; keep going. + return None + + try: + stream.close = _safe_close + except (AttributeError, TypeError): + # A stream that forbids setting instance attributes; leave it as-is. + pass + + def _setup_server_disk_logging(): """Tee stdout/stderr to ~/.unsloth/studio/logs/server/ and aim faulthandler at the same file so hard crashes (access violations / @@ -1037,6 +1115,11 @@ def _setup_server_disk_logging(): # the stderr the server already captures. os.environ.setdefault("PYTHONFAULTHANDLER", "1") + # Replacing the console streams orphans them from third-party "is this the + # live console?" checks, so guard their close() first (ipython/ipykernel#867). + _harden_console_close(sys.stdout) + _harden_console_close(sys.stderr) + sys.stdout = _TeeStream(sys.stdout, log_fh) sys.stderr = _TeeStream(sys.stderr, log_fh) @@ -1802,6 +1885,12 @@ def _build_arg_parser(): default = None, help = "Force server-side tools off for every request.", ) + parser.add_argument( + "--disable-dns-pinning", + action = "store_true", + help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens " + "DNS-rebinding protection; hostname and redirect validation remain enabled.", + ) parser.add_argument( "--parallel", "--n-parallel", @@ -1841,6 +1930,10 @@ if __name__ == "__main__": parser.error( "--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare" ) + if args.disable_dns_pinning: + os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1" + else: + os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0") kwargs = dict( host = args.host, diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py index 9fd8260bf2..be85fd56d1 100644 --- a/studio/backend/tests/test_amd_apu_unified_memory.py +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -2,7 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs -(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS.""" +(gfx1150/gfx1151/gfx1152), never for discrete AMD, NVIDIA, CPU or macOS.""" from __future__ import annotations @@ -35,6 +35,8 @@ def _fake_torch( [ ("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped) ("6.2.0", ["gfx1150"], True), # Strix Point APU + ("6.2.0", ["gfx1152"], True), # Krackan Point APU (Radeon 860M/840M) + ("6.2.0", ["gfx1152:sramecc-:xnack-"], True), # same, feature flags stripped ("6.2.0", ["gfx1100"], False), # discrete RDNA3 ("6.2.0", ["gfx1201"], False), # discrete RDNA4 ("6.2.0", ["gfx942"], False), # MI300X (data center) diff --git a/studio/backend/tests/test_colab_embed.py b/studio/backend/tests/test_colab_embed.py index dae0c7dae0..83b2a5a82d 100644 --- a/studio/backend/tests/test_colab_embed.py +++ b/studio/backend/tests/test_colab_embed.py @@ -336,17 +336,71 @@ def test_colab_login_html_includes_credentials(): html = colab._colab_login_html("unsloth", "alpha-beta-gamma-delta") assert "unsloth" in html assert "alpha-beta-gamma-delta" in html + # The username is fixed, so it reads inline rather than as its own field. + assert "Username:" not in html -def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch): +def test_shareable_link_html_embeds_password_under_the_link(): + """The credential belongs in the same card as the button it unlocks.""" + html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth") + assert "share.trycloudflare.com" in html + assert "secret-pass" in html + # Username is stated inline, not as its own labelled field. + assert "Username:" not in html + assert "unsloth" in html + # The password must sit after the link, not above it. + assert html.index("share.trycloudflare.com") < html.index("secret-pass") + + +def test_shareable_link_html_renders_the_url_as_a_link(): + """The printed URL is an anchor, using the popup-safe open the button uses.""" + html = colab._shareable_link_html("https://share.trycloudflare.com") + assert 'https://share.trycloudflare.com" in html + assert html.count("window.open(this.href,'_blank')") == 2 + + +def test_shareable_link_html_emphasises_the_password(): + """The password is the one thing to copy, so it is enlarged and underlined.""" + html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth") + pw_tag = html[html.index("Password") : html.index("secret-pass")] + assert "font-size: 24px" in pw_tag + assert "text-decoration: underline" in pw_tag + + +def test_shareable_link_html_password_has_no_adjacent_whitespace(): + """Whitespace beside the password is selected with it on a double click.""" + html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth") + before, after = html.split("secret-pass", 1) + assert before.endswith(">") + assert after.startswith("<") + # Label on its own line, so nothing shares the password's text node. + assert "Password:" not in html + # Plain selectable text: user-select overrides break double click to select. + assert "user-select" not in html + + +def test_shareable_link_html_omits_login_block_without_password(): + html = colab._shareable_link_html("https://share.trycloudflare.com") + assert "Password" not in html + + +def test_show_and_embed_folds_login_into_the_cloudflare_card(monkeypatch): + """One card, not two: the tunnel card carries the password itself.""" displayed: list[str] = [] ipython_display = SimpleNamespace( HTML = lambda html: SimpleNamespace(html = html), display = lambda html: displayed.append(html.html), ) + login_cards: list[tuple] = [] monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_show_colab_login_credentials", + lambda *args: login_cards.append(args), + ) monkeypatch.setattr( colab, "show_link", @@ -360,9 +414,74 @@ def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch): colab_login = ("unsloth", "secret-pass"), ) - assert len(displayed) == 2 + assert len(displayed) == 1 assert "share.trycloudflare.com" in displayed[0] - assert "secret-pass" in displayed[1] + assert "secret-pass" in displayed[0] + assert login_cards == [] + + +def test_show_and_embed_keeps_separate_login_card_without_tunnel(monkeypatch): + """No tunnel card to fold into, so the standalone login card still renders.""" + login_cards: list[tuple] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_show_colab_login_credentials", + lambda *args: login_cards.append(args), + ) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + colab._show_and_embed(8888, colab_login = ("unsloth", "secret-pass")) + + assert login_cards == [("unsloth", "secret-pass")] + + +def test_show_and_embed_skips_ready_card_when_tunnel_is_up(monkeypatch): + """The ready card only restates the tunnel card and prints a proxy URL that 404s.""" + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + colab._show_and_embed(8888, cloudflare_url = "https://share.trycloudflare.com") + + assert calls == [] + + +def test_show_and_embed_keeps_ready_card_without_tunnel(monkeypatch): + """Without a tunnel the ready card is the only guidance, so it must stay.""" + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + colab._show_and_embed(8888) + + assert calls == ["show_link"] def test_show_and_embed_skips_iframe_on_colab_when_cloudflare_ready(monkeypatch): diff --git a/studio/backend/tests/test_export_multi_gpu_device_map.py b/studio/backend/tests/test_export_multi_gpu_device_map.py new file mode 100644 index 0000000000..e483fbe728 --- /dev/null +++ b/studio/backend/tests/test_export_multi_gpu_device_map.py @@ -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 diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 6d1fac980b..0ab998af39 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -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 @@ -785,13 +810,21 @@ class TestLoadHubDownloadExclusion: def test_load_marker_precedes_hub_guard_and_unload(self): source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() - gguf_branch = source[source.index("if config.is_gguf:") :] + # _load_model_impl has more than one `if config.is_gguf:`, so anchor on + # the branch that actually owns the load marker rather than the first + # one in the file, which belongs to an earlier check. + marker = source.index("enter_context(gguf_load_in_flight") + gguf_branch_start = source.rindex("if config.is_gguf:", 0, marker) + gguf_branch = source[gguf_branch_start:] # The gguf_load_in_flight marker must be entered before the hub-download # guard and the unload so a concurrent load can't race the download - # manager. The llama_extra_args inheritance that used to sit between the - # marker and the guard now runs in _guard_chat_load_against_training, ahead - # of the GGUF branch, so it is no longer a landmark inside this slice. + # manager. The llama_extra_args inheritance moved out of the branch into + # _resolve_inherited_extra_args, which must still run BEFORE it: the + # inherited value (e.g. a carried --no-mmproj) shapes the guard's + # require_mmproj. Anchor on the call form so the assertion pins the + # endpoint's call site, not the function definition. + assert source.index("= _resolve_inherited_extra_args(") < gguf_branch_start assert ( gguf_branch.index("enter_context(gguf_load_in_flight") < gguf_branch.index("_hub_download_blocks_gguf_load") @@ -801,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 + ) diff --git a/studio/backend/tests/test_grouped_mm_rdna4_fallback.py b/studio/backend/tests/test_grouped_mm_rdna4_fallback.py new file mode 100644 index 0000000000..675b9c3210 --- /dev/null +++ b/studio/backend/tests/test_grouped_mm_rdna4_fallback.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Numerics + gating for the RDNA4 _grouped_mm CPU fallback (PRs #7276 / #7292). + +RDNA4 (gfx1200/gfx1201) ships a null HIP `_grouped_mm` kernel on ROCm <= 7.12 +(fixed in 7.13, ROCm/TheRock #5284). Training MoE models there crashes with +0xC0000005 on Windows and a plain segfault on Linux, so worker.py registers a +Python mm/bmm fallback on the CUDA dispatch key. + +The fallback is silent, GPU-gated, and reimplements a matmul: if it is wrong, an +RX 9070 user does not crash, they train on quietly wrong gradients. Until now the +only coverage was `assert '_gm_lib.impl("_grouped_mm"' in source` -- the math was +never executed once, in any suite. + +worker.py cannot be imported here (module-level structlog/backend imports), so +`_install_grouped_mm_cpu_fallback` is lifted out with ast and driven with a fake +`torch_mod` that forwards to real CPU torch. That also pins the op surface: the +fallback may only use the ops the fake exposes, and the registration is captured +instead of hitting a real CUDA dispatch key that CI runners do not have. + +The two gates around it are exec'd straight out of the source so this file tests +the shipped expressions rather than a copy of them. +""" + +import ast +import re +import textwrap +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +_WORKER_PATH = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" +_WORKER_SOURCE = _WORKER_PATH.read_text(encoding = "utf-8") + + +def _load_installer(): + """exec just _install_grouped_mm_cpu_fallback out of worker.py.""" + tree = ast.parse(_WORKER_SOURCE) + fn = [ + n + for n in tree.body + if isinstance(n, ast.FunctionDef) and n.name == "_install_grouped_mm_cpu_fallback" + ] + assert fn, "_install_grouped_mm_cpu_fallback not found in core/training/worker.py" + ns: dict = {} + exec(compile(ast.Module(body = fn, type_ignores = []), str(_WORKER_PATH), "exec"), ns) + return ns["_install_grouped_mm_cpu_fallback"] + + +_install_grouped_mm_cpu_fallback = _load_installer() + + +class _RecordingLibrary: + """Stands in for torch.library.Library: captures the registration instead of + binding it to a CUDA dispatch key no CI runner has.""" + + def __init__(self, namespace, kind): + self.namespace = namespace + self.kind = kind + self.registrations = [] + + def impl(self, name, fn, dispatch_key): + self.registrations.append((name, fn, dispatch_key)) + + +class _RecordingLogger: + def __init__(self): + self.info_calls = [] + self.warning_calls = [] + + def info(self, *args, **kwargs): + self.info_calls.append(args) + + def warning(self, *args, **kwargs): + self.warning_calls.append(args) + + +def _fake_torch(): + """Real CPU torch behind the exact op surface the fallback is allowed to use. + + Anything else the fallback reaches for raises AttributeError here, which is + the point: a new dependency has to be a deliberate edit, not a silent one.""" + return SimpleNamespace( + library = SimpleNamespace(Library = _RecordingLibrary), + mm = torch.mm, + bmm = torch.bmm, + matmul = torch.matmul, + cat = torch.cat, + zeros = torch.zeros, + ) + + +@pytest.fixture +def fallback(): + """The registered _grouped_mm implementation, plus the Library it landed on.""" + torch_mod = _fake_torch() + logger = _RecordingLogger() + lib = _install_grouped_mm_cpu_fallback(torch_mod, logger, "test") + assert lib.registrations, "the fallback registered nothing" + name, fn, key = lib.registrations[0] + return SimpleNamespace(fn = fn, lib = lib, logger = logger, name = name, key = key) + + +class TestRegistration: + """Where the override lands. Getting the namespace or dispatch key wrong is a + silent no-op: training still crashes on the null HIP kernel.""" + + def test_overrides_aten_grouped_mm_on_the_cuda_key(self, fallback): + assert fallback.lib.namespace == "aten" + assert fallback.lib.kind == "IMPL" + assert fallback.name == "_grouped_mm" + # ROCm dispatches through the CUDA key; "HIP"/"PrivateUse1" would not bind. + assert fallback.key == "CUDA" + + def test_registers_exactly_once(self, fallback): + assert len(fallback.lib.registrations) == 1 + + def test_returns_the_library_so_the_caller_can_keep_it_alive(self, fallback): + """A dropped Library is garbage collected and the override silently + unregisters mid-run; worker.py parks it in a module global.""" + assert isinstance(fallback.lib, _RecordingLibrary) + assert "_WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback(" in _WORKER_SOURCE + + def test_logs_the_patch_with_its_label(self, fallback): + assert fallback.logger.info_calls, "the patch must be visible in the run log" + assert "test" in fallback.logger.info_calls[0] + + +class TestUngroupedNumerics: + """offs=None: plain matmul, one path per rank combination. The 3-D case is + the regression #7292 fixed -- an unconditional mm() broke MoE experts.""" + + def test_2d_by_2d_matches_mm(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.mm(a, b)) + + def test_3d_by_3d_matches_bmm(self, fallback): + a = torch.randn(3, 6, 4) + b = torch.randn(3, 4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.bmm(a, b)) + + def test_3d_by_2d_matches_matmul(self, fallback): + a = torch.randn(3, 6, 4) + b = torch.randn(4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.matmul(a, b)) + + def test_2d_by_3d_matches_matmul(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(3, 4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.matmul(a, b)) + + def test_non_contiguous_inputs_are_handled(self, fallback): + """Transposed views reach _grouped_mm constantly; every path calls + .contiguous() and this catches it if one stops.""" + a = torch.randn(4, 6).t() + b = torch.randn(5, 4).t() + torch.testing.assert_close(fallback.fn(a, b), torch.mm(a, b)) + + +class TestGroupedNumerics: + """offs=[end-row of each group], the MoE token-routing layout.""" + + def test_matches_per_group_mm_with_3d_weights(self, fallback): + a = torch.randn(7, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([2, 5, 7]) + expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[1], a[5:7] @ b[2]], dim = 0) + torch.testing.assert_close(fallback.fn(a, b, offs), expected) + + def test_shared_2d_weight_is_reused_for_every_group(self, fallback): + a = torch.randn(7, 4) + b = torch.randn(4, 5) + offs = torch.tensor([2, 5, 7]) + torch.testing.assert_close(fallback.fn(a, b, offs), a @ b) + + def test_empty_group_produces_no_rows(self, fallback): + """An expert that routed zero tokens (offs[i] == offs[i-1]) must + contribute nothing, not a stray row.""" + a = torch.randn(5, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([2, 2, 5]) + expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[2]], dim = 0) + got = fallback.fn(a, b, offs) + assert got.shape == (5, 5) + torch.testing.assert_close(got, expected) + + def test_rows_past_the_last_offset_are_not_dropped(self, fallback): + """Trailing tokens beyond offs[-1] go through the last expert; dropping + them would silently shrink the output instead of raising.""" + a = torch.randn(7, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([2, 5]) + expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[1], a[5:7] @ b[-1]], dim = 0) + got = fallback.fn(a, b, offs) + assert got.shape[0] == a.shape[0] + torch.testing.assert_close(got, expected) + + def test_zero_rows_returns_an_empty_result_not_an_error(self, fallback): + a = torch.randn(0, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([], dtype = torch.int64) + got = fallback.fn(a, b, offs) + assert got.shape == (0, 5) + assert got.dtype == a.dtype + + def test_offsets_may_arrive_as_a_device_tensor_of_any_int_dtype(self, fallback): + a = torch.randn(4, 4) + b = torch.randn(2, 4, 5) + expected = torch.cat([a[0:2] @ b[0], a[2:4] @ b[1]], dim = 0) + for dtype in (torch.int32, torch.int64): + torch.testing.assert_close( + fallback.fn(a, b, torch.tensor([2, 4], dtype = dtype)), expected + ) + + +class TestBiasAndDtype: + def test_bias_is_added(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(4, 5) + bias = torch.randn(5) + torch.testing.assert_close(fallback.fn(a, b, None, bias), torch.mm(a, b) + bias) + + def test_bias_is_added_on_the_grouped_path_too(self, fallback): + a = torch.randn(4, 4) + b = torch.randn(2, 4, 5) + bias = torch.randn(5) + offs = torch.tensor([2, 4]) + expected = torch.cat([a[0:2] @ b[0], a[2:4] @ b[1]], dim = 0) + bias + torch.testing.assert_close(fallback.fn(a, b, offs, bias), expected) + + def test_out_dtype_is_honoured(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(4, 5) + got = fallback.fn(a, b, None, None, torch.float64) + assert got.dtype == torch.float64 + torch.testing.assert_close(got, torch.mm(a, b).to(torch.float64)) + + def test_promotion_from_bias_is_cast_back_to_the_input_dtype(self, fallback): + """Without the restore, a promoted result changes the autograd dtype + downstream of every MoE layer.""" + a = torch.randn(6, 4, dtype = torch.float32) + b = torch.randn(4, 5, dtype = torch.float32) + bias = torch.randn(5, dtype = torch.float64) + got = fallback.fn(a, b, None, bias) + assert got.dtype == torch.float32 + + def test_out_dtype_wins_over_the_input_dtype_restore(self, fallback): + a = torch.randn(6, 4, dtype = torch.float32) + b = torch.randn(4, 5, dtype = torch.float32) + bias = torch.randn(5, dtype = torch.float64) + got = fallback.fn(a, b, None, bias, torch.float64) + assert got.dtype == torch.float64 + + def test_bf16_inputs_stay_bf16(self, fallback): + """The dtype training actually runs in.""" + a = torch.randn(6, 4).to(torch.bfloat16) + b = torch.randn(4, 5).to(torch.bfloat16) + got = fallback.fn(a, b) + assert got.dtype == torch.bfloat16 + torch.testing.assert_close(got.float(), (a.float() @ b.float()), rtol = 2e-2, atol = 2e-2) + + +def _exec_source_snippet(anchor: str, last_line: str, **variables): + """Run a slice of worker.py verbatim, so the gate under test is the shipped + one and not a copy that can drift.""" + start = _WORKER_SOURCE.find(anchor) + assert start != -1, f"gate snippet not found in worker.py: {anchor!r}" + start = _WORKER_SOURCE.rfind("\n", 0, start) + 1 # keep the indent for dedent() + end = _WORKER_SOURCE.find(last_line, start) + assert end != -1, f"end of gate snippet not found: {last_line!r}" + snippet = textwrap.dedent(_WORKER_SOURCE[start : end + len(last_line)]) + ns = {"re": re, **variables} + exec(compile(snippet, str(_WORKER_PATH), "exec"), ns) + return ns + + +class TestLinuxHipVersionGate: + """PR #7292's Linux gate. Too low a floor keeps the slow Python fallback on + fixed ROCm 7.13+; too high reintroduces the segfault on 7.12.""" + + _ANCHOR = '_m = re.match(r"(\\d+)\\.(\\d+)", _hip_str)' + _LAST = '_hip_lt_713 = "rocmsdk" not in _ver' + + def _decide(self, hip_str, version): + ns = _exec_source_snippet(self._ANCHOR, self._LAST, _hip_str = hip_str, _ver = version.lower()) + return ns["_hip_lt_713"] + + @pytest.mark.parametrize( + "hip_str,version,affected", + [ + ("7.12.0", "2.10.0+rocm7.12.0", True), # the broken kernel + ("7.6.0", "2.9.0+rocm7.6.0", True), + ("6.4.0", "2.8.0+rocm6.4.0", True), + ("7.13.0", "2.11.0+rocm7.13.0", False), # AMD's fix + ("7.14.0", "2.11.0+rocm7.14.0", False), + ("8.0.0", "2.12.0+rocm8.0.0", False), + ], + ) + def test_torch_version_hip_decides_when_present(self, hip_str, version, affected): + assert self._decide(hip_str, version) is affected + + @pytest.mark.parametrize( + "version,affected", + [ + ("2.10.0+rocm7.12.0", True), + ("2.11.0+rocm7.13.0", False), + ("2.11.0+rocm7.14.0", False), + ], + ) + def test_falls_back_to_the_rocm_tag_in_torch_version(self, version, affected): + """AMD SDK / Radeon wheels leave torch.version.hip unset.""" + assert self._decide("", version) is affected + + def test_unknown_version_is_assumed_affected(self): + """Fallback is slow but correct; a missed guard is a crash.""" + assert self._decide("", "2.9.0+unknown") is True + + def test_rocmsdk_wheels_without_a_version_are_assumed_fixed(self): + """rocmsdk wheels post-date the gfx120X fix.""" + assert self._decide("", "2.10.0+rocmsdk20260107") is False + + +class TestLinuxRdna4NameMatch: + """The name regex is the fallback when a wheel omits gcnArchName.""" + + def _pattern(self): + """Read whatever pattern worker.py currently uses, not a copy of the one + it used when this test was written. Anchoring on the literal pattern text + would make a *widened* regex -- the dangerous edit, since it silently + forces the slow Python fallback onto RDNA3 users -- fail as "moved" + instead of being checked against the cases below.""" + m = re.search(r"re\.search\(r\"([^\"]+)\",\s*_lin_name\)", _WORKER_SOURCE) + assert m, "could not locate the RDNA4 device-name regex in worker.py" + return m.group(1) + + def test_name_is_lowercased_before_matching(self): + """The pattern is all-lowercase, so it only works against a lowercased + name. Device names arrive mixed case ("AMD Radeon RX 9070 XT").""" + assert self._pattern() == self._pattern().lower(), "pattern is not all-lowercase" + assert re.search( + r"_lin_name\s*=\s*\(getattr\(_props,\s*\"name\",\s*\"\"\)\s*or\s*\"\"\)\.lower\(\)", + _WORKER_SOURCE, + ), "worker.py must lowercase the device name before matching the RDNA4 pattern" + + def test_name_match_is_only_a_fallback_when_arch_is_unknown(self): + """gcnArchName is authoritative when present. Letting the name regex fire + alongside a known arch would misclassify any card whose marketing name + happens to look RDNA4.""" + assert re.search( + r"not _lin_arch and re\.search\(r\"[^\"]+\",\s*_lin_name\)", _WORKER_SOURCE + ), "the RDNA4 name regex must be guarded by `not _lin_arch`" + + @pytest.mark.parametrize( + "name,is_rdna4", + [ + ("AMD Radeon RX 9070 XT", True), + ("AMD Radeon RX 9060 XT", True), + ("Radeon RX9070", True), + ("AMD Radeon AI PRO R9700", True), + ("AMD Radeon RX 7900 XTX", False), # RDNA3, kernel is fine + ("AMD Radeon 8060S Graphics", False), # Strix Halo + ("AMD Radeon RX 6800 XT", False), + ("NVIDIA GeForce RTX 4090", False), + ], + ) + def test_matches_only_rdna4_cards(self, name, is_rdna4): + assert bool(re.search(self._pattern(), name.lower())) is is_rdna4 + + +class TestLinuxGateStructure: + """The block is a few hundred lines into run_training_process and can only be + checked structurally; these pin the parts a refactor would quietly drop.""" + + def _linux_block(self): + start = _WORKER_SOURCE.find("1f-linux") + assert start != -1, "the Linux ROCm gfx120X guard (#7292) is gone from worker.py" + end = _WORKER_SOURCE.find("1g.", start) + assert end != -1 + return _WORKER_SOURCE[start:end] + + def test_gated_on_linux_and_rocm(self): + block = self._linux_block() + assert 'sys.platform.startswith("linux")' in block + assert "_hw.IS_ROCM" in block, "guard must not run on NVIDIA/CPU hosts" + + def test_requires_both_rdna4_and_an_affected_hip(self): + block = self._linux_block() + assert "if _rdna4 and _hip_lt_713:" in block + + def test_scans_every_visible_device(self): + """device_map="balanced" can place layers on a later card, so checking + device 0 alone misses the RDNA4 GPU.""" + block = self._linux_block() + assert "for _i in range(_torch_lin.cuda.device_count()):" in block + + def test_matches_both_rdna4_arch_ids(self): + block = self._linux_block() + assert '("gfx1200", "gfx1201")' in block + + def test_failure_to_patch_is_non_fatal(self): + """A broken patch attempt must not take down the whole training run.""" + block = self._linux_block() + assert "except Exception" in block + assert "logger.warning" in block + + def test_windows_and_linux_share_one_implementation(self): + """Two copies of this fallback would drift; #7292 deliberately hoisted it.""" + assert _WORKER_SOURCE.count("def _install_grouped_mm_cpu_fallback(") == 1 + assert _WORKER_SOURCE.count("_install_grouped_mm_cpu_fallback(") >= 3 # def + win32 + linux + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 4332a440a5..45c8bcb032 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -250,6 +250,201 @@ class TestFlashAttnOff: assert _flash_off(["llama-server", "-fa"]) == ["llama-server", "-fa=off"] +_drop_env_v = LlamaCppBackend._drop_env_quantized_v_cache + + +class TestFlashAttnOffQuantizedKvCache: + """Only the V cache requires flash attention in llama.cpp (init aborts with + "V cache quantization requires flash_attn"); a quantized K cache runs fine + without FA. Studio launches FA on, so a quantized --cache-type-v is legal at + launch but would make the FA-off crash-recovery retry crash on init. The + fallback must reset a quantized V cache (main and draft) to f16 while leaving + the K cache and non-quantized (f16/bf16/f32) types unchanged -- resetting K + would needlessly enlarge it and can OOM a memory-constrained config.""" + + _QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"] + _NON_QUANTIZED = ["f16", "bf16", "f32"] + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_quantized_v_reset_k_preserved(self, qtype): + cmd = [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + qtype, + "--cache-type-v", + qtype, + ] + out = _flash_off(cmd) + assert out is not None + # FA flipped off AND the V axis reset to f16; the K axis is preserved so + # the FA-off retry keeps its memory budget (quantized K is FA-independent). + assert out[out.index("--flash-attn") + 1] == "off" + assert out[out.index("--cache-type-k") + 1] == qtype + assert out[out.index("--cache-type-v") + 1] == "f16" + assert len(out) == len(cmd) + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_quantized_draft_v_reset(self, qtype): + # The draft context shares the global --flash-attn flag, so its quantized + # V cache aborts too and must be reset; the draft K cache is preserved. + for v_flag, k_flag in ( + ("--cache-type-v-draft", "--cache-type-k-draft"), + ("--spec-draft-type-v", "--spec-draft-type-k"), + ("-ctvd", "-ctkd"), + ): + cmd = ["llama-server", "-fa", "on", k_flag, qtype, v_flag, qtype] + out = _flash_off(cmd) + assert out is not None + assert out[out.index(v_flag) + 1] == "f16" + assert out[out.index(k_flag) + 1] == qtype + + @pytest.mark.parametrize("ntype", _NON_QUANTIZED) + def test_nonquantized_cache_left_unchanged(self, ntype): + cmd = [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + ntype, + "--cache-type-v", + ntype, + ] + out = _flash_off(cmd) + assert out is not None + # Only FA flips; the non-quantized cache type is preserved verbatim. + assert out[out.index("--flash-attn") + 1] == "off" + assert out[out.index("--cache-type-k") + 1] == ntype + assert out[out.index("--cache-type-v") + 1] == ntype + + def test_equals_form_quantized_v_reset(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-v=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache-type-v=f16"] + + def test_equals_form_quantized_k_preserved(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-k=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache-type-k=q8_0"] + + def test_short_alias_v_reset_k_preserved(self): + out = _flash_off(["llama-server", "-fa", "on", "-ctk", "q4_0", "-ctv", "q4_0"]) + assert out == ["llama-server", "-fa", "off", "-ctk", "q4_0", "-ctv", "f16"] + + def test_asymmetric_cache_only_v_reset(self): + # Quantized V, non-quantized K: reset V, keep K untouched. + out = _flash_off( + [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + "f16", + "--cache-type-v", + "q8_0", + ] + ) + assert out[out.index("--cache-type-k") + 1] == "f16" + assert out[out.index("--cache-type-v") + 1] == "f16" + + def test_no_cache_flags_still_flips_fa(self): + out = _flash_off(["llama-server", "--flash-attn", "on", "-c", "4096"]) + assert out == ["llama-server", "--flash-attn", "off", "-c", "4096"] + + def test_quantized_k_only_still_flips_fa_but_keeps_k(self): + # A quantized K cache with no V flag is a valid FA-off launch; the retry + # must not touch the K cache (it would waste memory for nothing). + out = _flash_off(["llama-server", "--flash-attn", "on", "--cache-type-k", "q8_0"]) + assert out == ["llama-server", "--flash-attn", "off", "--cache-type-k", "q8_0"] + + def test_input_not_mutated(self): + cmd = ["llama-server", "--flash-attn", "on", "--cache-type-v", "q8_0"] + _flash_off(cmd) + assert cmd[-1] == "q8_0" + + @pytest.mark.parametrize( + "flag", + ["--cache_type_v", "--cache-type_v", "--cache_type-v"], + ) + def test_underscore_alias_v_reset(self, flag): + # llama.cpp normalizes '_' to '-' in any '--' long option before + # matching, so a pass-through --cache_type_v enables a quantized V cache + # and must be reset by the FA-off retry too (else init aborts). + out = _flash_off(["llama-server", "--flash-attn", "on", flag, "q8_0"]) + assert out is not None + assert out[out.index("--flash-attn") + 1] == "off" + # The user's flag spelling is preserved; llama.cpp normalizes it anyway. + assert out[out.index(flag) + 1] == "f16" + + def test_underscore_alias_draft_v_reset(self): + out = _flash_off(["llama-server", "-fa", "on", "--spec_draft_type_v", "q4_0"]) + assert out is not None + assert out[out.index("--spec_draft_type_v") + 1] == "f16" + + def test_underscore_alias_equals_form_v_reset(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache_type_v=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache_type_v=f16"] + + def test_underscore_value_not_normalized_for_nonquantized(self): + # Only the flag name is canonicalized; a non-quantized type value is + # matched verbatim and left untouched (no spurious reset). + out = _flash_off(["llama-server", "--flash-attn", "on", "--cache_type_v", "f16"]) + assert out[out.index("--cache_type_v") + 1] == "f16" + assert out[out.index("--flash-attn") + 1] == "off" + + def test_short_alias_underscore_not_applied(self): + # Short flags are never underscore-normalized by llama.cpp; -ctv still + # matches and resets, and an unrelated short token is left alone. + out = _flash_off(["llama-server", "-fa", "on", "-ctv", "q8_0"]) + assert out == ["llama-server", "-fa", "off", "-ctv", "f16"] + + +class TestDropEnvQuantizedVCache: + """The argv rewrite can't reach a cache type set purely through the + environment (Studio deliberately lets an env-only type reach the child), so + the FA-off retry separately drops a quantized V-cache env var. Only V is + dropped: a quantized K cache is FA-independent and must survive.""" + + _QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"] + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_drops_quantized_main_v_env(self, qtype): + env = {"LLAMA_ARG_CACHE_TYPE_V": qtype, "PATH": "/usr/bin"} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_CACHE_TYPE_V" not in env + assert env["PATH"] == "/usr/bin" + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_drops_quantized_draft_v_env(self, qtype): + env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V": qtype} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V" not in env + + def test_preserves_quantized_k_env(self): + # A quantized K cache runs without FA, so its env must not be dropped. + env = {"LLAMA_ARG_CACHE_TYPE_K": "q8_0", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q4_0"} + assert _drop_env_v(env) is False + assert env["LLAMA_ARG_CACHE_TYPE_K"] == "q8_0" + assert env["LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K"] == "q4_0" + + @pytest.mark.parametrize("ntype", ["f16", "bf16", "f32", "F16", " q8_0 "]) + def test_preserves_nonquantized_v_env(self, ntype): + # Non-quantized V env values (and whitespace/case variants of them) run + # fine without FA; only a genuinely quantized value is dropped. + if ntype.strip().lower() in ("q8_0",): + env = {"LLAMA_ARG_CACHE_TYPE_V": ntype} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_CACHE_TYPE_V" not in env + else: + env = {"LLAMA_ARG_CACHE_TYPE_V": ntype} + assert _drop_env_v(env) is False + assert env["LLAMA_ARG_CACHE_TYPE_V"] == ntype + + def test_noop_on_empty_env(self): + env = {} + assert _drop_env_v(env) is False + assert env == {} + + class TestNonProjectorDiagnostic: """_output_has_nonprojector_diagnostic gates the signal-only text-only retry: a hard crash that already names OOM / a bad arch / a TP limit must surface diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index cf9fde7118..92089d26ec 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -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") @@ -2239,7 +2268,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 +2305,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): diff --git a/studio/backend/tests/test_rag_project_source_upload.py b/studio/backend/tests/test_rag_project_source_upload.py new file mode 100644 index 0000000000..fd20816b56 --- /dev/null +++ b/studio/backend/tests/test_rag_project_source_upload.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Project sources upload: the path the create-project dialog drives.""" + +import os + +import pytest + +from core.rag import ingestion, store +from routes.rag import _sanitize_filename +from storage import rag_db + + +def _wait(job_id, timeout = 30.0): + import time + + deadline = time.time() + timeout + while time.time() < deadline: + status = ingestion.get_job_status(job_id) + if status and status["status"] in ("completed", "failed"): + return status + time.sleep(0.05) + raise AssertionError("ingestion did not finish in time") + + +def _ingest(project_id, filename, path): + return ingestion.start_ingestion( + store.project_scope(project_id), None, None, filename, path, project_id = project_id + ) + + +def test_project_document_persists_under_its_scope(rag_home, stub_embeddings, tmp_path): + path = tmp_path / "notes.txt" + path.write_text("alpha bravo charlie " * 50, encoding = "utf-8") + _, job_id = _ingest("P1", "notes.txt", str(path)) + assert _wait(job_id)["status"] == "completed" + + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, store.project_scope("P1")) + assert [d["filename"] for d in docs] == ["notes.txt"] + # Scoped: a sibling project cannot see it. + assert store.list_documents(conn, store.project_scope("P2")) == [] + assert store.search_lexical(conn, store.project_scope("P1"), "bravo", 5) + finally: + conn.close() + + +@pytest.mark.parametrize( + "raw", + [ + "x" * 300 + ".txt", + "y" * 512 + ".PDF", + "../" * 80 + "deep.md", + ], +) +def test_long_filenames_keep_their_extension(raw): + # _save_upload gates on the extension, so trimming it would reject the file. + out = _sanitize_filename(raw) + assert len(out) <= 200 + assert os.path.splitext(out)[1].lower() == os.path.splitext(raw)[1].lower() + + +@pytest.mark.parametrize( + "raw", + [ + "../../etc/passwd.txt", + "..\\..\\windows\\evil.txt", + "/absolute/notes.txt", + "C:\\Users\\me\\notes.txt", + ], +) +def test_sanitized_filenames_carry_no_path(raw): + out = _sanitize_filename(raw) + assert "/" not in out and "\\" not in out + + +@pytest.mark.parametrize("raw", ["." * 300, "noext" * 100, "a" * 100 + "." + "e" * 250]) +def test_sanitizer_degrades_safely(raw): + assert 0 < len(_sanitize_filename(raw)) <= 200 diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index ad46f6ee41..5cdbe4f2a5 100644 --- a/studio/backend/tests/test_rocm_oom_guard.py +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -80,6 +80,7 @@ class TestCanonicalGcnArchName: [ ("gfx1150", True), # Strix Point ("gfx1151", True), # Strix Halo + ("gfx1152", True), # Krackan Point (Radeon 860M/840M) ("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete ("gfx906", False), # MI50 — discrete server GPU ("gfx1201", False), # RX 9070 XT — discrete @@ -166,9 +167,15 @@ class TestDeviceNameFallback: # gfx1151 Gorgon Halo (Ryzen AI Max 400 refresh) "Radeon 8065S Graphics", # Ryzen AI Max+ 495 "AMD Radeon 8065S", + # gfx1152 Krackan Point (Ryzen AI 7 350 / AI 5 340) + "Radeon 860M", + "AMD Radeon 860M Graphics", + "Radeon 840M", + "AMD Radeon 840M Graphics", # case variants "RADEON 8060S GRAPHICS", "radeon 8050s", + "RADEON 860M", ], ) def test_unified_memory_detected(self, device_name: str) -> None: diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index a8c0c2305f..b491134045 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -85,6 +85,14 @@ def test_arg_parser_secure_polarity_and_not_secure_alias(): assert parser.parse_args(["--not-secure", "--secure"]).secure is True +def test_arg_parser_dns_pinning_opt_out_defaults_off(): + import run + + parser = run._build_arg_parser() + assert parser.parse_args([]).disable_dns_pinning is False + assert parser.parse_args(["--disable-dns-pinning"]).disable_dns_pinning is True + + def test_run_server_accepts_enable_tools_kwarg(): import inspect diff --git a/studio/backend/tests/test_server_disk_logging_outstream.py b/studio/backend/tests/test_server_disk_logging_outstream.py new file mode 100644 index 0000000000..0ff27666a0 --- /dev/null +++ b/studio/backend/tests/test_server_disk_logging_outstream.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for the Colab "OutStream has no attribute 'watch_fd_thread'" +startup crash. + +Field report (Colab): Unsloth Studio dies at server startup with +``❌ Unsloth Studio failed to start: 'OutStream' object has no attribute +'watch_fd_thread'``. + +Root cause chain: + * Colab's ipykernel ``OutStream`` is created with ``watchfd=False``, so it + never gains a ``watch_fd_thread``; the ``OutStream.close()`` shipped in the + affected ipykernel versions joins that thread unconditionally and raises + ``AttributeError`` (ipython/ipykernel#867). + * ``run._setup_server_disk_logging()`` replaces ``sys.stdout``/``sys.stderr`` + with a ``_TeeStream``. That changes the console object identity, so Colab's + ``absl`` logging handler -- which captured the ORIGINAL OutStream and whose + ``close()`` deliberately skips ``sys.stdout``/``sys.stderr`` -- no longer + recognizes it as the live console. + * ``run_server`` builds ``uvicorn.Config(...)``, whose ``configure_logging`` -> + ``logging.config.dictConfig`` -> ``logging.shutdown`` closes every existing + handler. The absl handler then calls ``OutStream.close()`` on the orphaned + stream, and the AttributeError aborts startup. + +These tests reproduce the mechanism with a stand-in OutStream (Colab-identical +constructs are not importable off Colab) and assert the tee/console path used at +startup survives it. +""" + +from __future__ import annotations + +import io +import logging +import sys +import weakref +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import run as run_mod # noqa: E402 + + +class _ColabOutStream(io.TextIOBase): + """Stand-in for Colab's ipykernel OutStream built with ``watchfd=False``: + no ``watch_fd_thread`` and an unguarded ``close()`` that joins it + (ipython/ipykernel#867).""" + + def __init__(self, name: str, sink: io.StringIO): + self.name = name + self._sink = sink + + def write(self, s): + return self._sink.write(s) + + def flush(self): + pass + + def writable(self): + return True + + def isatty(self): + return False + + def close(self): + # Never set because watchfd=False -> AttributeError, exactly as Colab. + self.watch_fd_thread.join() + + def __del__(self): + # io.TextIOBase.__del__ would call our buggy close() at GC (the harmless + # "Exception ignored" tail seen in Colab); silence it so the test is clean. + pass + + +class _WatchingOutStream(_ColabOutStream): + """OutStream with fd-watching ON: ``watch_fd_thread`` exists, close() is + well behaved and must keep working unchanged.""" + + def __init__(self, name: str, sink: io.StringIO): + super().__init__(name, sink) + self.close_ran = False + self.watch_fd_thread = type("_T", (), {"join": lambda self: None})() + + def close(self): + self.watch_fd_thread.join() + self.close_ran = True + + +class _AbslLikeHandler(logging.StreamHandler): + """Mirror of ``absl.logging.PythonHandler.close()``: close the captured + stream unless it is (still) one of the user-managed console streams.""" + + def close(self): + try: + user_managed = (sys.stderr, sys.stdout, sys.__stderr__, sys.__stdout__) + if self.stream not in user_managed and ( + not hasattr(self.stream, "isatty") or not self.stream.isatty() + ): + self.stream.close() + except ValueError: + pass + super().close() + + +class TestHardenConsoleClose: + def test_neutralizes_watchfd_false_close(self): + stream = _ColabOutStream("stdout", io.StringIO()) + with pytest.raises(AttributeError): + stream.close() # baseline: the ipykernel #867 bug is real + + stream = _ColabOutStream("stdout", io.StringIO()) + run_mod._harden_console_close(stream) + assert stream.close() is None # swallowed, no crash + + def test_healthy_close_still_runs_fully(self): + stream = _WatchingOutStream("stdout", io.StringIO()) + run_mod._harden_console_close(stream) + stream.close() + assert stream.close_ran is True + + def test_only_attributeerror_is_swallowed(self): + class _Boom: + def close(self): + raise ValueError("real teardown failure") + + stream = _Boom() + run_mod._harden_console_close(stream) + with pytest.raises(ValueError): + stream.close() + + def test_unrelated_attributeerror_still_propagates(self): + # Only #867 is neutralized; a genuine missing attribute during teardown + # must still surface instead of looking like a clean close. + class _Console: + def close(self): + return self.not_a_real_attribute + + stream = _Console() + run_mod._harden_console_close(stream) + with pytest.raises(AttributeError, match = "not_a_real_attribute"): + stream.close() + + def test_swallowed_across_attributeerror_message_shapes(self): + # Python 3.12 appends a "Did you mean" tail; the match must survive it, + # and pre-3.10 AttributeErrors carry no ``name``, only the message. + class _Suggesting: + def close(self): + raise AttributeError( + "'OutStream' object has no attribute 'watch_fd_thread'. " + "Did you mean: '_watch_pipe_fd'?" + ) + + stream = _Suggesting() + run_mod._harden_console_close(stream) + assert stream.close() is None + + def test_unsettable_close_is_left_alone(self): + # A stream whose close cannot be reassigned must not raise from hardening. + class _Frozen: + __slots__ = () + + def close(self): + return "ok" + + stream = _Frozen() + run_mod._harden_console_close(stream) # must not raise + assert stream.close() == "ok" + + +class TestTeeStreamClose: + def test_tee_close_over_buggy_stream_never_raises(self): + console = _ColabOutStream("stdout", io.StringIO()) + log = io.StringIO() + tee = run_mod._TeeStream(console, log) + tee.write("before-close") + tee.close() # must not raise despite the wrapped stream's broken close + assert log.getvalue() == "before-close" + + def test_tee_close_flushes_log(self): + class _FlushCounting(io.StringIO): + def __init__(self): + super().__init__() + self.flushes = 0 + + def flush(self): + self.flushes += 1 + super().flush() + + console, log = io.StringIO(), _FlushCounting() + tee = run_mod._TeeStream(console, log) + tee.write("x") + tee.close() + assert log.flushes >= 1 + + +class TestColabStartupRegression: + """End-to-end: the exact trigger -- an absl-style handler closing the + orphaned OutStream during the ``logging.shutdown`` that uvicorn's + ``uvicorn.Config`` -> ``dictConfig`` runs -- must not crash Studio, and the + tee must keep logging afterwards. + + ``logging.shutdown`` is driven over a LOCAL weakref list (identical code path + to ``logging.config._clearExistingHandlers``) so the global logging state and + pytest's own capture are untouched. + """ + + def _make_console_and_handlers(self, monkeypatch): + out_sink, err_sink = io.StringIO(), io.StringIO() + out_stream = _ColabOutStream("stdout", out_sink) + err_stream = _ColabOutStream("stderr", err_sink) + monkeypatch.setattr(sys, "stdout", out_stream) + monkeypatch.setattr(sys, "stderr", err_stream) + # absl-like handlers capture the ORIGINAL OutStreams (as in Colab). + handlers = [_AbslLikeHandler(sys.stdout), _AbslLikeHandler(sys.stderr)] + return out_sink, err_sink, out_stream, err_stream, handlers + + def test_baseline_reproduces_crash_without_fix(self, monkeypatch): + # Prove the test exercises the real path: swapping the console identity + # (what the tee does) makes the absl-like close hit #867. + _, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch) + try: + monkeypatch.setattr(sys, "stdout", io.StringIO()) + monkeypatch.setattr(sys, "stderr", io.StringIO()) + with pytest.raises(AttributeError, match = "watch_fd_thread"): + logging.shutdown([weakref.ref(h) for h in handlers]) + finally: + # Neutralize so a lingering handler can't crash global teardown. + run_mod._harden_console_close(out_stream) + run_mod._harden_console_close(err_stream) + for h in handlers: + try: + h.close() + except Exception: + pass + + def test_startup_survives_with_harden_and_tee(self, monkeypatch): + out_sink, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch) + + # Exactly what _setup_server_disk_logging does before serving: + run_mod._harden_console_close(sys.stdout) + run_mod._harden_console_close(sys.stderr) + log_fh = io.StringIO() + monkeypatch.setattr(sys, "stdout", run_mod._TeeStream(sys.stdout, log_fh)) + monkeypatch.setattr(sys, "stderr", run_mod._TeeStream(sys.stderr, log_fh)) + + # The close-storm uvicorn triggers via dictConfig -> logging.shutdown, + # closing the absl-like handlers over the (now orphaned) OutStreams. + logging.shutdown([weakref.ref(h) for h in handlers]) # must NOT raise + + # The tee still tees to both console and disk afterwards. + print("post-startup-line") + sys.stdout.flush() + assert "post-startup-line" in out_sink.getvalue() + assert "post-startup-line" in log_fh.getvalue() diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 00c7aeac69..23c70f8499 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -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. diff --git a/studio/backend/tests/test_web_fetch_extraction.py b/studio/backend/tests/test_web_fetch_extraction.py index b794ee3e81..d4c3d123c3 100644 --- a/studio/backend/tests/test_web_fetch_extraction.py +++ b/studio/backend/tests/test_web_fetch_extraction.py @@ -15,6 +15,8 @@ from __future__ import annotations import sys from pathlib import Path +import pytest + _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) @@ -715,6 +717,61 @@ def test_fetch_url_raw_missing_content_type_reported_empty(monkeypatch): assert content_type == "" +@pytest.mark.parametrize( + "disable_dns_pinning,expected_url", + [ + (False, "https://203.0.113.7:8443/page?q=1"), + (True, "https://example.com:8443/page?q=1"), + ], +) +def test_fetch_url_raw_dns_pinning_proxy_opt_out(monkeypatch, disable_dns_pinning, expected_url): + import email + import urllib.request + + import core.inference.tools as tools_mod + + class _FakeResp: + headers = email.message_from_string("Content-Type: text/plain\n") + + def __init__(self): + self._body = b"ok" + + def read(self, n = -1): + body, self._body = self._body, b"" + return body + + requested = [] + + class _FakeOpener: + def open( + self, + req, + timeout = None, + ): + requested.append(req) + return _FakeResp() + + resolved = [] + + def resolve(host, port): + resolved.append((host, port)) + return True, "", "203.0.113.7" + + monkeypatch.setenv("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "1" if disable_dns_pinning else "0") + monkeypatch.setattr(tools_mod, "_validate_and_resolve_host", resolve) + monkeypatch.setattr(urllib.request, "build_opener", lambda *handlers: _FakeOpener()) + + err, body, _content_type = tools_mod._fetch_url_raw( + "https://user:secret@example.com:8443/page?q=1" + ) + + assert err is None + assert body == "ok" + assert resolved == [("example.com", 8443)] + assert [req.full_url for req in requested] == [expected_url] + assert requested[0].get_header("Host") == "example.com:8443" + + def test_fetch_page_text_missing_content_type_html_sniffed(monkeypatch): # A header-less server returning an HTML body must still be converted. def fake_fetch( diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 174e6ef4dc..83602842af 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -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: " 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 diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index c95112c748..a31d9b6ced 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -105,7 +105,6 @@ import { archiveChatItem, ChatSearchDialog, clearNewChatDraft, - createChatProject, deleteChatProject, deleteChatItem, listStoredChatThreads, @@ -123,6 +122,7 @@ import { type ProjectRecord, type SidebarItem, } from "@/features/chat"; +import { NewProjectDialog } from "@/features/chat/components/new-project-dialog"; import { useAppearanceCustomStore, useSettingsDialogStore, @@ -696,7 +696,6 @@ export function AppSidebar() { }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); - const [projectNameDraft, setProjectNameDraft] = useState(""); const [projectCreateMoveTarget, setProjectCreateMoveTarget] = useState(null); const renameTrimmed = renameDraft.trim(); @@ -849,28 +848,26 @@ export function AppSidebar() { } } - async function commitCreateProject() { - const name = projectNameDraft.trim(); - if (!name) return; + // "New project" from a chat's menu moves that chat in and stays put; + // otherwise open the project, unless a slow upload outlasted the route the + // user was on when they hit create. + async function afterCreateProject( + project: ProjectRecord, + { stayedOnRoute }: { stayedOnRoute: boolean }, + ) { const moveTarget = projectCreateMoveTarget; + setProjectCreateMoveTarget(null); + if (!moveTarget) { + if (stayedOnRoute) openProject(project.id); + return; + } try { - const project = await createChatProject(name); - if (moveTarget) { - await moveChatItemToProject(moveTarget, project.id); - if (activeThreadId === moveTarget.id) { - useChatRuntimeStore.getState().setActiveProjectId(project.id); - } - } - setCreatingProject(false); - setProjectNameDraft(""); - setProjectCreateMoveTarget(null); - if (moveTarget) { - return; - } else { - openProject(project.id); + await moveChatItemToProject(moveTarget, project.id); + if (activeThreadId === moveTarget.id) { + useChatRuntimeStore.getState().setActiveProjectId(project.id); } } catch (err) { - toast.error(moveTarget ? "Failed to create and move chat" : "Failed to create project", { + toast.error("Failed to move chat to the new project", { description: err instanceof Error ? err.message : undefined, }); } @@ -1050,7 +1047,6 @@ export function AppSidebar() { { setProjectCreateMoveTarget(item); - setProjectNameDraft(""); setCreatingProject(true); }} > @@ -1393,7 +1389,6 @@ export function AppSidebar() { onClick={(e) => { e.stopPropagation(); setProjectCreateMoveTarget(null); - setProjectNameDraft(""); setCreatingProject(true); }} className="sidebar-row-action group-hover/projects-item:opacity-100 group-hover/projects-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto group-data-[collapsible=icon]:hidden" @@ -1855,6 +1850,18 @@ export function AppSidebar() { )} + {/* Collapsed rail has no room for the cog on the profile row, so it + sits above the avatar instead. */} + { + useSettingsDialogStore.getState().openDialog(); + closeMobileIfOpen(); + }} + /> @@ -2160,58 +2167,18 @@ export function AppSidebar() { - { setCreatingProject(open); - if (!open) { - setProjectNameDraft(""); - setProjectCreateMoveTarget(null); - } + if (!open) setProjectCreateMoveTarget(null); }} - > - - - - {projectCreateMoveTarget ? "Move to new project" : "New project"} - - - setProjectNameDraft(event.target.value)} - onKeyDown={(event) => { - if (event.key === "Enter") { - event.preventDefault(); - void commitCreateProject(); - } - }} - autoFocus - maxLength={120} - placeholder="Project name" - aria-label="Project name" - className="focus-visible:border-input focus-visible:ring-0" - /> - - - - - - + title={ + projectCreateMoveTarget ? "Move to new project" : "Create project" + } + submitLabel={projectCreateMoveTarget ? "Create and move" : "Create project"} + onCreated={afterCreateProject} + /> ); } diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index 97891e9358..2b01f7b719 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -346,6 +346,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ const [manualOpen, setManualOpen] = useState(false); const [dismissedWhileStreaming, setDismissedWhileStreaming] = useState(false); + const [retainStreamingHeight, setRetainStreamingHeight] = useState(false); const [duration, setDuration] = useState(0); const startTimeRef = useRef(null); @@ -361,13 +362,26 @@ 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]); + // Keep the streaming height cap until the automatic close finishes. Removing + // it on the completion frame expands long reasoning to its full height before + // the collapsible can close, which makes the entire chat jump. + useEffect(() => { + const timeout = window.setTimeout( + () => setRetainStreamingHeight(isReasoningStreaming), + isReasoningStreaming ? 0 : ANIMATION_DURATION, + ); + return () => window.clearTimeout(timeout); + }, [isReasoningStreaming]); + // Open while streaming (unless dismissed), or once manually opened. const isOpen = (isReasoningStreaming && !dismissedWhileStreaming) || manualOpen; const variant = isOpen ? "outline" : "ghost"; @@ -378,6 +392,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ if (isReasoningStreaming) { setDismissedWhileStreaming(!open); } else { + if (open) { + setRetainStreamingHeight(false); + } setManualOpen(open); } }, @@ -407,7 +424,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ aria-busy={isReasoningStreaming} streaming={isReasoningStreaming} > - + {children} diff --git a/studio/frontend/src/features/chat/api-provider-logo.tsx b/studio/frontend/src/features/chat/api-provider-logo.tsx index f9d574b8fa..7de9bea9a8 100644 --- a/studio/frontend/src/features/chat/api-provider-logo.tsx +++ b/studio/frontend/src/features/chat/api-provider-logo.tsx @@ -40,10 +40,10 @@ 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. - */ +// Monochrome logos vanish on a dark background. +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 +63,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, )} /> diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 7452cf3447..7e0544e2d1 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -185,6 +185,10 @@ import { listStoredChatThreads, } from "./utils/chat-history-storage"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; +import { + consumeProjectSourcesPending, + hasProjectSourcesPending, +} from "@/features/rag/components/project-source-dropzone"; const ProjectSourcesPanel = lazy(() => @@ -998,7 +1002,14 @@ function ProjectLanding({ const active = useChatActive(); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const initialActiveThreadRef = useRef(null); - const [projectTab, setProjectTab] = useState<"chats" | "sources">("chats"); + // Land on Sources when the project was just created with dropped files. + const [projectTab, setProjectTab] = useState<"chats" | "sources">(() => + hasProjectSourcesPending(projectId) ? "sources" : "chats", + ); + // Drop the marker once committed: React may replay the initializer above. + useEffect(() => { + consumeProjectSourcesPending(projectId); + }, [projectId]); const [pendingNewThreadId, setPendingNewThreadId] = useState( null, ); @@ -2676,7 +2687,7 @@ export function ChatPage({ config: meta?.config, nativePathToken: meta?.nativePathToken, nativePathExpiresAtMs: meta?.nativePathExpiresAtMs, - forceReload: isSameLoadedModel || undefined, + forceReload: meta?.forceReload ?? (isSameLoadedModel || undefined), }; await stageOrLoad(selection); })(); diff --git a/studio/frontend/src/features/chat/components/new-project-dialog.tsx b/studio/frontend/src/features/chat/components/new-project-dialog.tsx index 880129ac6c..6aca3d36c3 100644 --- a/studio/frontend/src/features/chat/components/new-project-dialog.tsx +++ b/studio/frontend/src/features/chat/components/new-project-dialog.tsx @@ -2,7 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useNavigate } from "@tanstack/react-router"; -import { useState } from "react"; +import { useEffect, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; import { @@ -12,31 +12,92 @@ import { DialogHeader, DialogTitle, } from "@/components/ui/dialog"; -import { Input } from "@/components/ui/input"; +import { + ProjectSourceDropzone, + type StagedSource, + uploadStagedSources, +} from "@/features/rag/components/project-source-dropzone"; import { toast } from "@/lib/toast"; +import { Folder02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { createChatProject } from "../hooks/use-chat-projects"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import type { ProjectRecord } from "../types"; -// Create-project dialog usable from the composer + menu. Creating opens the new -// project straight away rather than dropping the user on the projects list. +function currentRoute(): string { + if (typeof window === "undefined") return ""; + return window.location.pathname + window.location.search; +} + +// Create-project dialog for the composer, sidebar, and projects page. Creating +// opens the new project; `onCreated` overrides that for callers with their own +// follow-up (the sidebar's "move this chat to a new project"). export function NewProjectDialog({ open, onOpenChange, + title = "Create project", + submitLabel = "Create project", + onCreated, }: { open: boolean; onOpenChange: (open: boolean) => void; + title?: string; + submitLabel?: string; + onCreated?: ( + project: ProjectRecord, + context: { stayedOnRoute: boolean }, + ) => void | Promise; }) { const navigate = useNavigate(); const [name, setName] = useState(""); + const [staged, setStaged] = useState([]); + const [busy, setBusy] = useState(false); + // Uploads outlive this component, so a slow one must not yank the user to the + // new project after they have navigated away. + const mounted = useRef(true); + useEffect(() => { + // Set on setup, not just cleared on cleanup: StrictMode replays + // setup/cleanup/setup, which would otherwise leave this false forever. + mounted.current = true; + return () => { + mounted.current = false; + }; + }, []); + + function reset() { + setName(""); + setStaged([]); + } + + // Every close path routes through here: callers keep this mounted, so a draft + // left behind would resurface (and upload) on the next project. + function close() { + if (busy) return; + reset(); + onOpenChange(false); + } async function commitCreate() { const trimmed = name.trim(); - if (!trimmed) return; + if (!trimmed || busy) return; + setBusy(true); + // Sidebar callers keep this mounted across routes, so unmounting alone + // cannot tell whether the user has moved on during a slow upload. + const origin = currentRoute(); try { const project = await createChatProject(trimmed); + // Upload before closing so the Sources panel lists them on first fetch. + await uploadStagedSources(project.id, staged); + if (!mounted.current) return; + const stayedOnRoute = currentRoute() === origin; onOpenChange(false); - setName(""); + reset(); + if (onCreated) { + await onCreated(project, { stayedOnRoute }); + return; + } + if (!stayedOnRoute) return; const runtime = useChatRuntimeStore.getState(); runtime.setActiveThreadId(null); runtime.setActiveProjectId(project.id); @@ -45,6 +106,8 @@ export function NewProjectDialog({ toast.error("Failed to create project", { description: err instanceof Error ? err.message : undefined, }); + } finally { + setBusy(false); } } @@ -52,43 +115,59 @@ export function NewProjectDialog({ { - if (!next) setName(""); - onOpenChange(next); + if (next) { + onOpenChange(true); + return; + } + close(); }} > - + - New project + {title} - setName(e.target.value)} - onKeyDown={(e) => { - if (e.key === "Enter") { - e.preventDefault(); - void commitCreate(); - } - }} - autoFocus={true} - maxLength={120} - placeholder="Project name" - aria-label="Project name" - className="focus-visible:border-input focus-visible:ring-0" + {/* Name field: folder glyph in its own cell, divided from the input. */} +
+ + + +
+ - diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index bb19223a6a..48a6168555 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -62,6 +62,7 @@ import { import { isExternalModelId } from "../external-providers"; import { applyPerModelConfigToRuntime, + normalizeMaxSeqLength, type PerModelConfig, } from "@/features/model-picker"; import type { @@ -604,12 +605,19 @@ export function useChatModelRuntime() { async function performLoad(): Promise { if (abortCtrl.signal.aborted) throw new Error("Cancelled"); let previousWasUnloaded = false; + const pendingLoadConfig = + typeof selection !== "string" ? selection.config : undefined; + if (pendingLoadConfig) { + applyPerModelConfigToRuntime(pendingLoadConfig); + } const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const stateBeforeUnload = useChatRuntimeStore.getState(); let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; let approvedRemoteCodeFingerprint: string | null = null; - const maxSeqLength = stateBeforeUnload.params.maxSeqLength; + const maxSeqLength = + normalizeMaxSeqLength(pendingLoadConfig?.maxSeqLength) ?? + stateBeforeUnload.params.maxSeqLength; const previousActiveNativePathToken = stateBeforeUnload.activeNativePathToken; const previousIsGguf = @@ -643,34 +651,54 @@ export function useChatModelRuntime() { const previousActiveNativePathExpiresAtMs = stateBeforeUnload.activeNativePathExpiresAtMs; // Snapshot the load settings at click time, before the awaits below - // (validation, the trust dialog, unload). - const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride; - const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype; + // (validation, the trust dialog, unload). When the picker staged a + // config payload, prefer it over the store: React may not have + // flushed NumericValueInput's blur commit into state yet. + const loadChatTemplateOverride = + pendingLoadConfig?.chatTemplateOverride?.trim() + ? pendingLoadConfig.chatTemplateOverride + : stateBeforeUnload.chatTemplateOverride; + const loadKvCacheDtype = + pendingLoadConfig?.kvCacheDtype ?? stateBeforeUnload.kvCacheDtype; // gpuMemoryMode is a standing preference (kept across a model switch); // the rest are per-model knobs the reset below clears, so they are // re-baselined there in lock-step with the store. - let loadCustomContextLength = stateBeforeUnload.customContextLength; + let loadCustomContextLength = + pendingLoadConfig?.customContextLength ?? + stateBeforeUnload.customContextLength; const loadGgufContextLength = stateBeforeUnload.ggufContextLength; - const loadTensorParallel = stateBeforeUnload.tensorParallel; + const loadTensorParallel = + pendingLoadConfig?.tensorParallel ?? stateBeforeUnload.tensorParallel; const loadActivePresetSource = stateBeforeUnload.activePresetSource; const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant; - const loadGpuMemoryMode = stateBeforeUnload.gpuMemoryMode; - let loadGpuLayers = stateBeforeUnload.gpuLayers; - let loadNCpuMoe = stateBeforeUnload.nCpuMoe; + const loadGpuMemoryMode = + pendingLoadConfig?.gpuMemoryMode ?? stateBeforeUnload.gpuMemoryMode; + let loadGpuLayers = + pendingLoadConfig?.gpuLayers ?? stateBeforeUnload.gpuLayers; + let loadNCpuMoe = + pendingLoadConfig?.nCpuMoe ?? stateBeforeUnload.nCpuMoe; let loadSplitRatio = stateBeforeUnload.splitRatio; // Reconcile the persisted pick against the GPUs present now, so a stale // cross-host / now-hidden pick is dropped before /load rather than // rejected there. Warm the device cache first: load-on-selection can // run before any GPU hook mounted, and a cold cache would pass the // pick through unvalidated. validateGpuIds derives from this too. - if (stateBeforeUnload.selectedGpuIds != null) { + if ( + pendingLoadConfig?.selectedGpuIds !== undefined || + stateBeforeUnload.selectedGpuIds != null + ) { await ensureGpuDeviceCache(); } - let loadSelectedGpuIds = reconcilePersistedGpuIds( - stateBeforeUnload.selectedGpuIds, - ); - let loadSpeculativeType = stateBeforeUnload.speculativeType; - let loadSpecDraftNMax = stateBeforeUnload.specDraftNMax; + let loadSelectedGpuIds = + pendingLoadConfig?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds) + : reconcilePersistedGpuIds(stateBeforeUnload.selectedGpuIds); + let loadSpeculativeType = + pendingLoadConfig?.speculativeType != null + ? normalizeSpeculativeType(pendingLoadConfig.speculativeType) + : stateBeforeUnload.speculativeType; + let loadSpecDraftNMax = + pendingLoadConfig?.specDraftNMax ?? stateBeforeUnload.specDraftNMax; try { // Lightweight pre-flight validation: avoid unloading a working model // if the new identifier is clearly invalid (e.g. bad HF id / path). @@ -810,15 +838,23 @@ export function useChatModelRuntime() { // model loads at Auto/native, not the previous model's pin. customContextLength: null, }); - loadSpeculativeType = persistedSpeculativeType; - loadSpecDraftNMax = null; + loadSpeculativeType = + pendingLoadConfig?.speculativeType != null + ? normalizeSpeculativeType(pendingLoadConfig.speculativeType) + : persistedSpeculativeType; + loadSpecDraftNMax = pendingLoadConfig?.specDraftNMax ?? null; // Keep the click-time snapshot in lock-step with the store reset so // the load below sizes against the cleared per-model knobs, not the // previous model's (gpuMemoryMode is standing, so left as captured). - loadCustomContextLength = null; - loadSelectedGpuIds = null; - loadGpuLayers = GPU_LAYERS_AUTO; - loadNCpuMoe = 0; + // An explicit staged config from run-settings still wins. + loadCustomContextLength = + pendingLoadConfig?.customContextLength ?? null; + loadSelectedGpuIds = + pendingLoadConfig?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds) + : null; + loadGpuLayers = pendingLoadConfig?.gpuLayers ?? GPU_LAYERS_AUTO; + loadNCpuMoe = pendingLoadConfig?.nCpuMoe ?? 0; loadSplitRatio = null; } diff --git a/studio/frontend/src/features/chat/projects-page.tsx b/studio/frontend/src/features/chat/projects-page.tsx index e20e517787..192c4e2331 100644 --- a/studio/frontend/src/features/chat/projects-page.tsx +++ b/studio/frontend/src/features/chat/projects-page.tsx @@ -34,7 +34,6 @@ import { isTauri } from "@/lib/api-base"; import { isDownloadCancelled, pickNativeChatImport } from "@/lib/native-files"; import { toast } from "@/lib/toast"; import { - createChatProject, deleteChatProject, renameChatProject, useChatProjects, @@ -42,6 +41,7 @@ import { usePinnedProjectsStore, type ProjectRecord, } from "@/features/chat"; +import { NewProjectDialog } from "./components/new-project-dialog"; import { Delete02Icon, Download01Icon, @@ -124,7 +124,6 @@ export function ProjectsPage() { ); const [creating, setCreating] = useState(false); - const [nameDraft, setNameDraft] = useState(""); const [renaming, setRenaming] = useState(null); const [renameDraft, setRenameDraft] = useState(""); const [deleting, setDeleting] = useState(null); @@ -258,21 +257,6 @@ export function ProjectsPage() { navigate({ to: "/chat", search: { project: projectId } }); } - async function commitCreate() { - const name = nameDraft.trim(); - if (!name) return; - try { - const project = await createChatProject(name); - setCreating(false); - setNameDraft(""); - openProject(project.id); - } catch (err) { - toast.error("Failed to create project", { - description: err instanceof Error ? err.message : undefined, - }); - } - } - async function commitRename() { const target = renaming; const name = renameDraft.trim(); @@ -469,14 +453,7 @@ export function ProjectsPage() { - + @@ -511,10 +488,7 @@ export function ProjectsPage() { - - - - + {/* Create project (name + drag-and-drop sources) */} + {/* Rename project */} void; + inputRef?: Ref; }) { return (
@@ -146,6 +158,7 @@ function MaxSeqLengthSetting({
void; displayValue?: string; info?: ReactNode; + inputRef?: Ref; }) { return (
@@ -199,6 +214,7 @@ function AdvancedGpuSlider({ {info && {info}}
) => void; layerCount: number | null; moeLayerCount: number | null; + gpuLayersInputRef?: Ref; + moeLayersInputRef?: Ref; }) { const gpuDevices = useGpuDevices(); const mode = config.gpuMemoryMode ?? "auto"; @@ -322,6 +342,7 @@ function GpuMemorySettings({ <> ) => void; @@ -407,6 +431,8 @@ function GgufAdvancedSettings({ onEditTemplate: () => void; layerCount: number | null; moeLayerCount: number | null; + gpuLayersInputRef?: Ref; + moeLayersInputRef?: Ref; }) { return ( <> @@ -535,6 +561,8 @@ function GgufAdvancedSettings({ update={update} layerCount={layerCount} moeLayerCount={moeLayerCount} + gpuLayersInputRef={gpuLayersInputRef} + moeLayersInputRef={moeLayersInputRef} /> @@ -597,6 +625,10 @@ export function ModelConfigPage({ const [showAdvanced, setShowAdvanced] = useState(() => hasNonDefaultAdvanced(config), ); + const contextInputRef = useRef(null); + const maxSeqLengthInputRef = useRef(null); + const gpuLayersInputRef = useRef(null); + const moeLayersInputRef = useRef(null); const nativePathToken = target.meta.nativePathToken ?? (isActiveModel ? activeNativePathToken : null); @@ -744,11 +776,6 @@ export function ModelConfigPage({ ? { ...config, customContextLength: activeLoadedContext } : config : config; - // Load request needs a concrete max length; substitute the fallback here only, - // never in the persisted runtimeConfig. - const loadConfig = target.isGguf - ? runtimeConfig - : { ...runtimeConfig, maxSeqLength: maxSeqLengthValue }; const rememberChanged = remember !== savedRemember; const persistenceOnly = isActiveModel && atBaseline && rememberChanged; const primaryActionLabel = persistenceOnly @@ -760,18 +787,90 @@ export function ModelConfigPage({ : "Load model"; const handleRun = () => { - const defaultConfig = isDefaultConfig(runtimeConfig); + // Same-click Load/Reload: a numeric draft the user just typed is flushed only + // by that input's blur handler, which updates the parent config after this + // click closure already captured the stale value. Commit every numeric input + // imperatively so the staged load honors what the user just typed, not just + // the Context field. + const committedContext = target.isGguf + ? contextInputRef.current?.commit() + : undefined; + const committedMaxSeqLength = target.isGguf + ? undefined + : maxSeqLengthInputRef.current?.commit(); + const committedGpuLayers = target.isGguf + ? gpuLayersInputRef.current?.commit() + : undefined; + const committedMoeLayers = target.isGguf + ? moeLayersInputRef.current?.commit() + : undefined; + + const pendingPatch: Partial = {}; + if (committedContext != null) { + pendingPatch.customContextLength = committedContext; + } + if (committedMaxSeqLength != null) { + pendingPatch.maxSeqLength = clampMaxSeqLength( + committedMaxSeqLength, + MAX_SEQ_LENGTH_MAX, + ); + } + if (committedGpuLayers != null) { + pendingPatch.gpuLayers = committedGpuLayers; + } + if (committedMoeLayers != null) { + pendingPatch.nCpuMoe = committedMoeLayers; + } + const hasPending = + committedContext != null || + committedMaxSeqLength != null || + committedGpuLayers != null || + committedMoeLayers != null; + + const effectiveConfig = hasPending + ? { ...config, ...pendingPatch } + : config; + // pinFixedLayerContext above was computed from the render-time config, before + // the same-click GPU Layers draft was committed. Recompute it from + // effectiveConfig so committing a positive fixed-layer value still pins the + // fitted context; otherwise the saved config carries customContextLength: null + // and a later fresh load sends the native context with fixed layers (the OOM + // the pin exists to avoid). + const effectivePinFixedLayerContext = + target.isGguf && + effectiveConfig.gpuMemoryMode === "manual" && + effectiveConfig.gpuLayers != null && + effectiveConfig.gpuLayers >= 0 && + effectiveConfig.customContextLength == null && + activeLoadedContext != null; + const effectiveRuntimeConfig = hasPending + ? effectivePinFixedLayerContext + ? { ...effectiveConfig, customContextLength: activeLoadedContext } + : effectiveConfig + : runtimeConfig; + // Non-GGUF load substitutes the resolved max sequence length; recompute it + // from the committed draft so a same-click Max Seq Length edit is not lost. + const effectiveMaxSeqLengthValue = + committedMaxSeqLength == null + ? maxSeqLengthValue + : (normalizeMaxSeqLength(effectiveConfig.maxSeqLength) ?? + clampMaxSeqLength(DEFAULT_MAX_SEQ_LENGTH, nativeMaxSeqLength)); + // Recheck the committed draft so Save/Forget reloads when needed. + const effectiveAtBaseline = perModelConfigsEqual(effectiveConfig, baseline); + const effectivePersistenceOnly = + isActiveModel && effectiveAtBaseline && rememberChanged; + const defaultConfig = isDefaultConfig(effectiveRuntimeConfig); let saveFailed = false; if (remember) { saveFailed = !savePerModelConfig( target.id, target.ggufVariant, - runtimeConfig, + effectiveRuntimeConfig, ); } else { saveFailed = !deletePerModelConfig(target.id, target.ggufVariant); } - if (persistenceOnly) { + if (effectivePersistenceOnly) { if (saveFailed) { toast.error("Couldn't save settings for this model."); return; @@ -791,7 +890,10 @@ export function ModelConfigPage({ if (saveFailed) { toast.error("Couldn't save these settings, loading with them anyway."); } - onRun(loadConfig); + const effectiveLoadConfig = target.isGguf + ? effectiveRuntimeConfig + : { ...effectiveRuntimeConfig, maxSeqLength: effectiveMaxSeqLengthValue }; + onRun(effectiveLoadConfig); }; return ( @@ -838,6 +940,7 @@ export function ModelConfigPage({ setTemplateOpen(true)} layerCount={stagedDims?.layerCount ?? null} moeLayerCount={stagedDims?.moeLayerCount ?? null} + gpuLayersInputRef={gpuLayersInputRef} + moeLayersInputRef={moeLayersInputRef} /> )} @@ -914,6 +1019,7 @@ export function ModelConfigPage({ value={maxSeqLengthValue} max={maxSeqLengthMax} inputMax={MAX_SEQ_LENGTH_MAX} + inputRef={maxSeqLengthInputRef} onChange={(value) => update({ maxSeqLength: clampMaxSeqLength(value, MAX_SEQ_LENGTH_MAX), diff --git a/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx b/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx index 2489927fc2..be9aa7745c 100644 --- a/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx +++ b/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx @@ -2,7 +2,13 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { cn } from "@/lib/utils"; -import { useRef, useState } from "react"; +import { + forwardRef, + useEffect, + useImperativeHandle, + useRef, + useState, +} from "react"; export function snapToStep( value: number, @@ -28,44 +34,109 @@ function sanitizeNumeric(raw: string, allowNegative: boolean): string { return `${sign}${head}${tail}`; } -export function NumericValueInput({ - value, - min, - max, - step, - onChange, - displayValue, - className, - ariaLabel, - size: sizeAttr, - disabled = false, -}: { - value: number; - min?: number; - max?: number; - step: number; - onChange: (v: number) => void; - displayValue?: string; - className?: string; - ariaLabel?: string; - size?: number; - disabled?: boolean; -}) { +export type NumericValueInputHandle = { + /** Commit a valid focused/same-click draft; null when none is pending. */ + commit: () => number | null; +}; + +export const NumericValueInput = forwardRef< + NumericValueInputHandle, + { + value: number; + min?: number; + max?: number; + step: number; + onChange: (v: number) => void; + displayValue?: string; + className?: string; + ariaLabel?: string; + size?: number; + disabled?: boolean; + } +>(function NumericValueInput( + { + value, + min, + max, + step, + onChange, + displayValue, + className, + ariaLabel, + size: sizeAttr, + disabled = false, + }, + ref, +) { const [focused, setFocused] = useState(false); const [draft, setDraft] = useState(""); const cancelBlurCommitRef = useRef(false); + const draftRef = useRef(""); + const dirtyRef = useRef(false); + // Same-click Load: blur commits via onChange and clears dirtyRef before the + // button onClick runs, while parent `value` is still stale. Keep the blur + // result for one imperative commit(); clear when `value` catches up or on + // focus / external edits (Reset, slider). + const lastBlurCommittedRef = useRef(null); - const commit = (raw: string) => { + // The blur bridge is only valid across the single synchronous gesture that set + // it: blur commits during a button's mousedown and that button's onClick + // consumes it via commit() before React re-renders. Any settled render means the + // gesture is over, so drop the cache on every commit. Keying this on [value] + // alone missed a Reset (or other external edit) that restores the shown value + // unchanged when the blur did dispatch onChange (final !== value): value nets + // back to its prior number, so the effect never re-ran, the stale pin survived, + // and the next Load/Save replayed the override Reset had removed. + useEffect(() => { + lastBlurCommittedRef.current = null; + }); + + const commitDraft = (raw: string): number | null => { const parsed = Number.parseFloat(raw); if (!Number.isFinite(parsed)) { - return; + return null; } const final = snapToStep(parsed, step, min, max); if (final !== value) { onChange(final); } + return final; }; + useImperativeHandle( + ref, + () => ({ + commit: () => { + if (dirtyRef.current) { + const raw = draftRef.current; + const final = commitDraft(raw); + dirtyRef.current = false; + lastBlurCommittedRef.current = null; + if (final == null) { + draftRef.current = String(value); + } + if (focused) { + setFocused(false); + } + return final; + } + const blurCommitted = lastBlurCommittedRef.current; + if (blurCommitted != null) { + lastBlurCommittedRef.current = null; + if (focused) { + setFocused(false); + } + return blurCommitted; + } + if (focused) { + setFocused(false); + } + return null; + }, + }), + [draft, focused, max, min, onChange, step, value], + ); + const displayed = focused ? draft : (displayValue ?? String(value)); return ( @@ -82,7 +153,11 @@ export function NumericValueInput({ aria-label={ariaLabel} onFocus={(e) => { cancelBlurCommitRef.current = false; - setDraft(String(value)); + dirtyRef.current = false; + lastBlurCommittedRef.current = null; + const next = String(value); + draftRef.current = next; + setDraft(next); setFocused(true); const target = e.currentTarget; requestAnimationFrame(() => target.select()); @@ -90,24 +165,47 @@ export function NumericValueInput({ onBlur={() => { if (cancelBlurCommitRef.current) { cancelBlurCommitRef.current = false; - } else { - commit(draft); + lastBlurCommittedRef.current = null; + } else if (dirtyRef.current) { + const final = commitDraft(draftRef.current); + dirtyRef.current = false; + if (final == null) { + draftRef.current = String(value); + lastBlurCommittedRef.current = null; + } else { + draftRef.current = String(final); + // Only bridge the still-stale parent value when the blur actually + // dispatched onChange (final !== value). When final === value the + // parent is already current, so there is nothing to bridge; caching + // here would leave a stale pin that a later Reset or external edit + // (which doesn't change the displayed value) can never clear, so a + // following Load/Save would recreate the override Reset removed. + lastBlurCommittedRef.current = final !== value ? final : null; + } } setFocused(false); }} - onChange={(e) => - setDraft(sanitizeNumeric(e.target.value, (min ?? 0) < 0)) - } + onChange={(e) => { + dirtyRef.current = true; + lastBlurCommittedRef.current = null; + const next = sanitizeNumeric(e.target.value, (min ?? 0) < 0); + draftRef.current = next; + setDraft(next); + }} onKeyDown={(e) => { if (e.key === "Enter") { e.currentTarget.blur(); } else if (e.key === "Escape") { cancelBlurCommitRef.current = true; - setDraft(String(value)); + dirtyRef.current = false; + lastBlurCommittedRef.current = null; + const next = String(value); + draftRef.current = next; + setDraft(next); e.currentTarget.blur(); } }} className={cn(className)} /> ); -} +}); diff --git a/studio/frontend/src/features/model-picker/index.ts b/studio/frontend/src/features/model-picker/index.ts index d2b4785ec3..383d441f09 100644 --- a/studio/frontend/src/features/model-picker/index.ts +++ b/studio/frontend/src/features/model-picker/index.ts @@ -12,6 +12,7 @@ export { export { hfModelFitsDevice } from "./components/model-selector/recommended-fit"; export { NumericValueInput, + type NumericValueInputHandle, snapToStep, } from "./components/numeric-value-input"; export { SidebarModelConfig } from "./components/sidebar-model-config"; diff --git a/studio/frontend/src/features/rag/components/project-source-dropzone.tsx b/studio/frontend/src/features/rag/components/project-source-dropzone.tsx new file mode 100644 index 0000000000..61a66e985a --- /dev/null +++ b/studio/frontend/src/features/rag/components/project-source-dropzone.tsx @@ -0,0 +1,304 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; +import { File02Icon, FolderAddIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { XIcon } from "lucide-react"; +import { useCallback, useRef, useState } from "react"; +import { + invalidateProjectSources, + uploadProjectDocument, +} from "../api/rag-api"; +import { RAG_UPLOAD_ACCEPT } from "../types/rag"; +import { resolveVisionOverrides } from "./vision-overrides"; + +/** A file picked before the project exists, held until create commits. */ +export interface StagedSource { + id: string; + file: File; +} + +// Client-side dedup key; backend dedups authoritatively by content hash. +function fileSignature(file: File): string { + return `${file.name}|${file.size}|${file.lastModified}`; +} + +function formatSize(bytes: number): string { + if (!Number.isFinite(bytes) || bytes <= 0) return ""; + const units = ["B", "KB", "MB", "GB"]; + let value = bytes; + let unit = 0; + while (value >= 1024 && unit < units.length - 1) { + value /= 1024; + unit += 1; + } + const shown = + value >= 10 || unit === 0 + ? String(Math.round(value)) + : value.toFixed(1).replace(/\.0$/, ""); + return `${shown} ${units[unit]}`; +} + +const ACCEPTED_EXTS = new Set( + RAG_UPLOAD_ACCEPT.split(",").map((ext) => ext.trim().toLowerCase()), +); + +// `accept` only filters the picker, so a drop can carry anything. A folder +// arrives as an extension-less entry, which this rejects along with the types +// the backend would 400 on. +function isSupported(file: File): boolean { + const dot = file.name.lastIndexOf("."); + if (dot <= 0) return false; + return ACCEPTED_EXTS.has(file.name.slice(dot).toLowerCase()); +} + +/** Merge a selection into the staged list. Returns the names it would not take, + * so the caller can say so once instead of dropping them silently. */ +function addStagedSources( + staged: StagedSource[], + incoming: FileList | File[], +): { next: StagedSource[]; unsupported: string[]; duplicates: string[] } { + const seen = new Set(staged.map((entry) => fileSignature(entry.file))); + const next = [...staged]; + const unsupported: string[] = []; + const duplicates: string[] = []; + for (const file of Array.from(incoming)) { + if (!isSupported(file)) { + unsupported.push(file.name); + continue; + } + const signature = fileSignature(file); + if (seen.has(signature)) { + duplicates.push(file.name); + continue; + } + seen.add(signature); + next.push({ + id: `staged_${Math.random().toString(36).slice(2)}`, + file, + }); + } + return { next, unsupported, duplicates }; +} + +// Projects created with staged files, so the landing can open on Sources. +const projectsWithPendingSources = new Set(); + +function markProjectSourcesPending(projectId: string): void { + projectsWithPendingSources.add(projectId); +} + +/** Whether this project was just created with staged sources. Read-only, so it + * is safe in a render pass that React may replay. */ +export function hasProjectSourcesPending(projectId: string): boolean { + return projectsWithPendingSources.has(projectId); +} + +/** Drop the marker once the landing has committed. */ +export function consumeProjectSourcesPending(projectId: string): void { + projectsWithPendingSources.delete(projectId); +} + +/** Upload staged files to a new project. Indexing runs in the background; a + * per-file failure toasts and never blocks project creation. */ +export async function uploadStagedSources( + projectId: string, + staged: StagedSource[], +): Promise { + if (staged.length === 0) return; + invalidateProjectSources(projectId); + markProjectSourcesPending(projectId); + const { ocr, caption } = resolveVisionOverrides(); + const documentIds = new Set(); + const merged: string[] = []; + for (const { file } of staged) { + try { + const result = await uploadProjectDocument(projectId, file, ocr, caption); + // Same bytes under another name: the backend hashes content, so this is + // the document already uploaded. Say so rather than imply a new source. + if (documentIds.has(result.documentId)) merged.push(file.name); + else documentIds.add(result.documentId); + } catch (error) { + toast.error(`Couldn't upload ${file.name}`, { + description: error instanceof Error ? error.message : String(error), + }); + } + } + if (merged.length > 0) { + toast.info( + merged.length === 1 + ? `${merged[0]} matched a file already added` + : `${merged.length} files matched files already added`, + { description: "Identical contents are stored once." }, + ); + } + invalidateProjectSources(projectId); +} + +/** Create-project drop area: stages files until the project exists. */ +export function ProjectSourceDropzone({ + staged, + onChange, + disabled = false, +}: { + staged: StagedSource[]; + onChange: (next: StagedSource[]) => void; + disabled?: boolean; +}) { + const inputRef = useRef(null); + // Count enter/leave pairs: children fire dragleave on the parent. + const dragDepth = useRef(0); + const [dragging, setDragging] = useState(false); + + const addFiles = useCallback( + (files: FileList | File[]) => { + const { next, unsupported, duplicates } = addStagedSources(staged, files); + if (next.length !== staged.length) onChange(next); + if (unsupported.length > 0) { + toast.info( + unsupported.length === 1 + ? `Can't add ${unsupported[0]}` + : `Can't add ${unsupported.length} files`, + { description: `Supported types: ${RAG_UPLOAD_ACCEPT}` }, + ); + } + // Name, size and mtime can in principle match for two different files, so + // never drop one without saying so. + if (duplicates.length > 0) { + toast.info( + duplicates.length === 1 + ? `${duplicates[0]} is already added` + : `${duplicates.length} files were already added`, + ); + } + }, + [staged, onChange], + ); + + const endDrag = useCallback(() => { + dragDepth.current = 0; + setDragging(false); + }, []); + + return ( +
+

Sources

+ {/* Panel is the drop target; the inner button owns the click so staged + rows can carry their own remove buttons. */} +
{ + e.preventDefault(); + if (disabled) return; + dragDepth.current += 1; + setDragging(true); + }} + onDragOver={(e) => { + e.preventDefault(); + if (disabled) return; + e.dataTransfer.dropEffect = "copy"; + }} + onDragLeave={() => { + dragDepth.current = Math.max(0, dragDepth.current - 1); + if (dragDepth.current === 0) setDragging(false); + }} + onDrop={(e) => { + e.preventDefault(); + if (disabled) return; + endDrag(); + addFiles(Array.from(e.dataTransfer.files ?? [])); + }} + className={cn( + "rounded-[22px] border border-border transition-colors dark:border-white/10", + dragging && "border-primary/60 bg-primary/5", + disabled && "opacity-60", + )} + > + { + const files = Array.from(e.target.files ?? []); + e.target.value = ""; + addFiles(files); + }} + /> + {staged.length === 0 ? ( + + ) : ( +
+
    + {staged.map((entry) => ( +
  • + + + {entry.file.name} + + + {formatSize(entry.file.size)} + + +
    • + ))} +
+ +
+ )} +
+
+ ); +} diff --git a/studio/frontend/src/features/rag/components/use-rag-documents.ts b/studio/frontend/src/features/rag/components/use-rag-documents.ts index 8d6433d8c3..bdab7b0518 100644 --- a/studio/frontend/src/features/rag/components/use-rag-documents.ts +++ b/studio/frontend/src/features/rag/components/use-rag-documents.ts @@ -1,13 +1,8 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { useCallback, useEffect, useRef, useState } from "react"; -import { - CHAT_RAG_CAPTION_KEY, - CHAT_RAG_OCR_KEY, - useChatRuntimeStore, -} from "@/features/chat"; import { toast } from "@/lib/toast"; +import { useCallback, useEffect, useRef, useState } from "react"; import { deleteDocument, getJob, @@ -17,6 +12,7 @@ import { uploadThreadDocument, } from "../api/rag-api"; import type { DocumentStatus, RagDocument } from "../types/rag"; +import { resolveVisionOverrides } from "./vision-overrides"; export interface TrackedDocument extends RagDocument { progress?: number | null; @@ -263,18 +259,7 @@ export function useRagDocuments( tempId: string, ) => { try { - // Send vision-pass overrides only after the user has explicitly set them; - // otherwise backend env defaults own the ingest policy. - const state = useChatRuntimeStore.getState(); - const hasLocal = (key: string) => - typeof window !== "undefined" && - window.localStorage.getItem(key) !== null; - const ocr = hasLocal(CHAT_RAG_OCR_KEY) - ? state.ragOcrScanned - : undefined; - const caption = hasLocal(CHAT_RAG_CAPTION_KEY) - ? state.ragCaptionFigures - : undefined; + const { ocr, caption } = resolveVisionOverrides(); const result = activeScope.type === "kb" ? await uploadKnowledgeBaseDocument( diff --git a/studio/frontend/src/features/rag/components/vision-overrides.ts b/studio/frontend/src/features/rag/components/vision-overrides.ts new file mode 100644 index 0000000000..674484970b --- /dev/null +++ b/studio/frontend/src/features/rag/components/vision-overrides.ts @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + CHAT_RAG_CAPTION_KEY, + CHAT_RAG_OCR_KEY, + useChatRuntimeStore, +} from "@/features/chat"; + +function hasLocal(key: string): boolean { + if (typeof window === "undefined") return false; + try { + return window.localStorage.getItem(key) !== null; + } catch { + // Storage can be blocked outright (sandboxed context). These overrides are + // optional, so fall back to the backend defaults rather than failing the + // upload that asked for them. + return false; + } +} + +/** Ingest-time vision-pass overrides, sent only once the user has set them; + * otherwise backend env defaults own the policy. Shared by every upload path. */ +export function resolveVisionOverrides(): { + ocr: boolean | undefined; + caption: boolean | undefined; +} { + const state = useChatRuntimeStore.getState(); + return { + ocr: hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined, + caption: hasLocal(CHAT_RAG_CAPTION_KEY) + ? state.ragCaptionFigures + : undefined, + }; +} diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 7a3a3c58f7..f4ba98b1ce 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -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 ; case "api-keys": return ; + case "agents": + return ; case "about": return ; } @@ -222,6 +230,7 @@ export function SettingsDialog() { connections: null, data: null, "api-keys": null, + agents: null, about: null, }); @@ -249,9 +258,10 @@ export function SettingsDialog() { } }} className={cn( - // Cap at 880px but shrink to the viewport so it doesn't clip on - // iPad-portrait widths where a fixed width overflows. - "settings-surface !max-w-[min(880px,calc(100vw-2rem))] h-[560px] w-[min(880px,calc(100vw-2rem))] p-0 overflow-hidden", + // Cap at 960px but shrink to the viewport so it doesn't clip on + // iPad-portrait widths where a fixed width overflows. Height caps + // the same way so short viewports don't get a clipped dialog. + "settings-surface !max-w-[min(960px,calc(100vw-2rem))] h-[min(680px,calc(100dvh-2rem))] w-[min(960px,calc(100vw-2rem))] p-0 overflow-hidden", // Soft shadow, no outline ring. Pin --radius to the light value so // corner rounding matches in dark mode. "shadow-border rounded-xl ring-0 [--radius:1.1rem]", @@ -266,7 +276,9 @@ export function SettingsDialog() { {/* Keep tab content from expanding the dialog grid. */}
-