* fix(studio): set HIP_VISIBLE_DEVICES in apply_gpu_ids for ROCm training workers Training workers are spawned via multiprocessing spawn before detect_hardware() runs, so IS_ROCM is still False. If the user never set HIP_VISIBLE_DEVICES in their shell, _inherits_rocm_visibility is also False, leaving the worker with only CUDA_VISIBLE_DEVICES set. On ROCm hosts the HIP runtime honors HIP_VISIBLE_DEVICES over CUDA_VISIBLE_DEVICES, so the worker saw the full device list and torch raised "no usable HIP accelerator" on some setups. Fall back to probing torch.version.hip (a build-time attribute, safe to read before GPU init) to detect ROCm when neither IS_ROCM nor inherited env vars are available. Mirrors the existing fix in llama_cpp.py for llama-server subprocess GPU pinning. Fixes https://github.com/unslothai/unsloth/issues/5180 * test: tighten apply_gpu_ids ROCm fallback assertions Replace loose OR chain with exact string matches, split into three focused tests, and add a guard check for the try/except wrapper. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: detect ROCm unified memory (Strix Halo / AMD iGPU) via torch fallback amd-smi on iGPUs with shared/unified memory (e.g. Radeon 8060S on Strix Halo) reports only the dedicated VRAM slice (~512 MB) in its metric output, so get_visible_gpu_utilization() was returning usable_gb ≈ 0.35 GB instead of the full GTT pool (~128 GB). torch.cuda.mem_get_info() already surfaces the correct unified-pool size. Add _reconcile_rocm_unified_memory(): after amd-smi returns a valid result on a ROCm device, cross-check each device's vram_total_gb against torch.cuda.mem_get_info(). When torch reports a larger total, replace the amd-smi VRAM fields in-place. No-op for discrete AMD GPUs where the two sources agree. Fixes: "Falling back to all visible GPUs -- model may not fit" on AMD iGPU machines even when 100+ GB of unified memory is available. * Apply unified-memory reconciliation in get_gpu_utilization too The visible-GPU path was already corrected for AMD iGPUs with unified memory (Strix Halo / Radeon 8060S), but get_gpu_utilization was still returning the raw 512 MB amd-smi VRAM slice. Studio's /api/train/hardware endpoint and the live GPU monitor read from this primary path, so users continued seeing the wrong total even after auto_select_gpu_ids picked the right device. Refactor to share the per-device correction: * _apply_unified_memory_correction(metrics, torch_info) -- the actual replacement logic, in-place on a single metrics dict. * _reconcile_rocm_unified_memory(...) -- multi-device, iterates utilization["devices"] (visible-GPU path). * _reconcile_primary_rocm_unified_memory(...) -- single flat metrics dict (primary-GPU path), uses parent_visible_spec to pick the primary index, falls back to ordinal 0 when no visibility env is set. get_gpu_utilization now calls the primary reconciler under IS_ROCM, so both endpoints surface the real unified-memory pool on iGPUs while leaving discrete AMD GPUs untouched (torch_total <= smi_total -> no replace). * Use 'is not None' and log debug on torch.version.hip probe failures Two small follow-ups to the apply_gpu_ids ROCm fallback: 1. Match detect_hardware()'s 'getattr(torch.version, "hip", None) is not None' form so the entire codebase has one canonical 'this torch was built with HIP' check. On every shipping torch wheel hip is either None or a non-empty version string, so the new form agrees with the old bool() form on every real install. 2. Log the probe failure at debug level instead of swallowing it silently. The broad 'except Exception' is intentional (we never want apply_gpu_ids to crash a worker over a probe), but the silent pass made it impossible to tell whether the fallback was firing or being skipped. * fix(studio): honour HIP_VISIBLE_DEVICES in _get_parent_visible_gpu_spec before IS_ROCM is set When a user has HIP_VISIBLE_DEVICES set in their shell (e.g. "1" to select GPU 1) but detect_hardware() has not yet run in the Studio parent process, IS_ROCM is still False. _get_parent_visible_gpu_spec() was gated on IS_ROCM so it fell through to CUDA_VISIBLE_DEVICES (unset), saw all physical GPUs, and auto-selected index 0. apply_gpu_ids then overwrote HIP_VISIBLE_DEVICES with "0", making the intended GPU invisible to ROCm torch in the worker, which triggered the "no usable HIP accelerator" error (issue #5180). Apply the same _inherits_rocm_visibility pattern already used in apply_gpu_ids: check for HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES in the environment regardless of IS_ROCM so the correct GPU index is preserved. * fix(install): harden AMD ROCm GPU detection for multi-GPU and env-filtered setups The previous rocminfo awk pattern could miss discrete GPUs on machines where HIP_VISIBLE_DEVICES/ROCR_VISIBLE_DEVICES is used to mask an integrated GPU — the env vars filter rocminfo output but may not propagate into the install script subprocess, causing detection to fail entirely. Two changes: - Tighten rocminfo pattern from /gfx[0-9]/ && !/gfx000/ to /gfx[1-9][0-9]/ — simpler and correctly excludes the CPU agent (gfx000) without a negative lookahead - Add sysfs KFD topology fallback: reads /sys/class/kfd/kfd/topology/nodes/*/gpu_id which is a kernel-level view unaffected by HIP_VISIBLE_DEVICES or ROCR_VISIBLE_DEVICES Fixes detection failure reported in Discord by Chains (gfx1201 + iGPU machine where env var exclusion of the iGPU caused rocminfo to return no usable device). * Fix KFD sysfs awk fallback to read properties file The fallback added by this PR reads /sys/class/kfd/kfd/topology/nodes/*/gpu_id files but matches the literal token 'gpu_id' against their content. Those files contain only a single decimal value (e.g. '0' for CPU agents, '50432' for GPU agents), so the regex never matches and 'found' stays 0, making the fallback a no-op on every host. The properties file in the same directory contains key/value lines like 'gpu_id 50432' which is what the existing awk pattern expects. Reproduced with a synthetic sysfs layout: against gpu_id files awk exits 1; against properties files awk exits 0 when any node reports gpu_id > 0. * fix(setup.ps1): detect AMD ROCm GPU on Windows, bring to parity with setup.sh setup.ps1 only checked nvidia-smi and fell straight to "gpu: none" on AMD machines. setup.sh already probed rocminfo/amd-smi/hipconfig/hipinfo. Add three-tier detection mirroring install_llama_prebuilt.py's detect_host(): 1. hipinfo: gcnArchName in output confirms a real HIP GPU (not just SDK) 2. amd-smi list: "GPU: <digit>" data rows as fallback 3. WMI Win32_VideoController: last resort -- detects AMD GPU even without HIP SDK, then guides user to install it rather than silently going CPU Also corrects the "none" message to mention AMD ROCm alongside NVIDIA so users with AMD hardware understand the requirement. Fixes: rohit-style install where Strix Halo (Radeon 8060S) showed "gpu: none" even with the HIP SDK present. * fix(install.ps1): detect AMD ROCm GPU on Windows, bring to parity with setup.ps1 install.ps1 had the same nvidia-smi-only GPU detection as setup.ps1 before the setup.ps1 fix. Applies the same three-tier AMD detection: 1. hipinfo: gcnArchName confirms real HIP GPU 2. amd-smi list: GPU data rows as fallback 3. WMI Win32_VideoController: detects AMD GPU without HIP SDK and guides user to install it Fixes: install.ps1 showing "gpu: none" while setup.ps1 correctly showed "AMD GPU detected" on the same machine (reported by rohit, RX 7600 XT). * fix(install.ps1): suppress 'No NVIDIA GPU detected' when AMD GPU is present * feat: add Windows AMD ROCm PyTorch wheel installation install_python_stack.py: - Add _ROCM_WINDOWS_WHEEL_BASE and _ROCM_WINDOWS_RELEASES constants pointing to AMD repo.radeon.com (ROCm 7.2 -> torch 2.9.1+rocm7.2.1) - Extend _ensure_rocm_torch() with a Windows branch: detects ROCm via _has_rocm_gpu() / _detect_rocm_version(), requires Python 3.12 (cp312 is the only ABI AMD publishes for Windows), installs the direct wheel URL from repo.radeon.com install.ps1: - Capture ROCmVersion during AMD detection via hipconfig --version / amd-smi version (needed for wheel URL selection) - After Get-TorchIndexUrl, add an AMD wheel override block: when HasROCm and Python 3.12 detected, set ROCmTorchWheelUrl to AMD wheel URL - Expand torch install branch to handle ROCmTorchWheelUrl with uv pip install --force-reinstall --no-cache-dir * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: also install torchvision and torchaudio from AMD Windows repo AMD publishes matching torchvision-0.24.1+rocm7.2.1 and torchaudio-2.9.1+rocm7.2.1 cp312 wheels at the same repo.radeon.com release folder. Install all three in both install.ps1 and install_python_stack.py Windows ROCm path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feat: add ROCm 7.1.1 Windows wheel mapping AMD uses a different version string for 7.1.1 wheels: 2.9.0+rocmsdk20251116 (date-tagged) instead of +rocm7.1.1. Adds the 7.1.1 release folder to both install.ps1 and install_python_stack.py so users with ROCm 7.1 get ROCm torch instead of falling back to CPU. * fix: install rocm_sdk_core and rocm_sdk_libraries_custom alongside torch The AMD Windows torch wheels declare rocm[libraries]==<ver> as a hard dependency. Without installing rocm_sdk_core and rocm_sdk_libraries_custom from the same AMD release folder, uv cannot resolve the dependency and fails with 'No solution found'. Include all 5 wheels in one install call. * fix: expand ROCm wheel array to scalars for Invoke-InstallCommand @array splatting inside a scriptblock only works when the native command is prefixed with '&'. Invoke-InstallCommand uses '& $Command' to run the block, so @ROCmAllWheelUrls was not being expanded. Extract to scalar variables $rw0-$rw4 which are captured correctly by the closure. * fix: use --no-deps for AMD Windows torch wheel install uv's resolver looks up rocm[libraries]==0.1.dev0 on PyPI during dependency resolution before downloading any wheels, and fails because the package doesn't exist on PyPI. --no-deps skips resolution entirely and installs all 5 AMD wheels directly. The GPU runtime dependency is satisfied by the HIP SDK, not a Python package. * fix: setup.ps1 and install_python_stack.py now install ROCm torch on Windows setup.ps1 was always setting CuTag='cpu' for non-NVIDIA hosts and installing cpu-only PyTorch, overwriting the ROCm torch installed by install.ps1. Adds the same AMD wheel selection logic (ROCm version detection, Python 3.12 check, 5-wheel install with --no-deps) to setup.ps1's torch install block. install_python_stack.py: remove IS_WINDOWS guard from _ensure_rocm_torch() call site so the Windows path in _ensure_rocm_torch() is reachable during 'unsloth studio update' as well. * fix: suppress manual-install warning when ROCm torch already present; fix progress counter - Gate the 'must be installed manually' warning on torch.version.hip being empty so it doesn't fire when our ROCm torch install succeeded - Update _TOTAL counter to include the 3 ROCm steps on Windows now that _ensure_rocm_torch() is called there (fixes 10/9 display) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feat: add rocm step display in setup.ps1; fix warning and progress counter - Add 'rocm' step after 'cuda' in setup.ps1 showing ROCm version or HIP SDK missing - Move ROCm version detection up to GPU detection block so it's available early - Suppress 'must be installed manually' warning when torch.version.hip is set - Fix _TOTAL counter to include ROCm steps on Windows (fixes 10/9 display) * fix: detect AMD SDK ROCm torch via __version__ when torch.version.hip is unset AMD's repo.radeon.com wheels (e.g. 2.9.0+rocmsdk20251116) do not set torch.version.hip, leaving it None. All three probes that relied solely on torch.version.hip now also check for 'rocm' in torch.__version__.lower(): - hardware.py detect_hardware(): IS_ROCM was never set, causing the studio to report 'Hardware detected: CPU' even after AMD wheels were installed and HIP DLLs were on PATH. - install_python_stack.py _ensure_rocm_torch(): skip-if-already-installed probe would always reinstall on subsequent runs. - install_python_stack.py Windows AMD warning: suppression check always failed, so the 'must be installed manually' note kept appearing after a successful AMD wheel install. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * perf: drop --no-cache-dir from AMD ROCm torch wheel installs uv caches downloaded wheels by default; passing --no-cache-dir forced a full redownload of the ~2 GB torch wheel on every install run. CUDA installs never had this flag -- AMD was the only path affected. * fix: use install-state flag instead of subprocess probe for AMD Windows warning Replace the subprocess torch probe in the post-install warning block with a module-level _rocm_windows_torch_installed flag set by _ensure_rocm_torch(). Subprocess re-import of torch is unnecessary and fragile -- the install function already knows whether it succeeded. * fix: hoist global declaration to top of _ensure_rocm_torch Python requires the global statement to appear before any assignment to the variable within a function. Moving it to the function top fixes the SyntaxError on line 354. * fix: pass AMD torch install status via env var to suppress false warning setup.ps1 now sets UNSLOTH_ROCM_TORCH_INSTALLED=1 after a successful AMD wheel install. install_python_stack.py reads this at the top of _ensure_rocm_torch() to skip both the subprocess probe and the warning -- no re-import of torch needed, and the warning message now correctly says 'could not be auto-installed' rather than 'must be installed manually'. * fix: register ROCm DLL directory before torch import on Windows Python 3.8+ ignores PATH for extension DLL loading on Windows; amdhip64.dll and other HIP runtime DLLs must be registered via os.add_dll_directory(). Without this, torch.cuda.is_available() always returns False on AMD ROCm Windows even when HIP_PATH is correctly set in system environment variables. Reads HIP_PATH / ROCM_PATH env vars first, then falls back to scanning common ROCm install roots (C:\Program Files\AMD\ROCm, F:\ROCm, C:\ROCm). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: remove hardcoded non-standard ROCm paths from DLL directory scan Only use HIP_PATH/ROCM_PATH (set by AMD installer) and the standard C:\Program Files\AMD\ROCm\<version>\bin location. Custom drive paths like F:\ROCm are user-specific and should not be hardcoded. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: prevent torchao overrides step from overwriting AMD ROCm torch torchao==0.14.0 in overrides.txt declares torch as a dependency. Without --no-deps, uv resolves torch from PyPI and installs 2.11.0+cpu on top of the AMD ROCm wheels (2.9.0+rocmsdk20251116). This was the root cause of 'Hardware detected: CPU' -- the AMD wheels were installed but then immediately overwritten by the overrides step. When _rocm_windows_torch_installed is True, add --no-deps to the overrides pip_install call so torchao is installed without pulling in CPU torch. * fix: add rocm_sdk namespace tarball to Windows ROCm wheel installs torch/_rocm_init.py calls `import rocm_sdk` at startup, which requires the rocm namespace tarball (rocm-*.tar.gz) in addition to the SDK wheel packages. This tarball was missing from both install.ps1 and setup.ps1, causing ModuleNotFoundError on first torch import. - Add rocm-0.1.dev0.tar.gz to ROCm 7.1.1 install (provides rocm_sdk namespace) - Add rocm-7.2.1.tar.gz + rocm_sdk_devel to ROCm 7.2.1 install - Install tarball in a dedicated step before main SDK/torch wheels - Switch to @array splatting in install.ps1 scriptblock for dynamic wheel count - Remove --no-cache-dir from Python-side ROCm wheel install (prevents ~2GB redownload) * feat: enable ROCm 7.2 torch install + warn on gfx1151 with ROCm < 7.2 Chigoma333 (AMD Radeon 8060S / gfx1151, Strix Halo) confirmed that ROCm 7.1 segfaults when tensors are moved to GPU, but ROCm 7.2 + torch 2.11.0+rocm7.2 works fully including training. Changes: - Uncomment (7,2): "rocm7.2" in _ROCM_TORCH_INDEX (was blocked by <2.11.0) - Add _ROCM_TORCH_PKG_SPECS dict with per-tag version bounds: rocm7.2 → torch>=2.11.0,<2.12.0; all older tags → <2.11.0 - Add _detect_amd_gfx_codes() helper that parses rocminfo output - Warn on gfx1151/gfx1150 (Strix Halo) when ROCm < 7.2 is installed, pointing users at the known segfault and recommending upgrade - install.sh get_torch_index_url(): enable rocm7.2 case (previously capped to rocm7.1), cap unknown future tags to rocm7.2 - install.sh: override TORCH_CONSTRAINT to >=2.11.0,<2.12.0 when rocm7.2 index is selected, so pip can actually resolve torch 2.11.0 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: prefer Python 3.12 for AMD ROCm users when 3.13 is also installed After GPU detection, if ROCm HIP SDK is found and the selected Python is not 3.12, run a second pass to locate a 3.12 install via py.exe and PATH (catches uv-managed installs). Switch $DetectedPython to 3.12 so the venv is created with a compatible interpreter for the cp312-only AMD Windows torch wheels. NVIDIA and Intel GPU paths are unaffected -- the re-detection block only runs when $HasROCm is true. Fixes: #5301 * fix: also check uv-managed Python 3.12 for AMD ROCm #5301 * fix: hide amd-smi console popups on Windows, guard torch.distributed.is_initialized for ROCm #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: suppress remaining console popups on Windows, patch torch.distributed.is_initialized for ROCm #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: stub all missing torch.distributed attrs for ROCm Windows wheel #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: inject torch.distributed stub when C backend missing in ROCm Windows wheel #5301 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm/windows): pre-stub torch._C._distributed_c10d + raise amd-smi timeout Two fixes for Windows ROCm regressions reported by electroglyph on #5301: 1. worker.py — torch.distributed stub now fires unconditionally on Windows The previous stub only injected sys.modules in the except branch, meaning it was silently skipped when `import torch.distributed` happened to succeed (the C backend is lazily resolved). The crash then hit later when transformers/trl triggered the lazy load. Fix: on win32 we pre-populate sys.modules['torch._C._distributed_c10d'] AND set the attribute on the torch._C extension module *before* attempting the import, covering both the early-ImportError and lazy-load failure modes. 2. amd.py — increase amd-smi timeout from 5 s to 30 s on Windows (10 s Linux) amd-smi on Windows must cold-init the ROCm runtime on first invocation; 5 s was consistently too short, producing repeated 'Command timed out' warnings in the server log. 30 s gives enough headroom without blocking indefinitely on broken installs. 3. install.ps1 — widen Python 3.12 enforcement to ROCmGpuLabel (WMI-only path) Users whose HIP SDK is not on PATH were detected via WMI but not switched to Python 3.12 before the install started, causing a second pass. Guard now fires on (HasROCm -or ROCmGpuLabel). * fix(rocm): guard c10d stub, fix TorchIndexFamily for 7.1, clean dead code + comments - worker.py: wrap c10d stub injection in `if _c10d_key not in sys.modules` so Windows NVIDIA users with a real torch.distributed are never affected - install.ps1: fix Get-TauriTorchIndexFamily receiving hardcoded "rocm7.2" even when ROCm 7.1 wheels are installed; now branches on $ROCmVersion - main.py: remove dead `import ctypes as _ctypes` (ctypes is never called) - hardware.py, install_python_stack.py, worker.py, install.ps1: shorten verbose multi-line comment blocks throughout - tests: update 4 stale assertions that expected rocm7.2 to be absent/capped * fix(tests): match windows AMD warning assertion to actual source string * chore: trim verbose comment blocks across all ROCm-related files * fix: guard reconcile call against None numeric_ids; add torchvision lower bounds * fix(install.ps1): recreate venv with Python 3.12 after ROCm switch Venv was created with 3.13 before GPU detection ran; switching $DetectedPython to 3.12 had no effect since $VenvPython still pointed to the 3.13 interpreter inside the already-created venv. * ux: detect AMD GPU before Python selection to avoid double venv creation - Early hipinfo + WMI probe runs before Find-CompatiblePython so Python 3.12 is selected upfront when AMD is detected; venv is now created exactly once instead of 3.13 then immediately 3.12. - Post-venv recreation block replaced with a simple warning for the rare case where AMD was missed by the early probe. - setup.ps1: show venv's actual Python version (e.g. 3.12) instead of the system Python found by the pre-activation search (was showing 3.13). * fix(rocm/win): auto-stub all _distributed_c10d symbols via PEP-562 __getattr__ The bare ModuleType stub caused ImportError when torch._dynamo was imported (triggered by trainer.py accessing torch._dynamo.config at load time). torch._dynamo pulls in torch.distributed.fsdp._flat_param which does: from torch._C._distributed_c10d import FakeProcessGroup and potentially other symbols. Adding module __getattr__ auto-creates a stub class for any missing symbol so all such imports succeed without enumerating every individual symbol. Applied to both the primary stub and the fallback stub in the except branch. * chore: trim c10d stub comment * fix(rocm/win): auto-stub missing torch.distributed attrs (Store, ProcessGroup, …) * fix(rocm/win): pre-stub fsdp submodules in sys.modules; fix __getattr__ subpackage clash * feat(rocm/win): arch-aware wheel selector always picks newest ROCm release Replace HIP-SDK-version-gated wheel selection with GPU arch-based logic. Select-ROCmWheelRelease (PS) and _select_windows_rocm_release (Python) map gcnArchName → minimum ROCm version, then pick the newest available release that satisfies it (currently always rocm-rel-7.2.1 for any supported GPU). Wheels bundle their own ROCm runtime so the installed HIP SDK 7.1 does not prevent using 7.2.1 wheels on gfx1200 (RX 9060 XT) and similar RDNA 4 GPUs. Also installs the bitsandbytes Windows ROCm continuous-release wheel and sets BNB_ROCM_VERSION=72 in worker.py before ML imports so bnb loads the libbitsandbytes_rocm72.dll that ships in that wheel. * fix(rocm/win): stub class metaclass for ProcessGroup.BackendType; amd-smi circuit breaker torchao.float8.inference accesses ProcessGroup.BackendType as a class-level attribute. Plain type() stubs have no __getattr__ on the metaclass so this raises AttributeError. Introduce _StubClassMeta whose __getattr__ returns child stub classes, fixing the torchao import chain. Add an amd-smi circuit breaker in amd.py: after 3 consecutive failures the module stops spawning the process, eliminating the repeated Windows UAC / DiskPart elevation prompts caused by polling a non-functional amd-smi. Also guard BNB_ROCM_VERSION=72 behind a DLL existence check so bitsandbytes fails with its own detection message rather than a harder "DLL not found" when the Windows ROCm bnb wheel is not yet installed. * fix: stub __members__ so torchao float8 enum check doesn't crash on ROCm Windows torchao.float8.inference accesses ProcessGroup.BackendType.__members__ expecting a Python Enum registry dict. _StubClassMeta.__getattr__ was blocking all dunder attributes, causing AttributeError. Return {} for __members__ specifically so the isinstance/iteration checks pass cleanly. * fix: stub distributed tensor/functional_collectives to prevent missing C++ op crash on ROCm Windows torch._dynamo.trace_rules eagerly loads torch.distributed.tensor at import time, which pulls in _functional_collectives.py. That file registers Meta kernels for _c10d_functional C++ ops, but those ops are only registered by torch._C._distributed_c10d — a C extension absent from ROCm Windows wheels. Pre-stubbing the affected modules in sys.modules prevents the real import chain from running and avoids the "operator does not exist" crash. * fix: give mod stubs __path__ and pre-stub _tensor to fix 'not a package' import error _make_mod_stub now sets __path__=[] so Python treats stub modules as packages. Without it, any import of a submodule raises "is not a package". Also pre-stub torch.distributed._tensor and its submodules so that _tensor/__init__.py (which re-exports from torch.distributed.tensor) never runs and torchao's `from torch.distributed._tensor import DTensor` gets a harmless stub instead of crashing. * fix: stub torch.ops._c10d_functional namespace with hashable op sentinels torchao.dtypes.nf4tensor uses _c10d_functional ops as dict keys at import time (all_gather_into_tensor.default, wait_tensor.default) and torch.ops.c10d.scatter_.default. None of these ops are registered on ROCm Windows because torch._C._distributed_c10d (the C extension) doesn't ship. Replace the whole _c10d_functional namespace with a custom stub whose ops return hashable .default objects, so dict-key construction doesn't crash. Also inject a scatter_ stub into torch.ops.c10d if it's missing. * fix: stub entire torchao package on ROCm Windows instead of individual ops torchao is not supported on ROCm Windows and its import chain transitively requires torch._C._distributed_c10d (absent from the ROCm Windows wheel). Rather than stub each missing op one by one, stub the whole torchao package upfront. Unsloth uses bitsandbytes for quantization, not torchao, so this has no functional impact. transformers gracefully handles an importable-but- empty torchao by disabling TorchAoHfQuantizer. * fix: set __spec__ on mod stubs so importlib.util.find_spec doesn't raise Manually-injected sys.modules entries have __spec__=None by default. importlib.util.find_spec() raises ValueError when it finds a module in sys.modules with __spec__=None (transformers.utils.import_utils hits this when checking if torchao is available). Give every stub a minimal ModuleSpec(name, loader=None, is_package=True) to satisfy find_spec. * fix: add meta path finder to auto-stub subpackages of stub modules `import torchao.prototype` goes through the import machinery, not __getattr__, so an empty __path__ means ModuleNotFoundError. Rather than list every submodule explicitly, register a MetaPathFinder that intercepts any import whose parent is one of our stubs (detected by loader=None in the parent's ModuleSpec). Real installed packages always have a SourceFileLoader so they are never intercepted. Also register child stubs in sys.modules from __getattr__ as a belt-and-suspenders measure. * fix: use _unsloth_stub sentinel instead of loader=None for stub detection The import machinery overwrites module.__spec__ with the spec returned by find_spec (which has loader=_StubSubpackageLoader, not None), so the loader=None check broke for second-level subpackages. Switch to a custom _unsloth_stub object identity sentinel set directly on each stub module -- it survives __spec__ being replaced and correctly identifies stubs at any depth (torchao.prototype.safetensors, etc.). * refactor(rocm/win): switch to repo.amd.com arch-aware index, remove stubs AMD recommends repo.amd.com/rocm/whl/{arch}/ as the Windows ROCm wheel source. These wheels bundle their own ROCm runtime, support all Python versions (not just cp312), and include the full torch._C extension set (including _distributed_c10d) that the old repo.radeon.com wheel omitted. Changes: - install.ps1: remove Select-ROCmWheelRelease + hardcoded cp312 wheel URLs; remove Python 3.12 forced-preference logic; install via --index-url repo.amd.com/rocm/whl/{arch-family}/ - studio/setup.ps1: same -- remove Select-ROCmWheelRelease, switch to repo.amd.com arch-aware index URL - studio/install_python_stack.py: replace _ROCM_WINDOWS_RELEASES / _select_windows_rocm_release with _windows_rocm_index_url() using the _GFX_TO_AMD_INDEX_ARCH map; drop Python 3.12 restriction - studio/backend/core/training/worker.py: remove all stub machinery (_make_mod_stub, _StubSubpackageFinder, _StubSubpackageLoader, _StubClassMeta, torchao/fsdp/dtensor stubs, _c10d_functional ops stubs, BNB DLL detection) -- no longer needed with new wheel source * fix(rocm/win): restore _distributed_c10d + torchao stubs; fix BNB install repo.amd.com torch wheels also omit torch._C._distributed_c10d on Windows (RCCL is not shipped on Windows). torch/distributed/__init__.py imports from it unconditionally at module level, so the stub must land in sys.modules before any torch.distributed import. torchao (pulled in by transformers.quantizers) walks torchao.float8.distributed_utils -> torch.distributed._functional_collectives -> distributed_c10d at import time. Stubbing torchao up-front short-circuits that chain. worker.py: - Restore _make_mod_stub / _StubSubpackageFinder / _StubSubpackageLoader - Restore _StubClassMeta for ProcessGroup.BackendType attribute access - Restore _distributed_c10d stub with __getattr__ (Windows only) - Restore torchao stubs (5 modules, Windows only) install_python_stack.py: - BNB AMD wheel install was inside the early-return branch that fires when torch is already a ROCm build (installed by install.ps1). Move BNB install outside that branch so it always runs on Windows ROCm — the PyPI bitsandbytes has only CUDA DLLs and fails to load on ROCm. * worker: remove _distributed_c10d stub; stub only torchao The installed torch/distributed/__init__.py from repo.amd.com (torch==2.10.0+rocm7.12.0) is now properly guarded with `if is_available():`, so `import torch.distributed` alone is safe. The crash only comes via torchao's import chain: torchao.float8.distributed_utils → torch.distributed._functional_collectives (unguarded import) → torch.distributed.distributed_c10d → torch._C._distributed_c10d ← absent on Windows ROCm Stubbing torchao short-circuits the chain entirely. No need to stub _distributed_c10d. Remove _StubClassMeta and the _c10d stub block; keep only _make_mod_stub + _StubSubpackageFinder + torchao seeds. * fix: BNB AMD wheel skipped + torch.compile segfault on Windows ROCm install_python_stack.py: the UNSLOTH_ROCM_TORCH_INSTALLED=1 early-return path (set by setup.ps1 when it installed torch itself) returned before ever reaching the AMD BNB prerelease wheel install. The PyPI bitsandbytes==0.49.x ships only CUDA DLLs, so loading it on ROCm fails with "libbitsandbytes_rocm72.dll not found". Now installs the AMD Windows BNB wheel before returning on that path too. worker.py: torch._grouped_mm crashes on gfx1200 (null HIP kernel pointer, 0xC0000005) when torch.compile's JitDecomp system dispatches it during the first forward pass. Detect Windows ROCm via torch.version.hip (already in sys.modules from section 1e) and set TORCHDYNAMO_DISABLE=1 to bypass the broken kernel dispatch. * fix: BNB AMD wheel install fails uv wheel filename check The bitsandbytes continuous-release wheel is intentionally mismatched: filename encodes 1.33.7.preview (= 1.33.7rc0 in PEP 440) but wheel metadata reports 0.50.0.dev0. uv rejects this by default. Introduce _install_bnb_windows_rocm() helper that sets UV_SKIP_WHEEL_FILENAME_CHECK=1 only for this specific install, then restores the previous env value. Both BNB install call sites (the UNSLOTH_ROCM_TORCH_INSTALLED early-return path and the normal Windows ROCm path) now use this helper. * worker: patch _grouped_mm CUDA dispatch on Windows ROCm (gfx1200 null kernel) TORCHDYNAMO_DISABLE=1 stopped the compiler frontend but not the autograd JitDecomp system, which also dispatches _grouped_mm and hits the same null HIP kernel crash (0xC0000005). Verified that torch.library.Library("aten","IMPL").impl("_grouped_mm", fn, "CUDA") successfully overrides the broken HIP kernel with a Python mm fallback on torch==2.10.0+rocm7.12.0. Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor The fallback handles both the simple case (offs=None → torch.mm) and the grouped case (offs provided → split self by offsets, multiply each group against the corresponding slice of mat2, then cat results). Keep _WINDOWS_ROCM_GROUPED_MM_LIB alive at function scope to prevent the C++ dispatch registration from being freed by GC. * worker: fix torchao stub — return stub classes not modules for isinstance() peft/tuners/lora/torchao.py does: from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) The stub __getattr__ was returning stub modules, which isinstance() rejects with "arg 2 must be a type, a tuple of types, or a union". Add _StubTypeMeta metaclass whose __instancecheck__ always returns False, and _make_stub_type() to create stub classes via it. Change _make_mod_stub __getattr__ to return stub classes instead of stub modules for leaf attribute access, so isinstance() gets a valid type and returns False. _StubSubpackageFinder still handles import-style subpackage creation (those still need module objects in sys.modules); __getattr__ only fires for from-import or direct attribute access, which are the isinstance paths. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests: add coverage for Windows ROCm install paths and worker patches Add conftest.py to fix pre-existing sys.path issue that prevented test_rocm_support.py from running at all (install_python_stack.py imports from backend.utils.wheel_utils which needs studio/ on sys.path). New test classes cover everything added in this session: - TestWindowsRocmIndexUrl: arch → AMD pip index URL mapping (gfx120X-all, gfx1151, gfx1150, gfx110X-all, unknown → None, trailing slash) - TestDetectWindowsGfxArch: hipinfo output parsing, missing/timeout/bad returncode/no-gcnArchName paths - TestInstallBnbWindowsRocm: UV_SKIP_WHEEL_FILENAME_CHECK set+restored, env restored on exception, no-op when URL missing - TestRocmTorchInstalledEnvVar: UNSLOTH_ROCM_TORCH_INSTALLED=1 skips pip_install, calls _install_bnb_windows_rocm, sets flag - TestWorkerWindowsRocmPatches: _grouped_mm CUDA dispatch override, offs/grouped variant handling, GC-prevention sentinel, _StubTypeMeta __instancecheck__, _StubSubpackageFinder registration, torchao key submodule pre-stubbing, TORCHDYNAMO_DISABLE guard - TestRocmTorchPkgSpecs: rocm7.2 torch 2.11.x spec, default <2.11 cap, 3-tuple shape, _GFX_TO_AMD_INDEX_ARCH RDNA4/3.5/3 coverage * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tests: fix encoding, IS_WINDOWS patching, and wrong assertion - Add encoding="utf-8" to all read_text() calls (54 occurrences) so tests pass on Windows where the default codec is cp1252 and source files contain UTF-8 emoji (e.g. ⚠️ in install_python_stack.py) - Add @patch.object(stack_mod, "IS_WINDOWS", False) to Linux-path TestEnsureRocmTorch tests so they reach the Linux code path when run on a Windows machine instead of short-circuiting into the Windows branch - Fix test_grouped_mm_patch_guarded_by_windows_and_hip_check: the source uses getattr(_torch_for_rocm, "version", None) not torch.version, so check for '"version"' and '"hip"' substrings instead 137 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: pin BNB_ROCM_VERSION=72 for torch==2.11.0+rocm7.13.0 compatibility AMD's pip index now ships torch==2.11.0+rocm7.13.0 (ROCm 7.13). bitsandbytes auto-detects HIP 7.13 from torch.version.hip and looks for libbitsandbytes_rocm713.dll, which the AMD Windows prerelease wheel does not ship (it only ships rocm72.dll), causing a load error at training start. Fix: - worker.py section 1f: set BNB_ROCM_VERSION=72 (via setdefault) before section 2 ML imports, so bitsandbytes always loads rocm72.dll on Windows ROCm - install_python_stack.py: set BNB_ROCM_VERSION=72 in _install_bnb_windows_rocm() for any post-install imports; update comment to document root cause - tests: 4 new assertions covering the fix (141 passed, 2 skipped) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: detect BNB ROCm DLL suffix dynamically instead of hardcoding '72' BNB_ROCM_VERSION was pinned to '72' which works today (AMD wheel ships rocm72.dll) but would break again if AMD ships a future wheel with a different DLL suffix (e.g. rocm713.dll). Add _detect_bnb_rocm_dll_ver() to install_python_stack.py: scans the installed bitsandbytes package dir for libbitsandbytes_rocm{VER}.dll using importlib.util.find_spec (no BNB import needed) and returns the suffix. '72' remains the fallback when detection fails. Apply the same detection inline in worker.py section 1f. Both paths still respect a pre-set BNB_ROCM_VERSION (caller override wins). Tests: +8 cases covering detection logic and fallback (147 passed, 2 skipped). * fix: patch torch.distributed stubs in server process for Windows ROCm On Windows ROCm, torch.distributed ships without process-group helpers (is_initialized, is_available, get_rank, get_world_size). The worker subprocess already patches these in section 1e, but the main server process calls _determine_attention_impl_for_gpu_estimate() which calls unsloth's resolve_attention_implementation() → is_initialized(), causing: "Could not resolve attention implementation for '...': module 'torch.distributed' has no attribute 'is_initialized'" Fix: patch the missing attrs onto torch.distributed at the top of _determine_attention_impl_for_gpu_estimate, matching the same stubs already applied in worker.py section 1e. No-ops on Linux/CUDA where torch.distributed is fully populated. * fix: gate _grouped_mm dispatch patch on HIP < 7.13 AMD fixed the gfx1200 null HIP kernel in ROCm 7.13 (torch 2.11+). Users on the new wheel now get the real GPU _grouped_mm kernel for MoE workloads instead of the Python mm fallback. Changes: - worker.py: add _hip_ver_at_least() helper; wrap full _grouped_mm patch in `if not _hip_ver_at_least(7, 13):` with else branch that logs the skip reason; update section-1f comment to document the fix - test_rocm_support.py: add 5 tests covering the helper definition, the (7, 13) gate expression, the else branch, the skip log message, and the AMD-format version string parsing (.split(".")[:2]) Verified: torch==2.11.0+rocm7.13.0 — 3D batch and grouped (offs) variants both succeed; null crash only present on rocm7.12 and earlier. * fix: stub is_torchelastic_launched on torch.distributed for Windows ROCm resolve_attention_implementation calls is_torchelastic_launched() which does not exist in the incomplete torch.distributed shipped with the Windows ROCm wheel, causing a warning on every model config load in the server process. Add it to the stub table alongside the four helpers already patched in _determine_attention_impl_for_gpu_estimate. Also adds two tests: one confirming the new stub and one confirming all five core distributed helpers are covered. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: explicit warnings on AMD ROCm arch/version fallbacks + Fast-Install arg order setup.ps1: - Fix Fast-Install argument order: packages before flags, consistent with all other Fast-Install calls in the file (was: Fast-Install --force-reinstall --index-url $url torch ...) (now: Fast-Install torch torchvision torchaudio --force-reinstall --index-url $url) - Add explicit [WARN] substep when $HasROCm is true but arch mapping fails: - GPU arch detected but not in supported wheel list → names the arch and lists supported families so user knows exactly what to report - HIP SDK present (amd-smi path) but gcnArchName unreadable → instructs user to re-install the HIP SDK; previously fell back silently to CPU install.sh: - Add [WARN] to stderr before silent CPU fallback when AMD GPU is confirmed (rocminfo/amd-smi) but ROCm version cannot be read from any source (amd-smi, /opt/rocm/.info/version, hipconfig, dpkg, rpm) - Add [WARN] to stderr when ROCm version is too old (< 6.0) with upgrade link install.ps1 and setup.sh: no changes needed (already handle these paths correctly) * fix: robust gfx arch detection for Strix Halo / HIP-runtime-only installs Covers users who have the HIP runtime (amd-smi available) but not the full HIP SDK (no hipinfo), which is common on Strix Halo iGPU systems. Without this, $ROCmGfxArch stays null and the installer silently falls back to CPU-only PyTorch despite a working GPU. Detection waterfall (setup.ps1 + install.ps1): 1. hipinfo gcnArchName -- full HIP SDK (existing, unchanged) 2. amd-smi list gfx pattern -- newer amd-smi versions embed arch 3. amd-smi static --asic -- ROCm 6+ ASIC details with GFX target 4. UNSLOTH_ROCM_GFX_ARCH env -- manual override escape hatch 5. GPU name → arch table -- best-effort from marketing name: 890M / Strix Halo → gfx1151 (RDNA 3.5 iGPU, Strix Halo) 880M / Strix Point → gfx1150 (RDNA 3.5 iGPU, Strix Point) 780M / Phoenix → gfx1103 (RDNA 3 iGPU) RX 7900/7800/7700 → gfx1100 (RDNA 3 desktop) RX 9070 XT / 9080 → gfx1201 (RDNA 4) RX 9070 / 9060 XT → gfx1200 (RDNA 4) When arch is inferred from name, a Cyan substep tells the user to set UNSLOTH_ROCM_GFX_ARCH to skip inference on future installs. WMI block intentionally does not set $HasROCm (no runtime confirmation). Tests: 11 new tests in TestStrixHaloGfxArchDetection covering all five detection levels, WMI safety, and gfx regex in both ps1 files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: resolve hipinfo/hipconfig via HIP_PATH/ROCM_PATH when not on PATH AMD HIP SDK sets HIP_PATH on Windows but does not always add the bin directory to PATH. Get-Command hipinfo therefore silently fails and detection falls through to WMI, which cannot provide a gfx arch, leaving the user with a CPU-only PyTorch install and no warning. Changes: - setup.ps1 / install.ps1: before falling through to amd-smi, attempt to locate hipinfo.exe and hipconfig.exe under $env:HIP_PATH\bin (then $env:ROCM_PATH\bin) when Get-Command returns nothing - Emit a [WARN] with the resolved path and a one-liner to permanently fix PATH via SetEnvironmentVariable - Emit a [WARN] when HIP_PATH/ROCM_PATH is set but the exe is still not found (incomplete SDK install) - Emit a [WARN] with the first hipinfo output line when hipinfo runs but returns a non-zero exit code (e.g. "no ROCm-capable device detected") - 18 new tests in TestHipSdkEnvPathResolution; total 183 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feat: print HIP SDK path and full hipconfig version in terminal on AMD detection Both install.ps1 and setup.ps1 now emit substeps under the gpu step when AMD ROCm is detected: gpu AMD ROCm (gfx1200) HIP SDK: C:\Program Files\AMD\ROCm\7.1 hipconfig: 7.1.51803-d3a86bd04 Previously only the gpu label (e.g. "AMD ROCm (gfx1200)") was shown with no indication of where the SDK was found or which exact build was active. The full hipconfig build string (e.g. 7.1.51803-d3a86bd04 instead of just 7.1) is now stored in ROCmVersionFull and also used in setup.ps1's 'rocm' step label. 9 new tests in TestHipSdkDetectedSubstep; total 192 passed, 2 skipped * fix: Strix rocm7.1 segfault bypass + Ubuntu 24.04 HIP gcc-install-dir Issue 1 (install.sh): gfx1151/gfx1150 + ROCm 7.1 causes a segfault in torch._grouped_mm (moe_utils.py:167). The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so _amd_gpu_radeon=true silently lands on the broken combo. When Strix Halo/Point is detected and TORCH_INDEX_URL is rocm7.1, override to rocm7.2 PyTorch index, update TORCH_CONSTRAINT, and set _amd_gpu_radeon=false to bypass the Radeon repo entirely. Emits a clear [WARN] explaining the segfault and linking to the ROCm upgrade docs. Issue 2 (setup.sh): ROCm 7.x ships clang-20 which on Ubuntu 24.04+ picks /usr/lib/gcc/x86_64-linux-gnu/14/ (runtime dir, no C++ headers), causing 'cstdlib file not found' and a failed llama.cpp HIP build. Iterate gcc versions 14→11 to find the first install dir that has both runtime and /usr/include/c++/<ver> headers, then pass --gcc-install-dir to clang via CMAKE_HIP_FLAGS. Fix confirmed by h34v3nzc0dex (llama.cpp 417/417 clean). 11 new tests across TestStrixRocm71Override and TestSetupShGccInstallDir; total 203 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: BNB_ROCM_VERSION in server process + torch._C._distributed_c10d stubs Two errors visible in training logs on Windows ROCm: 1. Server process bitsandbytes crash: "Configured ROCm binary not found at libbitsandbytes_rocm713.dll" The installed BNB wheel ships rocm72.dll (not rocm713.dll). The training worker already sets BNB_ROCM_VERSION=72 via DLL detection but the server process (main.py) imported bitsandbytes before that ran. Fix: add the same DLL-scan + BNB_ROCM_VERSION assignment to main.py inside the existing win32 guard, before any downstream import can pull in bitsandbytes. 2. torch.distributed import failure: "No module named 'torch._C._distributed_c10d'; torch._C is not a package" torch._C is a C extension on Windows ROCm — Python cannot do submodule imports from it, so torch.distributed fails to import before our attribute stubs could ever run. Fix: inject empty ModuleType stubs for _distributed_c10d, _distributed_autograd and _distributed_rpc into sys.modules inside the win32 guard in hardware.py BEFORE importing torch.distributed, so the import succeeds and our attribute stubs take effect. 9 new tests in TestServerStartupRocmFixes; total 212 passed, 2 skipped * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(win32): populate distributed c10d stub with dummy symbols torch.distributed tries to `from torch._C._distributed_c10d import FakeProcessGroup` (and ProcessGroup, Work, Store, etc.). The previous empty ModuleType stub caused an AttributeError on those names. Populate every stub with a _Dummy class for each known symbol so the import chain completes silently on Windows ROCm where torch._C is a compiled extension and its _distributed_c10d submodule doesn't exist. Adds four new tests in TestServerStartupRocmFixes covering FakeProcessGroup, ProcessGroup, setattr population, and all three _distributed_* siblings. * fix(win32): distinguish HIP SDK installed vs GPU not ROCm-accessible Previously, when hipinfo was found but exited non-zero (e.g. "no ROCm-capable device detected"), both install.ps1 and setup.ps1 fell through to the WMI-label-only branch and printed "AMD GPU detected -- HIP SDK not found" -- factually wrong since the SDK binary is present. Add $HipSdkInstalled flag (set true when hipinfo binary is found, regardless of exit code). When HipSdkInstalled && !HasROCm: - Show "AMD GPU detected -- not ROCm-accessible (HIP <ver>)" instead - Explain this is a driver issue, not an SDK issue, with a link - Still run hipconfig version capture so version shows in output - CPU-only hint now says "GPU not ROCm-accessible" not "require HIP SDK" Also applies to setup.ps1 (same detection block, same branches). Adds TestHipSdkInstalledButDeviceInaccessible (11 tests). * fix(win32): scope ROCm workarounds to AMD hosts only Three Codex-flagged issues where Windows ROCm workarounds incorrectly applied to Windows CUDA (NVIDIA) machines: main.py (P1): BNB_ROCM_VERSION was set unconditionally on all win32 hosts. On NVIDIA, bitsandbytes sees BNB_ROCM_VERSION and looks for a ROCm DLL that doesn't exist, breaking bitsandbytes initialisation. Fix: gate the block on HIP_PATH/ROCM_PATH being present (ROCm hosts only). worker.py (P2): torchao stubs were seeded for all win32 runs, shadowing real torchao on Windows CUDA and silently disabling torchao quantization for NVIDIA users. Fix: gate on HIP_PATH/ROCM_PATH (win32 ROCm only). install_python_stack.py (P1): _detect_windows_gfx_arch() only checked shutil.which("hipinfo"), skipping the HIP_PATH/ROCM_PATH fallback that the PowerShell installers use. On installs where the HIP SDK bin dir is not on PATH, _ensure_rocm_torch() returned early without installing ROCm wheels or bitsandbytes. Fix: mirror the env-var fallback. * fix(linux): route Strix + ROCm 7.1 to AMD arch-specific index Instead of falling back to pytorch.org/rocm7.2, the Strix override now routes to repo.amd.com/rocm/whl/gfx1151/ (or gfx1150/) which serves torch 2.11.0+rocm7.13.0 -- AMD's build containing the actual _grouped_mm kernel fix, verified on real gfx1151 hardware by h34v3nzc0dex. This exercises the real GPU kernel path rather than the rocm7.2 workaround. UNSLOTH_AMD_ROCM_MIRROR can override the base URL for air-gapped installs. Also teaches _tauri_torch_index_family to recognise AMD arch-specific URLs (repo.amd.com/rocm/whl/gfx*) and return the rocm7.13 family label so _tauri_gpu_branch correctly classifies these installs as rocm. Suggested by h34v3nzc0dex based on hardware-verified probe results. * fix(studio/rocm): gate ROCm-only side-effects on active torch runtime Address five edge cases flagged during PR review: 1. studio/backend/main.py: BNB_ROCM_VERSION was set whenever HIP_PATH or ROCM_PATH was present in the environment. A Windows CUDA user who once installed the HIP SDK and reverted to a CUDA torch wheel still has those env vars set, so bitsandbytes would try to load libbitsandbytes_rocm72.dll against a CUDA torch and crash. Now probe torch.version.hip inside the env-var guard (worker.py already does this). 2. studio/backend/main.py: os.add_dll_directory returned handles were discarded. Per CPython docs, the directory leaves the DLL search list when the handle is garbage collected. Retain handles in module-level _ROCM_DLL_HANDLES list so they survive process lifetime. 3. studio/install_python_stack.py: _install_bnb_windows_rocm() returned None regardless of pip_install_try outcome, and the caller flipped _rocm_windows_torch_installed to True unconditionally. On a failed BNB install the post-install "manual install may be required" warning was suppressed and the user was misled. Helper now returns bool; caller gates on it. 4. studio/install_python_stack.py: _detect_windows_gfx_arch returned the raw capture group, so mixed-case hipinfo output ("Gfx1151") missed the lowercase keys in _GFX_TO_AMD_INDEX_ARCH and silently fell back to CPU torch. Lowercase the token. 5. studio/install_python_stack.py: UNSLOTH_ROCM_TORCH_INSTALLED=1 early- return trusted the env var even when the venv was wiped between runs. Subprocess-probe torch importability first; fall through to the full install path if the probe fails. Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py (adds one new test for case 5 fall-through). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): worker.py parity + don't roll back ROCm torch on bnb failure Addresses findings from a 10x reviewer pass on the prior fix commit: 1. studio/backend/core/training/worker.py (parity with main.py): - Gate the torchao stub block on torch.version.hip / 'rocm' in torch.__version__ instead of HIP_PATH / ROCM_PATH env-var presence. Same root cause as main.py: HIP SDK env vars stick around on CUDA hosts. - Add module-level Windows ROCm DLL registration block. Worker subprocesses inherit env vars but not the parent's add_dll_directory handles, so the first `import torch` in the worker could fail to find amdhip64.dll when HIP_PATH\bin is not on PATH. Mirrors main.py setup. Handles retained at module scope via _ROCM_DLL_HANDLES. - Promote _WINDOWS_ROCM_GROUPED_MM_LIB to module scope with `global` in run_training_process so the torch.library.Library registration survives past function return / mid-run garbage collection. - Harden _torch_has_hip() to also accept 'rocm' in torch.__version__ (AMD SDK / Radeon wheels may not set torch.version.hip). 2. studio/install_python_stack.py: - Don't roll back ROCm torch when bitsandbytes install fails. The prior commit gated _rocm_windows_torch_installed on _install_bnb_windows_rocm() returning True; if torch installed successfully but bnb failed, the flag stayed False and later install steps could overwrite ROCm torch with the generic CPU torch wheel. Set the flag after torch install; surface bnb failure as a separate warning instead. - _detect_windows_gfx_arch now probes in three tiers: UNSLOTH_ROCM_GFX_ARCH env-var override (matches the PowerShell installer), then hipinfo (PATH or HIP_PATH\bin), then amd-smi (`static --asic`, `list`). Without the amd-smi fallback, runtime-only Radeon installs without hipinfo on PATH made `studio update` return early and leave the venv on CPU torch. - Linux torch-already-rocm probe in _ensure_rocm_torch now matches the Windows probe shape: accepts torch.version.hip OR 'rocm' in torch.__version__ to cover AMD SDK / Radeon Linux wheels. 3. studio/backend/utils/hardware/hardware.py: - apply_gpu_ids() final-fallback torch probe accepts 'rocm' in torch.__version__ in addition to torch.version.hip, matching detect_hardware(). AMD SDK wheels could otherwise leak through with CUDA-only visibility masks on a spawned ROCm worker. Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py (no test changes needed; the probe shape that prints the hip version (or 'rocm' sentinel) preserves the existing non-empty-string contract). Not addressed in this commit (deferred or out of scope): - Tag drift / lemonade checksum (PR 5303 surface, not this PR). - install.sh rocm7.2.1 URL: small fix, separate. - install.ps1 / setup.ps1 'Radeon 8060S' marketing-name fallback table. - Strix Halo + ROCm 7.1 routing asymmetry in Python update path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): robustness pass - rocm tag normalisation, Strix routing parity, hardened detection Robustness pass on top of76137b2d. Four targeted fixes: 1. install.sh ROCm-tag routing normalisation. `rocm7.2.1` would route to https://download.pytorch.org/whl/rocm7.2.1 which does not exist (PyTorch publishes major.minor URLs only). Same for any future patch-level tag. Normalise every rocm{maj.min}* pattern to the bare {maj.min} index URL. 2. install.ps1 + studio/setup.ps1 marketing-name fallback. The gfx1151 row matched 890M / Strix Halo / HX 37x / HX 38x / AI 9 HX but not the actual retail name 'AMD Radeon 8060S Graphics' shipped by OEMs (Ryzen AI MAX+ 395). Add '8060S' to the regex. 3. install_python_stack.py Strix + ROCm 7.1 routing parity with install.sh. The shell installer reroutes Strix Halo / Point + ROCm 7.1 to repo.amd.com/rocm/whl/{gfx}/ (which serves torch 2.11.0+rocm7.13.0 with the upstream _grouped_mm fix). The Python `studio update` path only warned and still installed the broken generic rocm7.1 wheel. Mirror the override: detect gfx1151/gfx1150 on ROCm 7.1, route to the AMD per-gfx index, honour UNSLOTH_AMD_ROCM_MIRROR override. 4. _detect_windows_gfx_arch amd-smi parsing tightened. The amd-smi fallback added in the prior commit used a bare `\bgfx[1-9][0-9a-z]{2,3}\b` match against the lowercased stdout, which could pick up stray gfx references in warnings / device-name strings. Anchor on labelled lines first (Target_Graphics_Version, ASIC, Arch, gfx) and fall back to the bare match only when no labelled line is present. Tests: 231 passed, 1 skipped in tests/studio/install/test_rocm_support.py; sim_5301 23 cases pass (6 new sims for the Strix override + amd-smi parsing). * fix(studio/rocm): multi-GPU selection, Strix sibling handling, defensive cleanups Round 4 robustness pass based on 5 parallel Opus reviewers of head21773215. Seven items from across regression / edge-case / error-paths / architecture reviews: 1. studio/backend/main.py BNB gate: aligned with the broad ROCm check used everywhere else in this PR (torch.version.hip OR 'rocm' in __version__). AMD SDK / Radeon Linux wheels do not always populate torch.version.hip; without this, main.py would silently skip BNB_ROCM_VERSION while worker.py set it. 2. studio/install_python_stack.py _install_bnb_windows_rocm: init _ok = False before the try block. Without this, if pip_install_try itself raises (e.g. OSError on uv binary missing), the finally block restored env vars correctly but the subsequent `if not _ok:` raised UnboundLocalError, masking the original exception. 3. studio/install_python_stack.py _detect_windows_gfx_arch: - Rewrote to use re.findall (not re.search) on both hipinfo and amd-smi output, dedup tokens preserving order, and select via new _pick_visible_index() helper. - HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES (first comma entry, integer) now picks the right GPU on multi-AMD-GPU hosts. Out-of-range or non-int values fall back to the first GPU (matches detect_host behaviour in install_llama_prebuilt.py). 4. studio/install_python_stack.py Strix override now consults the runtime target before flipping: - Previous behaviour intersected gfx_codes with {gfx1151, gfx1150} and picked the first Strix arch, ignoring whether HIP_VISIBLE_DEVICES selected a non-Strix sibling (e.g. discrete RX 7900 in a mixed APU+dGPU box). Could install Strix-specific wheels onto a gfx1100 dGPU. - Now resolves the runtime gfx via _pick_visible_index() and only overrides when that runtime target is in the Strix set. 5. studio/backend/main.py + studio/backend/core/training/worker.py: ROCm version dir scan no longer sorts lexically. Previous sort placed "10.0" before "7.0" alphabetically, which would mis-prioritise ROCm 10.x bin dirs once AMD ships them. New _ver_key() splits on "." and sorts numerically with a string fallback. 6. install.sh Strix override URL: replaced ${var%/} (strips one trailing slash) with a while-loop that strips all trailing slashes, matching Python's .rstrip("/"). A user setting UNSLOTH_AMD_ROCM_MIRROR with "http://corp/whl///" no longer ends up with "http://corp/whl///gfx1151/" which strict pip proxies (artifactory, sonatype) 404 on. 7. studio/install_python_stack.py: bumped torch import probe timeout from 30s to 90s. PyTorch's lazy .so loading can take 60-90s on cold NFS or USB-backed venvs. The shorter timeout was producing a false "torch missing" classification and reinstalling a working ROCm torch. Tests: 231 passed, 1 skipped. sim_5301 30 cases pass (added 7 new sims for multi-GPU detection, Strix sibling handling, and _ok-init regression). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): worker BNB/grouped_mm broad gate, install.sh Strix visibility, runtime-only ROCm detection Round-5 robustness pass based on 20 parallel reviewers of head96b9e465. 1. studio/backend/core/training/worker.py - BNB version pin / dynamo disable / _grouped_mm fallback block was still gated on torch.version.hip alone despite the torchao stub block above already using the broad check. AMD SDK / Radeon Windows wheels (torch.__version__ contains "rocm" but torch.version.hip is None) silently skipped the Windows ROCm runtime patches. Aligned to the same broad check (8/20 reviewers). 2. studio/backend/core/training/worker.py - _hip_ver_at_least() now also parses the ROCm version out of torch.__version__ (e.g. "2.11.0+rocm7.13.0") when torch.version.hip is missing, so the kernel-fix gate is correct for SDK / Radeon wheels too. 3. studio/backend/core/training/worker.py - _grouped_mm_safe_impl with offs=None now picks torch.bmm/matmul for 3-D inputs instead of always calling torch.mm. The real _grouped_mm accepts 3-D batched matmul; the prior fallback raised "self must be a matrix" on MoE workloads (2/20). 4. studio/backend/main.py - dropped the HIP_PATH / ROCM_PATH env-var gate from the BNB block; probe torch directly. Runtime-only Radeon / AMD SDK Windows installs do not set those SDK env vars but still ship ROCm torch (5/20 reviewers). 5. install.sh - Strix override now collects every gfx token from rocminfo / amd-smi (in enumeration order), then indexes by HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non- Strix dGPU host where the user selected the dGPU does NOT get rerouted to the Strix per-gfx index. Mirrors the Python update path (5/20 reviewers). 6. install.sh - Strix detection chain now also probes `amd-smi static --asic`, matching the PowerShell installer (1/20). Closes the gap on runtime-only Strix hosts where `amd-smi list` does not surface a gfx token. 7. studio/install_python_stack.py - _has_rocm_gpu() now has the sysfs KFD topology fallback (/sys/class/kfd/kfd/topology/nodes/*/gpu_id), matching install.sh. On minimal package-managed installs without rocminfo / amd-smi GUI tools, `studio update` can now detect the GPU and repair the venv instead of returning early (2/20). 8. studio/install_python_stack.py - _detect_amd_gfx_codes() now falls back to `amd-smi list` and `amd-smi static --asic` when rocminfo is missing (2/20). Strix routing on runtime-only Radeon hosts now matches what install.sh has done for a while. 9. studio/install_python_stack.py - Strix override now applies even when has_hip_torch is True. The whole point of the override is to repair an existing broken torch.version.hip == "7.1" install; skipping the reinstall left users on the known _grouped_mm segfaulting stack (3/20). Tests: 231 passed, 1 skipped. sim_5301 30 cases pass. sim_cross 12 pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/rocm): code review hardening pass - main.py: numeric DLL sort (string sort picked rocm72 over rocm713); add basename() to regex; log warning on detection failure; log info when BNB_ROCM_VERSION is set (mirrors worker.py) - worker.py: explicit len-guard in _hip_ver_at_least() with warning logs instead of silent IndexError/ValueError swallow - hardware.py: isinstance(result, dict) guard before result.get() in _smi_query() to prevent AttributeError on non-dict backend returns - amd.py: round() before int() on parsed GPU IDs; log warning when truncation occurs (defensive against malformed amd-smi output) - setup.sh: quote --gcc-install-dir value in CMAKE_HIP_FLAGS so paths with spaces do not break the CMake argument - install.ps1, setup.ps1: apply colon-split + ToLower() to hipinfo gcnArchName match (consistent with each other and with setup.sh) - install.sh: tighten ROCm tag case patterns to explicit rocmX.Y|rocmX.Y.* to avoid unintended prefix matches * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/training): GPU OOM guard to prevent system freeze on VRAM exhaustion On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can cause a HIP driver hang that freezes the entire system rather than raising a recoverable Python exception. Two-part fix: - set_per_process_memory_fraction(0.90) caps the HIP/CUDA allocator at 90% of VRAM so PyTorch raises OutOfMemoryError before hitting the hardware limit, keeping the driver alive and the system responsive - top-level exception handler detects OOM errors by type and message and surfaces a clear actionable message to the UI (reduce max_seq_length, enable gradient_checkpointing, lower batch size) instead of the raw CUDA/HIP error string * fix(studio/rocm): OOM guard ROCm-only + unified memory, multi-GPU arch selection OOM guard (worker.py): - Scope to _hw.IS_ROCM only -- NVIDIA CUDA has a graceful OOM path and does not need the allocator cap - Detect unified memory by comparing torch VRAM against psutil system RAM; use 0.80 on unified-memory APUs (gfx1151 Strix Halo) where the GPU pool is carved from host RAM, 0.90 on discrete cards Multi-GPU arch selection: - install.ps1 / setup.ps1: replace -match (first hit only) with [regex]::Matches() to collect all gcnArchName entries, then index by HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES - install_python_stack.py: index into full token list before dedup so HIP_VISIBLE_DEVICES=2 on [gfx1100, gfx1100, gfx1151] resolves gfx1151 - install.sh: remove awk dedup from gfx token collection for same reason GCC multiarch (setup.sh): - Only append -linux-gnu when gcc -print-multiarch does not already return the full triple, fixing double-suffix on Ubuntu 24.04 * fix(tests): update ROCm version cap expectations from rocm7.1 to rocm7.2 Daniel's normalisation commit updated the cap from rocm7.1 to rocm7.2 since PyTorch now publishes that index and rocm7.2 ships torch 2.11.0. Test expectations were stale. * fix(tests): correct MLX smoke test losses_per_step assertion logging_steps=1 with max_steps=30 produces 30 loss entries, not 7. The assertion was stale from a previous config. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio/worker): detect unified-memory APU by GPU name not VRAM/RAM ratio The previous heuristic (VRAM > 50 % of system RAM) false-positived on discrete cards in low-RAM systems — e.g. RX 9060 XT 16 GB on a 16 GB or 24 GB machine would trip the unified-memory path and log "unified memory host" when it should say "discrete". AMD iGPUs (gfx1150/gfx1151 Strix Halo, Strix Point, etc.) expose names with a digit+M suffix ("AMD Radeon 890M"), while discrete cards use "RX NNNN [XT|XTX]" naming. Matching that suffix is reliable across all current ROCm-capable AMD consumer GPUs and does not require psutil. Also includes the device name in the log line to ease future debugging. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install/setup.ps1): force array on hipinfo gcnArchName parse to fix single-GPU arch truncation When [regex]::Matches() finds exactly one match, PowerShell's pipeline unwraps the result to a scalar string. Indexing a scalar string with [0] returns the first *character*, so a one-GPU system would parse gcnArchName "gfx1200" as "g", which is not in the supported arch map and triggers the CPU-only fallback. Wrapping with @() forces the result to remain an array regardless of match count. On a single-GPU machine the arch is now correctly read as "gfx1200" (or whatever the full name is) so the ROCm wheel index is selected. Reproducer: hipinfo exits 0 and outputs exactly one gcnArchName line. Without @(), $_hipAllArches = "gfx1200" (String); $_hipAllArches[0] = 'g'. With @(), $_hipAllArches = @("gfx1200") (Object[]); $_hipAllArches[0] = "gfx1200". * fix(studio/rocm): classify unified-memory APU via VRAM/RAM ratio, not arch list Replace the gcnArchName allowlist {gfx1150, gfx1151} with a psutil-based heuristic: unified APUs expose the entire system RAM as the HIP pool (ratio ≥ 0.90), discrete cards are well below that. No arch name required — future APUs classify correctly without code changes. Also removes the stale import re / \d[Mm]\b device-name regex that5d84704left behind, and logs vram/sys GiB for easier on-hardware verification. Addresses h34v3nzc0dex review: Radeon 8060S (gfx1151, 128 GiB unified) now correctly gets 0.80 cap instead of 0.90. * fix(studio/rocm): revert to gcnArchName for unified-memory APU classification VRAM/RAM ratio >= 0.90 false-positives on machines where discrete VRAM equals system RAM (e.g. RX 9060 XT 16 GB + 16 GB system RAM → ratio 1.0, incorrectly classified as unified → wrong 0.80 cap applied). gcnArchName is the correct signal: naming-independent, stable within a product family, and already parsed throughout this PR. Unified set is {gfx1150, gfx1151} (Strix Point + Strix Halo). * fix(studio/llama-prebuilt): resolve hipinfo via HIP_PATH/ROCM_PATH on Windows shutil.which("hipinfo") returns None when the HIP SDK bin dir is not on PATH -- the HIP SDK installer sets HIP_PATH/ROCM_PATH but does not always add the bin dir to PATH. This caused has_rocm=False in the prebuilt asset selector, so AMD ROCm machines got the CPU llama.cpp zip instead of the HIP one, silently running all chat inference on CPU. Add _resolve_exe() that falls back to %HIP_PATH%\bin and %ROCM_PATH%\bin when shutil.which() finds nothing, mirroring the same fallback already present in setup.ps1. * fix(studio/llama-prebuilt): pass --has-rocm from setup.ps1 to skip re-detection The Python prebuilt installer re-detects ROCm independently via shutil.which("hipinfo"), which fails when hipinfo is not on PATH (HIP SDK sets HIP_PATH but doesn't always add the bin dir to PATH). This caused has_rocm=False and downloaded the CPU llama.cpp zip even on confirmed AMD ROCm machines. setup.ps1 already performs reliable ROCm detection with its own HIP_PATH/ROCM_PATH fallback. Add --has-rocm flag to install_llama_prebuilt.py so setup.ps1 can forward its result directly, and pass it whenever $HasROCm is true. The Python script then overrides has_rocm=True in the HostInfo without re-probing. * fix(studio/llama-prebuilt): add HIP asset to simple-policy Windows path direct_upstream_release_plan (used by --simple-policy, which setup.ps1 always passes) only checked has_usable_nvidia on Windows and fell straight to CPU for AMD ROCm machines, ignoring has_rocm entirely. The --has-rocm override had no effect because the simple-policy code path never reached resolve_asset_choice where has_rocm was checked. Add an elif branch for has_rocm that tries the upstream HIP asset (llama-TAG-bin-win-hip-radeon-x64.zip) before falling through to the CPU fallback, consistent with the non-simple-policy path. * fix(studio/setup.ps1): auto-remove mismatched llama.cpp install kind When an existing llama.cpp install is the wrong kind for the current GPU (e.g. windows-cpu on an AMD ROCm machine that should have windows-hip), the prebuilt installer skips on tag match and never upgrades. Read install_kind from UNSLOTH_PREBUILT_INFO.json before invoking the installer and remove the directory if the kind doesn't match, forcing a fresh download of the correct variant. * fix(studio/setup.ps1): show live PyTorch install output in verbose mode for ROCm The ROCm torch reinstall (setup.ps1 phase) always silently captured output, so in --verbose mode the torch downgrade mid-install (2.11.0+rocm → 2.10.0 → 2.11.0+rocm) looked like the final state was 2.10.0. Match the CPU/CUDA blocks which show live uv output when $script:UnslothVerbose is set. * fix(rocm/windows): set ROCBLAS_TENSILE_LIBPATH for bundled rocblas.dll The llama.cpp ROCm prebuilt bundles rocblas.dll next to the binary but not the Tensile kernel library files it depends on at runtime (rocblas/library/TensileLibrary*.dat + *.hsaco). The bundled DLL searches for these files relative to its own location by default, i.e. <binary_dir>/rocblas/library/, which does not exist in the prebuilt install tree. This causes a silent crash on the very first GEMM (prefill) with no output from llama-server, seen by the caller as WinError 10054 / 10061. Model load and the single-token warmup pass because they use simpler code paths that do not trigger rocBLAS GEMM. Fix: set ROCBLAS_TENSILE_LIBPATH in the subprocess env to <HIP_PATH>/bin/rocblas/library so the bundled DLL finds the kernel files from the system ROCm installation. Uses setdefault so a user- supplied env var is never overwritten. No-ops on CUDA and CPU (no HIP_PATH) and on Linux (win32 branch only). Reproducer log: rocBLAS error: Cannot read .../Release/rocblas/library/TensileLibrary.dat rocBLAS error: Could not initialize Tensile host: directory_iterator: The system cannot find the path specified. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install.sh): restore gfx token dedup in Strix multi-GPU awk indexer536a54dfremoved the per-source `| awk '!seen[$0]++'` dedup from the _gfx_all collection step but left the indexer awk as bare NF, so on a mixed-arch host (e.g. dGPU gfx1100 + Strix iGPU gfx1151) where rocminfo emits each gfx token twice (Name: field + ISA triple), HIP_VISIBLE_DEVICES=1 indexed vals[1] = the second gfx1100 occurrence instead of gfx1151, triggering the Strix routing on the wrong GPU. Add !seen[$0]++ to the indexer awk so duplicate tokens from the same GPU collapse to one entry before the HIP_VISIBLE_DEVICES index is applied -- matching exactly what the Python side does with dict.fromkeys() in _detect_amd_gfx_codes(). The comment above the block ("skip duplicates") already documented this as the intended behaviour. * fix(studio/install): correct _TOTAL progress count on Windows base_total += 3 fired for all non-macOS platforms including Windows, but flash-attn (line 1620) and ROCm torch final (line 1705) are both guarded by 'not IS_WINDOWS and not IS_MACOS', so on Windows with torch enabled _TOTAL was 13 while only 11 _progress() calls actually execute. Split into +1 for the ROCm torch check (all non-macOS) and +2 for the two Linux-only steps, so Windows gets _TOTAL=11 and Linux gets 14. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install.ps1): enforce torch>=2.11.0 for gfx120X and Strix on Windows The AMD arch-specific index (repo.amd.com/rocm/whl/gfx120X-all/ and gfx1151/) publishes torch wheels from 2.7.1 through 2.11.0. Without a version floor pip can resolve to torch 2.10.0+rocm7.12 on RDNA 4 (gfx120X) or torch 2.10.0+rocm7.1 on Strix (gfx1151/gfx1150), both of which have a null-pointer crash in torch._C._grouped_mm (TheRock issues #5284 / #3284). torch 2.11.0+rocm7.13 contains the fix. Add $ROCmTorchFloor alongside $ROCmIndexUrl: set to torch>=2.11.0 for the two affected arch families, null for all others. Wire it into the uv pip install call so the broken wheels are never selected. * fix(rocm/windows): address Codex nits - deterministic DLL suffix, CUDA llama.cpp kind, HIP_VISIBLE_DEVICES arch indexing - install_python_stack.py / worker.py: _detect_bnb_rocm_dll_ver() and the inline worker probe now collect ALL libbitsandbytes_rocm*.dll suffixes and return max() by numeric value instead of stopping at the first glob hit. Filesystem glob order is not guaranteed; this ensures '713' always wins over '72' when both variants are present in the wheel. - setup.ps1 (expectedKind): add 'windows-cuda' branch so NVIDIA hosts are not treated as 'windows-cpu'. Previously an existing windows-cuda prebuilt was always considered a mismatch on non-ROCm machines, forcing an unnecessary re-download on every update. - setup.ps1 (amd-smi gfx arch): collect ALL gfx tokens from amd-smi list output in GPU order and honour HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES when selecting which arch to use. On mixed-arch AMD systems where the visible GPU is not the first enumerated one, this prevents installing an incompatible wheel index. Falls back to index 0 (same as before) when the visibility var is unset or is a comma-separated list. - test_rocm_support.py: add test_picks_highest_suffix_when_multiple_dlls to cover the multi-DLL case that was previously untested. * fix(rocm): misleading amd-smi log, BNB spec consistency, torch ceiling for AMD index amd.py: split 'returncode != 0 or not stdout' into two separate branches. Previously, exit-0 with empty output logged 'amd-smi returned code 0' (which reads as success, not a warning) and incorrectly incremented the circuit-breaker counter. Now: non-zero exit logs the code and counts toward the limit as before; empty stdout on exit 0 logs at DEBUG level and does not penalise the counter (amd-smi --json always emits at least [] on exit 0, so this branch is rare and is not a tool failure). main.py: replace spec.origin / os.path.dirname() with spec.submodule_search_locations to match install_python_stack.py and worker.py. For normal wheel installs both approaches reach the same directory, but using submodule_search_locations is the canonical way and handles editable bitsandbytes installs correctly. Also use max() by numeric suffix (same as the other two sites) instead of a sort-then-break loop. install.ps1: add <2.12.0 ceiling to the torch constraint for gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix). AMD actively publishes new versions on their per-arch index; without a ceiling, a future 2.12.0+rocmX.Y wheel would be pulled in automatically before being validated on these architectures. The ceiling matches the existing Linux install_python_stack.py constraint for the same arches. Bump both when 2.12.x is confirmed working. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm): torch floor in setup.ps1, torchvision pin for Strix, rocmsdk in _hip_ver_at_least setup.ps1: add \ (mirrors install.ps1) and derive \ from it. Previously the AMD index install called 'Fast-Install torch torchvision torchaudio --force-reinstall --index-url \' with no version constraint, so pip could resolve torch 2.10.0+rocm7.12 for gfx1151/gfx1200 -- the exact broken wheel the PR is meant to avoid. Now gfx120X and Strix enforce 'torch>=2.11.0,<2.12.0', matching install.ps1 and the Linux constraint. install_python_stack.py: pin torchvision and torchaudio in _strix_override_pkgs. The Strix Linux override uses --index-url (exclusive, no PyPI fallback); bare unversioned 'torchvision' and 'torchaudio' could resolve a build from AMD's index targeting a different torch major, causing ABI/version mismatches at runtime. Now pinned to '>=0.26.0,<0.27.0' and '>=2.11.0,<2.12.0' respectively, matching _ROCM_TORCH_CONSTRAINT['rocm7.2']. worker.py: extend _hip_ver_at_least to handle AMD SDK wheel version strings. The fallback regex r'rocm(\d+)\.(\d+)' cannot match '2.9.0+rocmsdk20251116' (no rocmX.Y component), so the function always returned False on SDK/Radeon wheels -- installing the Python _grouped_mm workaround on wheels that already have the working HIP kernel. Added a second check: if the version string contains '+rocmsdk', assume >= 7.13 (the rocmsdk format post-dates the gfx120X null-kernel fix) and skip the fallback. * fix(rocm): warn on OOB HIP_VISIBLE_DEVICES, bail on empty numeric_ids mask - setup.ps1: when HIP/ROCR_VISIBLE_DEVICES names an index beyond the detected GPU count, emit a yellow warning and fall back to GPU 0 instead of silently reading allGfxArches[-1] (wrong arch) - hardware.py _reconcile_primary_rocm_unified_memory: distinguish numeric_ids=None (no env var, use torch ordinal 0) from numeric_ids=[] (empty mask / HIP_VISIBLE_DEVICES=-1, no GPU visible); bail out early in the empty case to avoid querying torch.device(0) incorrectly * fix(rocm): gate StubSubpackageFinder on win32 ROCm, add gcnArchName fallbacks - worker.py _StubSubpackageFinder: the meta_path append was running on every platform on every call to run_training_process; moved it inside the if _is_win32_rocm: block since stubs are only seeded there and the finder is a pure accumulation on Linux/Windows CUDA - worker.py OOM guard: AMD SDK / Radeon wheels may not populate gcnArchName, causing Strix Halo to be misclassified as discrete and get the 0.90 cap (12.8 GB OS headroom) instead of 0.80 (25.6 GB); now tries gcn_arch_name / arch_name / gfx_arch_name variants first, then falls back to device-name matching (890M -> Strix Halo, 880M -> Strix Point) with a debug log when the fallback fires * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm): pin torchvision/torchaudio in setup.ps1, remove -Unique from arch array - setup.ps1 ROCm torch install: torchvision and torchaudio were passed bare alongside pinned torch>=2.11.0,<2.12.0 for gfx1151/gfx1200 arches. AMD publishes packages independently so a future torchvision 0.27 (for torch 2.12) on the same arch index would cause pip ResolutionImpossible or an ABI-incompatible install. Added torchvisionFloorMap and torchaudioFloorMap mirroring install_python_stack.py's strix override (torchvision>=0.26.0,<0.27.0, torchaudio>=2.11.0,<2.12.0) and derived ROCmVisionSpec/ROCmAudioSpec used in all three Fast-Install call sites. - setup.ps1 amd-smi arch detection: Select-Object -Unique was collapsing same-arch multi-GPU arrays (e.g. two gfx1151 APUs -> 1-element array) causing HIP_VISIBLE_DEVICES=1 to trigger a false out-of-range warning and fall back to GPU 0 even though the correct GPU would have been at index 1. Removed -Unique; added comment noting the positional-index assumption and its non-contiguous-GPU limitation. * fix(rocm): add 8060s/8050s to OOM guard device-name fallback, extract classifier helper Path 3 of the OOM guard device-name fallback only checked for 890m/880m (gfx1150 Strix Point SKU names). Strix Halo (gfx1151) ships as Radeon 8060S (Ryzen AI MAX+ 395) and Radeon 8050S (cut-down SKU) -- neither matches, so the fallback returned is_unified=False and applied the 0.90 fraction instead of 0.80, leaving ~12.8 GiB OS headroom on a 128 GiB pool instead of ~25.6 GiB. Fix: add 8060s and 8050s to the name-match set. Also correct the comment that mislabelled 890M as a Strix Halo name (it is Strix Point). Refactor: extract the three-path classifier into _rocm_classify_unified_memory() so it can be unit-tested directly. Add 31 test cases in test_rocm_oom_guard.py covering all three paths and the regression case (Radeon 8060S Graphics). Reported-by: h34v3nzc0dex * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(rocm): pass explicit dtype on bf16-unsupported hardware (RDNA2) dtype=None lets unsloth auto-detect the model dtype. On RDNA2 (gfx103x, e.g. RX 6600) is_bfloat16_supported() incorrectly returns True, so unsloth picks bf16 and the first bf16 kernel dispatch triggers: LLVM ERROR: Cannot select: intrinsic %llvm.amdgcn.fdot2.bf16.bf16 Replace every dtype=None in load_model() with _auto_dtype which resolves to None when bf16 is supported (all modern NVIDIA + RDNA3+) and torch.float16 otherwise. This gives RDNA2 users a working float16 training path without touching NVIDIA behaviour at all. Fixes: https://github.com/unslothai/unsloth/issues/5337 * fix: reduce log noise for expected non-issues on Windows ROCm Three log lines fired at warning/error level for conditions that are completely expected on a Windows HIP SDK-only setup: amd.py - amd-smi WinError 2 (FileNotFoundError): downgrade warning -> debug. amd-smi ships with Adrenalin, not the HIP SDK; absence is normal. - 'disabling' message: downgrade warning -> info with clearer text 'not available (not installed; expected on HIP SDK-only systems); GPU VRAM polling disabled' hardware.py - torch.distributed.Store missing: downgrade warning -> debug. The distributed stub added in this PR intentionally omits Store; the attention-impl fallback to eager is expected and non-actionable. worker.py - causal-conv1d: add early Windows exit (info) in both _ensure_causal_conv1d_fast_path and _causal_conv1d_install hook; no cp313/win_amd64 wheel exists, so the install always fails. - FLA: add early Windows exit (info) in _ensure_flash_linear_attention_unconditional; triton dependency has no cp313/win_amd64 wheel. - Defense-in-depth: _install_package_wheel_first non-HIP PyPI failure logs info+debug on Windows instead of error; FLA failure logs info+debug on Windows instead of warning. * [AMD] FIx installation of bitsandbytes when it's from .dev and skip rebuilding llama.cpp if we build it manually. * fix: use force_pip for Windows ROCm bitsandbytes prebuilt wheel install uv rejects the bnb continuous-release wheel due to filename/metadata version mismatch (1.33.7.preview vs 0.50.0.dev0). Switch to force_pip=True (pip bypass) instead of the UV_SKIP_WHEEL_FILENAME_CHECK env var workaround -- cleaner and consistent with how the Linux path handles it. BNB_ROCM_VERSION is still set post-install to the detected DLL suffix so the worker subprocess loads the correct libbitsandbytes_rocm{VER}.dll even when torch.version.hip reports a newer HIP version than the wheel ships. * fix: three small correctness fixes found in PR review - _install_bnb_windows_rocm: use UV_SKIP_WHEEL_FILENAME_CHECK=1 with try/finally instead of force_pip=True so the env var is always restored and the failing CI test passes - _determine_attention_impl_for_gpu_estimate: gate torch._C distributed stubs on IS_ROCM so Windows CUDA users keep the real extension - install.ps1 amd-smi fallback: collect all gfx tokens and index by HIP_VISIBLE_DEVICES, matching the hipinfo path on multi-GPU hosts * fix: stub torchao in export subprocess on Windows ROCm On Windows, the ROCm build of PyTorch ships without the distributed C extension (torch._C._distributed_c10d). torchao, which is pulled in transitively by transformers.quantizers at import time, walks into torch.distributed._functional_collectives -> distributed_c10d and crashes with: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package This only affected the export subprocess because the training subprocess already applied an identical torchao stub (introduced separately to fix the same root cause). The export subprocess had no such guard and died during 'Importing Unsloth...' before any model loading could happen. Fix: apply the same _StubSubpackageFinder / torchao stub pattern to the export subprocess entry point, gated on Windows ROCm detection, before any import of transformers or unsloth_zoo. Root cause tracked in ROCm/TheRock#3284 (libuv / torch.distributed missing on Windows ROCm builds). Ref: https://github.com/ROCm/TheRock/issues/3284 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install.sh, setup.sh: add GPU arch step logging to match PS1 scripts Both shell scripts were missing the step "gpu" terminal log block that install.ps1 and setup.ps1 emit. This adds equivalent output: GPU label with gfx arch (e.g. "AMD ROCm (gfx1151)"), ROCm root path, hipconfig version, and marketing name substep. Includes the same gfx arch detection chain (rocminfo → amd-smi list → amd-smi static --asic), UNSLOTH_ROCM_GFX_ARCH env override, and name-based arch inference table (Strix Halo/Point, RDNA 3/4) as the PS1 versions. install.sh also replaces bare echo blocks for the AMD ROCm and CPU-only cases with formatted substep output. * Fix BNB_ROCM_VERSION gate, ROCm GPU mask preference, APU unified memory and Release build for PR #5301 - main.py: gate BNB_ROCM_VERSION on the rocm bnb DLL or HIP_PATH/ROCM_PATH instead of importing torch on every Windows host - hardware.py: prefer HIP/ROCR visible-device masks only on ROCm hosts so a stale mask cannot override CUDA_VISIBLE_DEVICES on NVIDIA - llama_cpp.py: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1 only for unified-memory APUs (gfx1150/gfx1151) - setup.sh: pass -DCMAKE_BUILD_TYPE=Release for the HIP source build - add test_amd_apu_unified_memory.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: guard recompile_limit + fix AMD VRAM monitor fallback trainer.py: torch._dynamo.config.recompile_limit does not exist in some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2 wheels). Guard the assignment so training doesn't crash on RDNA2/RDNA3. hardware.py: when amd-smi/nvidia-smi is unavailable or returns no usable data (HIP SDK-only Windows, Docker, unexpected JSON format), the existing fallback used torch.cuda.memory_allocated() which is process-specific and reads near-zero even with a fully loaded model. Switch to torch.cuda.mem_get_info() via _torch_get_per_device_info() which reports system-wide VRAM occupancy so the GPU monitor shows real usage on all AMD systems without requiring amd-smi. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: Windows VRAM monitor via Performance Counter API When amd-smi/nvidia-smi is unavailable on Windows, query dedicated GPU VRAM via Windows Performance Counters (same source as Task Manager). This gives system-wide cross-process usage, fixing the near-zero reading caused by torch.cuda.mem_get_info only seeing the Studio server process. Linux fallback path unchanged (mem_get_info is system-wide on ROCm). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: rename to _rocm_windows_perf_counter_vram_gb, scope to IS_ROCM Function is AMD ROCm specific — amd-smi absent on Windows when only the HIP SDK is installed. Scoped to IS_ROCM so NVIDIA Windows path is untouched (nvidia-smi handles that case). * fix: AMD VRAM monitor — Linux DRM sysfs + Windows perf counter Linux: read /sys/class/drm/card*/device/mem_info_vram_used|total for system-wide GPU memory across all processes. No tools required, always present on Linux AMD systems. Windows: Windows Performance Counter API (already added). Both paths are gated on IS_ROCM and only fire when amd-smi is absent. torch mem_get_info remains as last resort (process-local). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: AMD GPU monitor — utilization, temperature, and power for Windows and Linux fallback paths - Windows: GPU utilization via \GPU Engine(*engtype_3D*)\Utilization Percentage perf counter - Windows: temperature and power via ADL (atiadlxx.dll, ships with Adrenalin) - Linux: GPU utilization via DRM sysfs gpu_busy_percent - Linux: temperature via hwmon temp1_input (millidegrees C) - Linux: power via hwmon power1_average / power1_input (microwatts) All paths are no-op fallbacks (None) when the source is unavailable. Mirrors what nvidia-smi provides on the CUDA path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: remove ADL ctypes — does not support AMD iGPU (Strix Halo) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Erland366 <erland.pg366@gmail.com> Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
5584 lines
244 KiB
Python
5584 lines
244 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||
|
||
"""
|
||
llama-server inference backend for GGUF models.
|
||
|
||
Manages a llama-server subprocess and proxies chat completions
|
||
through its OpenAI-compatible /v1/chat/completions endpoint.
|
||
"""
|
||
|
||
import atexit
|
||
import contextlib
|
||
import json
|
||
import os
|
||
import re
|
||
import struct
|
||
import structlog
|
||
from loggers import get_logger
|
||
import shutil
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Generator, Iterable, List, Optional
|
||
from urllib.parse import urlparse
|
||
|
||
import httpx
|
||
|
||
from core.inference.llama_server_args import (
|
||
parse_cache_override,
|
||
parse_ctx_override,
|
||
resolve_cache_type_kv,
|
||
resolve_requested_ctx,
|
||
)
|
||
from core.tool_healing import (
|
||
_TC_END_TAG_RE,
|
||
_TC_FUNC_CLOSE_RE,
|
||
_TC_FUNC_START_RE,
|
||
_TC_JSON_START_RE,
|
||
_TC_PARAM_CLOSE_RE,
|
||
_TC_PARAM_START_RE,
|
||
_TOOL_ALL_PATS,
|
||
_TOOL_CLOSED_PATS,
|
||
parse_tool_calls_from_text,
|
||
strip_tool_call_markup,
|
||
)
|
||
from utils.native_path_leases import child_env_without_native_path_secret
|
||
from utils.subprocess_compat import (
|
||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||
)
|
||
from core.inference.tool_call_parser import (
|
||
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
|
||
)
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
# ── Pre-compiled patterns for plan-without-action re-prompt ──
|
||
# Forward-looking intent signals that indicate the model is
|
||
# describing what it *will* do rather than giving a final answer.
|
||
_INTENT_SIGNAL = re.compile(
|
||
r"(?i)("
|
||
# Direct intent: "I'll ...", "I will ...", "Let me ...", "I am going to ..."
|
||
# Handles both straight and curly apostrophes.
|
||
# Excludes "I can", "I should", "I want to", "let's" which
|
||
# appear frequently in direct answers / explanations.
|
||
# Negative lookahead drops negated forms ("I will not", "I'll never")
|
||
# so a refusal doesn't trigger a re-prompt.
|
||
r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)"
|
||
r"|"
|
||
# Step/plan framing: "First ...", "Step 1:", "Here's my plan"
|
||
r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))"
|
||
r"|"
|
||
# "Now I" / "Next I" patterns
|
||
r"\b(?:now i|next i)\b"
|
||
r")"
|
||
)
|
||
_MAX_REPROMPTS = 3
|
||
|
||
# Without max_tokens, llama-server defaults to n_predict = n_ctx (up to
|
||
# 262144 for Qwen3.5), producing many-minute zombie decodes when cancel
|
||
# fails. t_max_predict_ms is a wall-clock backstop applied unconditionally,
|
||
# but the llama.cpp README notes it ONLY fires after a newline has been
|
||
# generated -- a model stuck in a long unbroken non-newline sequence is
|
||
# unbounded by it. So we still want a token cap as the front-line limiter.
|
||
#
|
||
# The cap is the model's effective context length when we know it,
|
||
# falling back to a generous floor when metadata is unavailable. 4096 was
|
||
# too low: Qwen3 / gpt-oss reasoning traces routinely exceed it, and any
|
||
# OpenAI-API caller that omits max_tokens (langchain, llama-index, raw
|
||
# curl) sees responses silently truncated mid-sentence.
|
||
_DEFAULT_MAX_TOKENS_FLOOR = 32768
|
||
_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min
|
||
_REPROMPT_MAX_CHARS = 2000
|
||
|
||
# ── Pre-compiled patterns for GGUF shard detection ───────────
|
||
_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
|
||
_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$")
|
||
|
||
|
||
# ── Sliding-window-pattern resolver ───────────────────────────
|
||
# Resolves the per-layer SWA mask when a GGUF reports a sliding window
|
||
# but no `sliding_window_pattern` field. Tier order in
|
||
# `_resolve_swa_pattern`: GGUF metadata, on-disk cache, bootstrap dict
|
||
# below, transformers introspection, HF Hub config.json, legacy 1/4
|
||
# fallback. Period N means layer i is SWA iff `(i + 1) % N != 0`,
|
||
# matching transformers. Skipped on purpose: phi3 (no key/val length
|
||
# in GGUF, window >= ctx anyway), qwen2 family (converter strips
|
||
# sliding_window when use_sliding_window=False), mistral v0.1/v0.2
|
||
# (all-SWA can't be expressed as a period).
|
||
_BOOTSTRAP_SWA_DEFAULTS: dict[str, int] = {
|
||
"gemma2": 2, # Gemma2Config.sliding_window_pattern
|
||
"gemma3": 6, # Gemma3TextConfig.sliding_window_pattern
|
||
"gemma3n": 5, # text_config.layer_types: SWA*4 + FULL
|
||
"gpt_oss": 2, # text_config.layer_types: alternating
|
||
"cohere2": 4, # Cohere2Config.sliding_window_pattern
|
||
}
|
||
|
||
# Process-wide cache backed by JSON on disk. Values are int period or
|
||
# list[bool] mask. Lazy-loaded.
|
||
_SWA_CACHE: Optional[dict] = None
|
||
_SWA_CACHE_LOCK = threading.Lock()
|
||
|
||
|
||
def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool:
|
||
"""Quick DNS check. Runs on a daemon thread so concurrent sockets
|
||
in the same process are not affected by socket.setdefaulttimeout."""
|
||
result: list[Optional[bool]] = [None]
|
||
|
||
def _probe() -> None:
|
||
try:
|
||
socket.gethostbyname(host)
|
||
result[0] = False
|
||
except Exception:
|
||
result[0] = True
|
||
|
||
t = threading.Thread(target = _probe, daemon = True)
|
||
t.start()
|
||
t.join(timeout)
|
||
# Thread still running -> resolver wedged -> treat as dead.
|
||
return True if result[0] is None else result[0]
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def _hf_offline_if_dns_dead():
|
||
"""Set HF_HUB_OFFLINE for the body of this block only when DNS to
|
||
huggingface.co fails. Restores the env on exit so a transient
|
||
resolver hiccup at the start of one load can't quarantine the whole
|
||
process. Respects an explicit user setting (no-op if already set)."""
|
||
if "HF_HUB_OFFLINE" in os.environ:
|
||
yield False
|
||
return
|
||
if not _probe_dns_dead():
|
||
yield False
|
||
return
|
||
|
||
transformers_was_set = "TRANSFORMERS_OFFLINE" in os.environ
|
||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||
if not transformers_was_set:
|
||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||
logger.warning("huggingface.co unreachable; using local HF cache for this load.")
|
||
try:
|
||
yield True
|
||
finally:
|
||
os.environ.pop("HF_HUB_OFFLINE", None)
|
||
if not transformers_was_set:
|
||
os.environ.pop("TRANSFORMERS_OFFLINE", None)
|
||
|
||
|
||
def _swa_cache_path() -> Path:
|
||
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
|
||
base = Path(home) if home else Path.home() / ".unsloth" / "studio"
|
||
return base / "swa_cache.json"
|
||
|
||
|
||
def _load_swa_cache() -> dict:
|
||
global _SWA_CACHE
|
||
with _SWA_CACHE_LOCK:
|
||
if _SWA_CACHE is not None:
|
||
return _SWA_CACHE
|
||
try:
|
||
with open(_swa_cache_path()) as f:
|
||
_SWA_CACHE = json.load(f)
|
||
if not isinstance(_SWA_CACHE, dict):
|
||
_SWA_CACHE = {}
|
||
except (FileNotFoundError, json.JSONDecodeError, OSError):
|
||
_SWA_CACHE = {}
|
||
return _SWA_CACHE
|
||
|
||
|
||
def _save_swa_cache(cache: dict) -> None:
|
||
try:
|
||
path = _swa_cache_path()
|
||
path.parent.mkdir(parents = True, exist_ok = True)
|
||
tmp = path.with_suffix(".json.tmp")
|
||
with open(tmp, "w") as f:
|
||
json.dump(cache, f, indent = 2, sort_keys = True)
|
||
tmp.replace(path)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def _period_from_layer_types(layer_types: list) -> Optional[int]:
|
||
"""Smallest period N where `(i+1) % N != 0` matches the SWA mask,
|
||
or None if no fixed period fits."""
|
||
if not layer_types:
|
||
return None
|
||
is_swa = ["full" not in str(t).lower() for t in layer_types]
|
||
n = len(is_swa)
|
||
for N in range(1, n + 1):
|
||
if all(((i + 1) % N != 0) == is_swa[i] for i in range(n)):
|
||
return N
|
||
return None
|
||
|
||
|
||
def _fetch_swa_entry_from_hf(repo_id: str) -> Optional[object]:
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
|
||
cfg_path = hf_hub_download(repo_id, "config.json", repo_type = "model")
|
||
with open(cfg_path) as f:
|
||
cfg = json.load(f)
|
||
except Exception:
|
||
return None
|
||
|
||
src = cfg.get("text_config") if isinstance(cfg.get("text_config"), dict) else cfg
|
||
period = src.get("sliding_window_pattern")
|
||
if isinstance(period, int) and period > 0:
|
||
return period
|
||
lt = src.get("layer_types")
|
||
if isinstance(lt, list) and lt:
|
||
return _period_from_layer_types(lt) or [
|
||
"full" not in str(t).lower() for t in lt
|
||
]
|
||
return None
|
||
|
||
|
||
def _arch_aliases(arch: str) -> tuple:
|
||
# GGUF emits `falcon-h1`; HF model_type is `falcon_h1`. Normalise both ways.
|
||
seen = []
|
||
for a in (arch, arch.replace("-", "_"), arch.replace("_", "-")):
|
||
if a and a not in seen:
|
||
seen.append(a)
|
||
return tuple(seen)
|
||
|
||
|
||
def _swa_entry_from_config_obj(cfg) -> Optional[object]:
|
||
src = getattr(cfg, "text_config", None) or cfg
|
||
period = getattr(src, "sliding_window_pattern", None)
|
||
if isinstance(period, int) and period > 0:
|
||
return period
|
||
lt = getattr(src, "layer_types", None)
|
||
if isinstance(lt, list) and lt:
|
||
return _period_from_layer_types(lt) or [
|
||
"full" not in str(t).lower() for t in lt
|
||
]
|
||
return None
|
||
|
||
|
||
_SWA_PATTERN_SOURCE_RE = re.compile(
|
||
r"sliding_window_pattern\s*(?::\s*[\w\[\], ]*)?\s*=\s*(\d+)"
|
||
)
|
||
|
||
|
||
def _resolve_swa_entry_from_transformers(arch: str) -> Optional[object]:
|
||
"""Default-instantiate the matching Config; on failure, regex-parse
|
||
its source for `sliding_window_pattern = N`."""
|
||
try:
|
||
from transformers.models.auto.configuration_auto import (
|
||
CONFIG_MAPPING,
|
||
CONFIG_MAPPING_NAMES,
|
||
)
|
||
except Exception:
|
||
return None
|
||
|
||
cfg_class = None
|
||
for alias in _arch_aliases(arch):
|
||
if alias in CONFIG_MAPPING_NAMES:
|
||
try:
|
||
cfg_class = CONFIG_MAPPING[alias]
|
||
break
|
||
except Exception:
|
||
cfg_class = None
|
||
if cfg_class is None:
|
||
return None
|
||
|
||
try:
|
||
if (entry := _swa_entry_from_config_obj(cfg_class())) is not None:
|
||
return entry
|
||
except Exception:
|
||
pass
|
||
|
||
import inspect
|
||
|
||
candidates = [cfg_class]
|
||
text_cfg_class = getattr(cfg_class, "sub_configs", {}).get("text_config")
|
||
if text_cfg_class is not None:
|
||
candidates.append(text_cfg_class)
|
||
for cls in candidates:
|
||
try:
|
||
src = inspect.getsource(cls)
|
||
except (OSError, TypeError):
|
||
continue
|
||
if m := _SWA_PATTERN_SOURCE_RE.search(src):
|
||
period = int(m.group(1))
|
||
if period > 0:
|
||
return period
|
||
return None
|
||
|
||
|
||
def _resolve_swa_pattern(
|
||
arch: Optional[str],
|
||
n_layers: Optional[int],
|
||
source_repo_candidates: tuple = (),
|
||
*,
|
||
allow_network: Optional[bool] = None,
|
||
) -> Optional[list]:
|
||
if not arch or not n_layers:
|
||
return None
|
||
if allow_network is None:
|
||
allow_network = os.environ.get("UNSLOTH_STUDIO_OFFLINE", "0") not in (
|
||
"1",
|
||
"true",
|
||
"True",
|
||
"yes",
|
||
)
|
||
|
||
cache = _load_swa_cache()
|
||
|
||
def _entry_to_mask(entry):
|
||
if isinstance(entry, int) and entry > 0:
|
||
return [(i + 1) % entry != 0 for i in range(n_layers)]
|
||
if isinstance(entry, list) and entry:
|
||
return [bool(entry[i % len(entry)]) for i in range(n_layers)]
|
||
return None
|
||
|
||
def _persist(entry):
|
||
with _SWA_CACHE_LOCK:
|
||
cache[arch] = entry
|
||
_save_swa_cache(cache)
|
||
|
||
if (entry := cache.get(arch)) is not None:
|
||
if (mask := _entry_to_mask(entry)) is not None:
|
||
return mask
|
||
|
||
if (entry := _BOOTSTRAP_SWA_DEFAULTS.get(arch)) is not None:
|
||
return _entry_to_mask(entry)
|
||
|
||
entry = _resolve_swa_entry_from_transformers(arch)
|
||
if entry is not None:
|
||
_persist(entry)
|
||
return _entry_to_mask(entry)
|
||
|
||
# Tier 3: live HF fetch (with persistent caching of the result)
|
||
if allow_network:
|
||
for repo_id in source_repo_candidates:
|
||
if not repo_id:
|
||
continue
|
||
entry = _fetch_swa_entry_from_hf(repo_id)
|
||
if entry is not None:
|
||
_persist(entry)
|
||
return _entry_to_mask(entry)
|
||
|
||
return None
|
||
|
||
|
||
def _hf_repo_from_url(url: Optional[str]) -> Optional[str]:
|
||
"""Strip `https://huggingface.co/owner/name(/...)` to `owner/name`."""
|
||
if not url or "huggingface.co/" not in url:
|
||
return None
|
||
tail = url.split("huggingface.co/", 1)[1].rstrip("/")
|
||
parts = tail.split("/")
|
||
if len(parts) < 2:
|
||
return None
|
||
return f"{parts[0]}/{parts[1]}"
|
||
|
||
|
||
# Model size extraction — lazy import to avoid pulling in transformers
|
||
# at module level. See PR description for the full explanation.
|
||
def _extract_model_size_b(model_id: str):
|
||
from utils.models import extract_model_size_b
|
||
|
||
return extract_model_size_b(model_id)
|
||
|
||
|
||
_TOOL_TEMPLATE_MARKERS = (
|
||
"{%- if tools %}",
|
||
"{%- if tools -%}",
|
||
"{% if tools %}",
|
||
"{% if tools -%}",
|
||
'"role" == "tool"',
|
||
"'role' == 'tool'",
|
||
'message.role == "tool"',
|
||
"message.role == 'tool'",
|
||
)
|
||
|
||
|
||
def detect_reasoning_flags(
|
||
chat_template: Optional[str],
|
||
model_identifier: Optional[str] = None,
|
||
*,
|
||
log_source: Optional[str] = None,
|
||
) -> dict:
|
||
"""Classify a chat template's reasoning and tool-calling capabilities.
|
||
|
||
Returns a dict with the same five keys populated by the GGUF sniffer:
|
||
``supports_reasoning``, ``reasoning_style``
|
||
(``"enable_thinking"`` | ``"reasoning_effort"``),
|
||
``reasoning_always_on``, ``supports_preserve_thinking``, and
|
||
``supports_tools``. Used by both the llama-server backend at load
|
||
time and the safetensors/transformers paths in ``routes/inference``
|
||
so the two agree on what the frontend will see.
|
||
"""
|
||
flags = {
|
||
"supports_reasoning": False,
|
||
"reasoning_style": "enable_thinking",
|
||
"reasoning_always_on": False,
|
||
"supports_preserve_thinking": False,
|
||
"supports_tools": False,
|
||
}
|
||
if not chat_template:
|
||
return flags
|
||
tpl = chat_template
|
||
prefix = f"{log_source}: " if log_source else ""
|
||
|
||
if "enable_thinking" in tpl:
|
||
flags["supports_reasoning"] = True
|
||
flags["reasoning_style"] = "enable_thinking"
|
||
logger.info(f"{prefix}model supports reasoning (enable_thinking)")
|
||
elif "reasoning_effort" in tpl:
|
||
# gpt-oss / Harmony templates use reasoning_effort
|
||
# ("low" | "medium" | "high") instead of a boolean.
|
||
flags["supports_reasoning"] = True
|
||
flags["reasoning_style"] = "reasoning_effort"
|
||
logger.info(f"{prefix}model supports reasoning (reasoning_effort)")
|
||
elif "thinking" in tpl:
|
||
# DeepSeek uses 'thinking' instead of 'enable_thinking'
|
||
normalized_id = (model_identifier or "").lower()
|
||
if "deepseek" in normalized_id:
|
||
flags["supports_reasoning"] = True
|
||
logger.info(f"{prefix}model supports reasoning (DeepSeek thinking)")
|
||
|
||
# Hardcoded <think> tags or reasoning_content in the template mean
|
||
# thinking is always on (no toggle to disable it).
|
||
if not flags["supports_reasoning"]:
|
||
if ("<think>" in tpl and "</think>" in tpl) or "reasoning_content" in tpl:
|
||
flags["supports_reasoning"] = True
|
||
flags["reasoning_always_on"] = True
|
||
logger.info(f"{prefix}model always reasons (<think> tags in template)")
|
||
|
||
# preserve_thinking is an independent kwarg on some Qwen templates
|
||
# that keeps historical <think> blocks in prior assistant turns.
|
||
if "preserve_thinking" in tpl:
|
||
flags["supports_preserve_thinking"] = True
|
||
logger.info(f"{prefix}model supports preserve_thinking")
|
||
|
||
if any(marker in tpl for marker in _TOOL_TEMPLATE_MARKERS):
|
||
flags["supports_tools"] = True
|
||
logger.info(f"{prefix}model supports tool calling")
|
||
|
||
return flags
|
||
|
||
|
||
def _is_mtp_model_name(
|
||
model_identifier: Optional[str],
|
||
gguf_path: Optional[str] = None,
|
||
) -> bool:
|
||
"""Name-based MTP detector. Fallback for the metadata signal."""
|
||
for cand in (model_identifier, Path(gguf_path).name if gguf_path else None):
|
||
if cand and "-mtp" in cand.lower():
|
||
return True
|
||
return False
|
||
|
||
|
||
def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool:
|
||
"""User passed --spec-type / --spec-default? llama-server takes a
|
||
single --spec-type (comma-separated to chain), so suppress
|
||
auto-emit when this is true."""
|
||
if not extra_args:
|
||
return False
|
||
for raw in extra_args:
|
||
tok = str(raw)
|
||
if not tok.startswith("--"):
|
||
continue
|
||
flag = tok.split("=", 1)[0]
|
||
if flag in ("--spec-type", "--spec-default"):
|
||
return True
|
||
return False
|
||
|
||
|
||
def _build_ngram_mod_flags(
|
||
caps: Optional[dict],
|
||
n_match: int = 24,
|
||
n_min: int = 48,
|
||
n_max: int = 64,
|
||
) -> list[str]:
|
||
"""Emit the right ngram-mod knob flags for the running llama-server.
|
||
|
||
Post-rename builds expose ``--spec-ngram-mod-n-{match,min,max}``;
|
||
pre-rename builds expose the legacy ``--spec-ngram-size-n`` /
|
||
``--draft-min`` / ``--draft-max``. ``caps`` comes from
|
||
``probe_server_capabilities``; ``ngram_mod_flavor`` tells us which
|
||
set is real (vs a removal-stub entry). Returns ``[]`` when neither
|
||
set is available so the caller can drop ngram-mod entirely.
|
||
"""
|
||
flavor = caps.get("ngram_mod_flavor") if caps else None
|
||
if flavor == "new":
|
||
return [
|
||
"--spec-ngram-mod-n-match",
|
||
str(n_match),
|
||
"--spec-ngram-mod-n-min",
|
||
str(n_min),
|
||
"--spec-ngram-mod-n-max",
|
||
str(n_max),
|
||
]
|
||
if flavor == "legacy":
|
||
# Legacy llama.cpp before the spec arg rename: same knobs lived
|
||
# under --spec-ngram-size-n (lookup length) and the generic
|
||
# --draft-min / --draft-max (ngram size N range).
|
||
return [
|
||
"--spec-ngram-size-n",
|
||
str(n_match),
|
||
"--draft-min",
|
||
str(n_min),
|
||
"--draft-max",
|
||
str(n_max),
|
||
]
|
||
return []
|
||
|
||
|
||
# Canonical Speculative Decoding modes exposed by the Studio chat UI.
|
||
# The dropdown renders five options (auto, mtp, ngram, mtp+ngram, off);
|
||
# the load API also accepts legacy values that the original Switch and
|
||
# external callers emit (default, draft-mtp, ngram-mod, ngram-simple).
|
||
_CANONICAL_SPEC_MODES = {"auto", "mtp", "ngram", "mtp+ngram", "off", "ngram-simple"}
|
||
_LEGACY_SPEC_MODE_MAP = {
|
||
"default": "auto",
|
||
"draft-mtp": "mtp",
|
||
"ngram-mod": "ngram",
|
||
}
|
||
|
||
|
||
def _canonicalize_spec_mode(value):
|
||
"""Map any accepted ``speculative_type`` input onto a canonical mode.
|
||
|
||
Returns one of ``auto``, ``mtp``, ``ngram``, ``mtp+ngram``, ``off``,
|
||
``ngram-simple``, or ``None`` (callers treat ``None`` as ``auto``).
|
||
Unknown strings collapse to ``auto`` so a stale UI value or typo
|
||
falls back to the safe platform-aware path.
|
||
"""
|
||
if value is None:
|
||
return None
|
||
if not isinstance(value, str):
|
||
return None
|
||
stripped = value.strip().lower()
|
||
if not stripped:
|
||
return None
|
||
if stripped in _CANONICAL_SPEC_MODES:
|
||
return stripped
|
||
if stripped in _LEGACY_SPEC_MODE_MAP:
|
||
return _LEGACY_SPEC_MODE_MAP[stripped]
|
||
# llama.cpp comma-chains are emitted by old persisted state e.g.
|
||
# "ngram-mod,draft-mtp"; collapse the most common one explicitly.
|
||
pieces = [p.strip() for p in stripped.split(",") if p.strip()]
|
||
has_mtp = any(p in ("mtp", "draft-mtp") for p in pieces)
|
||
has_ngram = any(p in ("ngram", "ngram-mod") for p in pieces)
|
||
if has_mtp and has_ngram:
|
||
return "mtp+ngram"
|
||
if has_mtp:
|
||
return "mtp"
|
||
if has_ngram:
|
||
return "ngram"
|
||
return "auto"
|
||
|
||
|
||
def _backfill_usage_from_timings(usage, timings):
|
||
"""Synthesize ``usage`` from llama-server's ``timings`` when the
|
||
OpenAI-style usage block is missing or reports zero tokens.
|
||
|
||
The Studio chat UI computes generation t/s from
|
||
``meta.usage.completion_tokens / totalStreamTime``. llama-server
|
||
always populates ``timings.predicted_n`` (true decoded count) and
|
||
``timings.prompt_n``, but the ``usage`` field on the final SSE chunk
|
||
can be absent or zero on some server builds / streaming
|
||
configurations, which makes the UI fall back to wall-clock t/s and
|
||
dilute speculative-decoding speedups.
|
||
"""
|
||
if not timings:
|
||
return usage
|
||
if usage and usage.get("completion_tokens"):
|
||
return usage
|
||
predicted_n = timings.get("predicted_n")
|
||
prompt_n = timings.get("prompt_n")
|
||
if predicted_n is None and prompt_n is None:
|
||
return usage
|
||
out = dict(usage or {})
|
||
if not out.get("completion_tokens") and predicted_n is not None:
|
||
out["completion_tokens"] = predicted_n
|
||
if not out.get("prompt_tokens") and prompt_n is not None:
|
||
out["prompt_tokens"] = prompt_n
|
||
out["total_tokens"] = int(out.get("prompt_tokens") or 0) + int(
|
||
out.get("completion_tokens") or 0
|
||
)
|
||
return out
|
||
|
||
|
||
class LlamaCppBackend:
|
||
"""
|
||
Manages a llama-server subprocess for GGUF model inference.
|
||
|
||
Lifecycle:
|
||
1. load_model() — starts llama-server with the GGUF file
|
||
2. generate_chat_completion() — proxies to /v1/chat/completions, streams back
|
||
3. unload_model() — terminates llama-server subprocess
|
||
"""
|
||
|
||
def __init__(self):
|
||
self._process: Optional[subprocess.Popen] = None
|
||
self._port: Optional[int] = None
|
||
self._model_identifier: Optional[str] = None
|
||
self._gguf_path: Optional[str] = None
|
||
self._hf_repo: Optional[str] = None
|
||
self._hf_variant: Optional[str] = None
|
||
self._is_vision: bool = False
|
||
self._healthy = False
|
||
# Set by _classify_gpu_offload after _wait_for_health.
|
||
self._gpu_offload_active: Optional[bool] = None
|
||
self._context_length: Optional[int] = None
|
||
self._effective_context_length: Optional[int] = None
|
||
self._max_context_length: Optional[int] = None
|
||
self._chat_template: Optional[str] = None
|
||
self._chat_template_override: Optional[str] = None
|
||
self._supports_reasoning: bool = False
|
||
self._reasoning_always_on: bool = False
|
||
self._reasoning_style: str = "enable_thinking"
|
||
self._supports_preserve_thinking: bool = False
|
||
self._supports_tools: bool = False
|
||
self._cache_type_kv: Optional[str] = None
|
||
self._reasoning_default: bool = True
|
||
self._speculative_type: Optional[str] = None
|
||
# Canonical UI-facing mode the user requested: one of
|
||
# ``auto``/``mtp``/``ngram``/``mtp+ngram``/``off``/``ngram-simple``.
|
||
# Round-tripped through the status API so the dropdown reflects
|
||
# the picked mode rather than the resolved internal flag set
|
||
# (auto on a 27B MTP GGUF resolves to draft-mtp but the dropdown
|
||
# should still read "Auto").
|
||
self._requested_spec_mode: Optional[str] = None
|
||
# User-supplied --spec-draft-n-max override (None = platform default).
|
||
self._spec_draft_n_max: Optional[int] = None
|
||
# KV-cache estimation fields (populated by _read_gguf_metadata)
|
||
self._n_layers: Optional[int] = None
|
||
self._n_kv_heads: Optional[int] = None
|
||
self._n_kv_heads_by_layer: Optional[list[int]] = None
|
||
self._n_heads: Optional[int] = None
|
||
self._embedding_length: Optional[int] = None
|
||
# Architecture-aware KV fields for 5-path estimation
|
||
self._kv_key_length: Optional[int] = None
|
||
self._kv_value_length: Optional[int] = None
|
||
self._sliding_window: Optional[int] = None
|
||
self._sliding_window_pattern: Optional[list[bool]] = None
|
||
self._full_attention_interval: Optional[int] = None
|
||
self._kv_lora_rank: Optional[int] = None
|
||
self._key_length_mla: Optional[int] = None
|
||
self._kv_key_length_swa: Optional[int] = None
|
||
self._kv_value_length_swa: Optional[int] = None
|
||
self._ssm_inner_size: Optional[int] = None
|
||
self._ssm_state_size: Optional[int] = None
|
||
# Last N layers reuse KV from earlier layers and don't allocate
|
||
# their own cache (Gemma 3n / Gemma 4: <arch>.attention.shared_kv_layers).
|
||
self._shared_kv_layers: Optional[int] = None
|
||
# MTP head count (llama.cpp #22673); >0 enables --spec-type draft-mtp.
|
||
self._nextn_predict_layers: Optional[int] = None
|
||
self._lock = threading.Lock()
|
||
# Wraps load_model() end-to-end so concurrent loads serialise
|
||
# and never coexist as two llama-server processes (#5401).
|
||
self._serial_load_lock = threading.Lock()
|
||
# Last extra_args / requested n_ctx, preserved across unload so
|
||
# the chat UI's /unload+/load Apply path can inherit them (#5401).
|
||
# ``_extra_args_source`` records the (model_identifier, hf_variant)
|
||
# the stored args came from so the route can refuse cross-model
|
||
# inheritance.
|
||
self._extra_args: Optional[List[str]] = None
|
||
self._extra_args_source: Optional[tuple[str, Optional[str]]] = None
|
||
self._requested_n_ctx: int = 0
|
||
self._stdout_lines: list[str] = []
|
||
self._stdout_thread: Optional[threading.Thread] = None
|
||
# llama-server tee log (see _drain_stdout / _kill_process).
|
||
self._llama_log_fh = None
|
||
self._llama_log_path: Optional[Path] = None
|
||
self._cancel_event = threading.Event()
|
||
self._api_key: Optional[str] = None
|
||
# True once a probe has completed; cleared on transient failure.
|
||
self._is_audio: bool = False
|
||
self._audio_type: Optional[str] = None
|
||
self._audio_probed: bool = False
|
||
# Monotonic timestamp set in _kill_process; read by load_model
|
||
# to decide whether to wait for the VRAM reclaim to finish.
|
||
self._last_kill_monotonic: float = 0.0
|
||
|
||
self._kill_orphaned_servers()
|
||
atexit.register(self._cleanup)
|
||
|
||
# ── Properties ────────────────────────────────────────────────
|
||
|
||
@property
|
||
def is_loaded(self) -> bool:
|
||
return self._process is not None and self._healthy
|
||
|
||
@property
|
||
def is_active(self) -> bool:
|
||
"""True if a llama-server process exists (loading or loaded)."""
|
||
return self._process is not None
|
||
|
||
@property
|
||
def base_url(self) -> str:
|
||
return f"http://127.0.0.1:{self._port}"
|
||
|
||
@property
|
||
def model_identifier(self) -> Optional[str]:
|
||
return self._model_identifier
|
||
|
||
@property
|
||
def is_vision(self) -> bool:
|
||
return self._is_vision
|
||
|
||
@property
|
||
def hf_variant(self) -> Optional[str]:
|
||
return self._hf_variant
|
||
|
||
@property
|
||
def extra_args(self) -> Optional[List[str]]:
|
||
"""Extra llama-server flags from the last load. Copy; None = never
|
||
set, [] = explicitly cleared. Used by the route for inheritance."""
|
||
return list(self._extra_args) if self._extra_args is not None else None
|
||
|
||
@property
|
||
def requested_n_ctx(self) -> int:
|
||
"""n_ctx the last load was invoked with (not the effective cap).
|
||
0 means Auto. Used by the route to detect Auto-vs-explicit flips."""
|
||
return self._requested_n_ctx
|
||
|
||
@property
|
||
def extra_args_source(self) -> Optional[tuple[str, Optional[str]]]:
|
||
"""(model_identifier, hf_variant) the stored extra_args came from.
|
||
``None`` if no extras have ever been recorded. Used by the route
|
||
to refuse cross-model inheritance (#5401)."""
|
||
return self._extra_args_source
|
||
|
||
@property
|
||
def context_length(self) -> Optional[int]:
|
||
"""Return the effective context length the server is running at."""
|
||
return self._effective_context_length or self._context_length
|
||
|
||
@property
|
||
def max_context_length(self) -> Optional[int]:
|
||
"""Return the largest context that fits on this hardware at load time.
|
||
|
||
This is the "safe zone" threshold the UI renders warnings
|
||
against. For a model whose weights fit on some GPU subset, it
|
||
is the binary-search cap from ``_fit_context_to_vram`` for that
|
||
subset. For a model whose weights exceed 90% of every GPU
|
||
subset, it is the 4096 fallback -- the spec's default when the
|
||
model will not fit. The UI slider ceiling is
|
||
``native_context_length``; dragging above ``max_context_length``
|
||
triggers the "might be slower" warning.
|
||
"""
|
||
return self._max_context_length or self._context_length
|
||
|
||
@property
|
||
def native_context_length(self) -> Optional[int]:
|
||
"""Return the model's native context length from GGUF metadata."""
|
||
return self._context_length
|
||
|
||
def load_progress(self) -> Optional[dict]:
|
||
"""Return live model-load progress, or None if not loading.
|
||
|
||
While llama-server is warming up, its process is typically in
|
||
kernel state D (disk sleep) mmap'ing the weight shards into
|
||
page cache before pushing layers to VRAM. During that window
|
||
``/api/inference/status`` only reports ``loading``, which gives
|
||
the UI nothing to display besides a spinner that looks stuck
|
||
for minutes on large MoE models.
|
||
|
||
This method samples ``/proc/<pid>/status VmRSS`` against the
|
||
sum of the GGUF shard sizes so the UI can render a real bar
|
||
and compute rate / ETA. Returns ``None`` when no load is in
|
||
flight (no process, or process already healthy).
|
||
|
||
Shape::
|
||
|
||
{
|
||
"phase": "mmap" | "ready",
|
||
"bytes_loaded": int, # VmRSS of the llama-server
|
||
"bytes_total": int, # sum of shard file sizes
|
||
"fraction": float, # bytes_loaded / bytes_total, 0..1
|
||
}
|
||
|
||
Linux-only in the current implementation. On macOS/Windows the
|
||
equivalent would be a different API; this returns ``None`` on
|
||
platforms where ``/proc/<pid>/status`` is unavailable.
|
||
"""
|
||
proc = self._process
|
||
if proc is None:
|
||
return None
|
||
pid = proc.pid
|
||
if pid is None:
|
||
return None
|
||
|
||
# Sum up shard sizes (primary + any extras sitting alongside).
|
||
bytes_total = 0
|
||
gguf_path = self._gguf_path
|
||
if gguf_path:
|
||
primary = Path(gguf_path)
|
||
try:
|
||
if primary.is_file():
|
||
bytes_total += primary.stat().st_size
|
||
except OSError:
|
||
pass
|
||
# Extra shards live alongside the primary with the same prefix
|
||
# before the shard index (e.g. ``-00001-of-00004.gguf``).
|
||
try:
|
||
parent = primary.parent
|
||
stem = primary.name
|
||
m = _SHARD_RE.match(stem)
|
||
prefix = m.group(1) if m else None
|
||
if prefix and parent.is_dir():
|
||
for sibling in parent.iterdir():
|
||
if (
|
||
sibling.is_file()
|
||
and sibling.name.startswith(prefix)
|
||
and sibling.name != stem
|
||
and sibling.suffix == ".gguf"
|
||
):
|
||
try:
|
||
bytes_total += sibling.stat().st_size
|
||
except OSError:
|
||
pass
|
||
except OSError:
|
||
pass
|
||
|
||
# Read VmRSS from /proc/<pid>/status. Kilobytes on Linux.
|
||
bytes_loaded = 0
|
||
try:
|
||
with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
|
||
for line in f:
|
||
if line.startswith("VmRSS:"):
|
||
kb = int(line.split()[1])
|
||
bytes_loaded = kb * 1024
|
||
break
|
||
except (FileNotFoundError, PermissionError, ValueError, OSError):
|
||
return None
|
||
|
||
phase = "ready" if self._healthy else "mmap"
|
||
fraction = 0.0
|
||
if bytes_total > 0:
|
||
fraction = min(1.0, bytes_loaded / bytes_total)
|
||
return {
|
||
"phase": phase,
|
||
"bytes_loaded": bytes_loaded,
|
||
"bytes_total": bytes_total,
|
||
"fraction": round(fraction, 4),
|
||
}
|
||
|
||
@property
|
||
def chat_template(self) -> Optional[str]:
|
||
return self._chat_template
|
||
|
||
@property
|
||
def chat_template_override(self) -> Optional[str]:
|
||
return self._chat_template_override
|
||
|
||
@property
|
||
def supports_reasoning(self) -> bool:
|
||
return self._supports_reasoning
|
||
|
||
@property
|
||
def reasoning_always_on(self) -> bool:
|
||
return self._reasoning_always_on
|
||
|
||
@property
|
||
def reasoning_style(self) -> str:
|
||
return self._reasoning_style
|
||
|
||
@property
|
||
def supports_preserve_thinking(self) -> bool:
|
||
return self._supports_preserve_thinking
|
||
|
||
@property
|
||
def reasoning_default(self) -> bool:
|
||
return self._reasoning_default
|
||
|
||
def _reasoning_kwargs(self, enable_thinking: bool) -> dict:
|
||
if self._reasoning_style == "reasoning_effort":
|
||
return {"reasoning_effort": "high" if enable_thinking else "low"}
|
||
return {"enable_thinking": enable_thinking}
|
||
|
||
def _request_reasoning_kwargs(
|
||
self,
|
||
enable_thinking: Optional[bool],
|
||
reasoning_effort: Optional[str] = None,
|
||
preserve_thinking: Optional[bool] = None,
|
||
) -> Optional[dict]:
|
||
"""Build chat_template_kwargs from per-request reasoning fields.
|
||
|
||
Produces a merged dict covering the active model's reasoning style
|
||
(``enable_thinking`` or ``reasoning_effort``) plus the independent
|
||
``preserve_thinking`` kwarg when the template supports it.
|
||
"""
|
||
kwargs: dict = {}
|
||
# Always-on reasoning models hardcode <think> tags in their template
|
||
# and do not consume enable_thinking / reasoning_effort -- skip.
|
||
if self._supports_reasoning and not self._reasoning_always_on:
|
||
if self._reasoning_style == "reasoning_effort":
|
||
if reasoning_effort in ("low", "medium", "high"):
|
||
kwargs["reasoning_effort"] = reasoning_effort
|
||
elif enable_thinking is not None:
|
||
kwargs["reasoning_effort"] = "high" if enable_thinking else "low"
|
||
else:
|
||
if enable_thinking is not None:
|
||
kwargs["enable_thinking"] = enable_thinking
|
||
if self._supports_preserve_thinking and preserve_thinking is not None:
|
||
kwargs["preserve_thinking"] = preserve_thinking
|
||
return kwargs or None
|
||
|
||
@property
|
||
def supports_tools(self) -> bool:
|
||
return self._supports_tools
|
||
|
||
@property
|
||
def cache_type_kv(self) -> Optional[str]:
|
||
return self._cache_type_kv
|
||
|
||
@property
|
||
def speculative_type(self) -> Optional[str]:
|
||
return self._speculative_type
|
||
|
||
@property
|
||
def requested_spec_mode(self) -> Optional[str]:
|
||
"""Canonical UI-facing mode the user requested (see field doc)."""
|
||
return self._requested_spec_mode
|
||
|
||
@property
|
||
def spec_draft_n_max(self) -> Optional[int]:
|
||
"""User --spec-draft-n-max override active on the load, or None
|
||
when the platform default (6 GPU / 3 CPU) is in effect."""
|
||
return self._spec_draft_n_max
|
||
|
||
# ── Binary discovery ──────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_llama_server_binary() -> Optional[str]:
|
||
"""
|
||
Locate the llama-server binary.
|
||
|
||
Search order:
|
||
1. LLAMA_SERVER_PATH environment variable (direct path to binary)
|
||
1b. UNSLOTH_LLAMA_CPP_PATH env var (custom llama.cpp install dir)
|
||
2. ~/.unsloth/llama.cpp/llama-server (make build, root dir)
|
||
3. ~/.unsloth/llama.cpp/build/bin/llama-server (cmake build, Linux)
|
||
4. ~/.unsloth/llama.cpp/build/bin/Release/llama-server.exe (cmake build, Windows)
|
||
5. ./llama.cpp/llama-server (legacy: make build, root dir)
|
||
6. ./llama.cpp/build/bin/llama-server (legacy: cmake in-tree build)
|
||
7. llama-server on PATH (system install)
|
||
8. ./bin/llama-server (legacy: extracted binary)
|
||
"""
|
||
import os
|
||
import sys
|
||
|
||
binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"
|
||
|
||
# 1. Env var — direct path to binary
|
||
env_path = os.environ.get("LLAMA_SERVER_PATH")
|
||
if env_path and Path(env_path).is_file():
|
||
return env_path
|
||
|
||
# 1b. UNSLOTH_LLAMA_CPP_PATH — custom llama.cpp install directory
|
||
custom_llama_cpp = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
||
if custom_llama_cpp:
|
||
custom_dir = Path(custom_llama_cpp)
|
||
# Root dir (make builds)
|
||
root_bin = custom_dir / binary_name
|
||
if root_bin.is_file():
|
||
return str(root_bin)
|
||
# build/bin/ (cmake builds on Linux)
|
||
cmake_bin = custom_dir / "build" / "bin" / binary_name
|
||
if cmake_bin.is_file():
|
||
return str(cmake_bin)
|
||
# build/bin/Release/ (cmake builds on Windows)
|
||
if sys.platform == "win32":
|
||
win_bin = custom_dir / "build" / "bin" / "Release" / binary_name
|
||
if win_bin.is_file():
|
||
return str(win_bin)
|
||
|
||
# 2-4. Match installer layout: env-mode -> $STUDIO_HOME/llama.cpp;
|
||
# default/HOME-redirect -> ~/.unsloth/llama.cpp (sibling of studio).
|
||
legacy_llama = Path.home() / ".unsloth" / "llama.cpp"
|
||
try:
|
||
from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433
|
||
|
||
_resolved_sr = _sr()
|
||
_legacy_studio = Path.home() / ".unsloth" / "studio"
|
||
try:
|
||
_is_legacy = _resolved_sr.resolve() == _legacy_studio.resolve()
|
||
except (OSError, ValueError):
|
||
_is_legacy = _resolved_sr == _legacy_studio
|
||
if _is_legacy:
|
||
search_roots = [legacy_llama]
|
||
else:
|
||
# why: _kill_orphaned_servers excludes the legacy root in custom
|
||
# mode; discovery must match so we never spawn a server we then
|
||
# refuse to clean up. UNSLOTH_LLAMA_CPP_PATH (handled earlier)
|
||
# is the explicit way to share a build across roots.
|
||
search_roots = [_resolved_sr / "llama.cpp"]
|
||
except (ImportError, OSError, ValueError):
|
||
search_roots = [legacy_llama]
|
||
_seen_roots: set[str] = set()
|
||
_unique_roots: list[Path] = []
|
||
for r in search_roots:
|
||
k = str(r)
|
||
if k not in _seen_roots:
|
||
_seen_roots.add(k)
|
||
_unique_roots.append(r)
|
||
for unsloth_home in _unique_roots:
|
||
home_root = unsloth_home / binary_name
|
||
if home_root.is_file():
|
||
return str(home_root)
|
||
home_linux = unsloth_home / "build" / "bin" / binary_name
|
||
if home_linux.is_file():
|
||
return str(home_linux)
|
||
if sys.platform == "win32":
|
||
home_win = unsloth_home / "build" / "bin" / "Release" / binary_name
|
||
if home_win.is_file():
|
||
return str(home_win)
|
||
|
||
# 5–6. Legacy: in-tree build (older setup.sh / setup.ps1 versions)
|
||
project_root = Path(__file__).resolve().parents[4]
|
||
# Root dir (make builds)
|
||
root_path = project_root / "llama.cpp" / binary_name
|
||
if root_path.is_file():
|
||
return str(root_path)
|
||
# build/bin/ (cmake builds)
|
||
build_path = project_root / "llama.cpp" / "build" / "bin" / binary_name
|
||
if build_path.is_file():
|
||
return str(build_path)
|
||
if sys.platform == "win32":
|
||
win_path = (
|
||
project_root / "llama.cpp" / "build" / "bin" / "Release" / binary_name
|
||
)
|
||
if win_path.is_file():
|
||
return str(win_path)
|
||
|
||
# 7. System PATH
|
||
system_path = shutil.which("llama-server")
|
||
if system_path:
|
||
return system_path
|
||
|
||
# 8. Legacy: extracted to bin/
|
||
bin_path = project_root / "bin" / binary_name
|
||
if bin_path.is_file():
|
||
return str(bin_path)
|
||
|
||
return None
|
||
|
||
# ── llama-server capability probe ─────────────────────────────
|
||
|
||
# Cached on (path, mtime); `unsloth studio update` bumps mtime.
|
||
_capability_cache: dict[tuple[str, int], dict[str, object]] = {}
|
||
|
||
@classmethod
|
||
def probe_server_capabilities(
|
||
cls, binary: Optional[str] = None
|
||
) -> dict[str, object]:
|
||
"""Parse `llama-server --help` for feature flags. Returns
|
||
{found, mtp_token, supports_mtp, ngram_mod_flavor,
|
||
supports_ngram_mod, spec_draft_n_max_flag}.
|
||
|
||
``ngram_mod_flavor`` is ``"new"`` when the binary exposes the
|
||
post-rename ``--spec-ngram-mod-n-match / -n-min / -n-max`` as
|
||
real args, ``"legacy"`` when only the pre-rename
|
||
``--spec-ngram-size-n / --draft-min / --draft-max`` are real
|
||
(the rename ships with stub removal entries for the legacy
|
||
names; we tell stubs apart by the "argument has been removed"
|
||
description), or ``None`` if neither set is usable.
|
||
|
||
``spec_draft_n_max_flag`` is the actual flag name the binary
|
||
accepts: ``--spec-draft-n-max`` on post-rename builds, or
|
||
``--draft-max`` on legacy. ``None`` means n_max cannot be set.
|
||
"""
|
||
bin_path = binary or cls._find_llama_server_binary()
|
||
if not bin_path or not Path(bin_path).is_file():
|
||
return {
|
||
"found": False,
|
||
"mtp_token": None,
|
||
"supports_mtp": False,
|
||
"ngram_mod_flavor": None,
|
||
"supports_ngram_mod": False,
|
||
"spec_draft_n_max_flag": None,
|
||
}
|
||
try:
|
||
mtime = int(Path(bin_path).stat().st_mtime)
|
||
except OSError:
|
||
mtime = 0
|
||
cache_key = (bin_path, mtime)
|
||
cached = cls._capability_cache.get(cache_key)
|
||
if cached is not None:
|
||
return cached
|
||
|
||
mtp_token: Optional[str] = None
|
||
ngram_mod_flavor: Optional[str] = None
|
||
spec_draft_n_max_flag: Optional[str] = None
|
||
try:
|
||
result = subprocess.run(
|
||
[bin_path, "--help"],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 10,
|
||
check = False,
|
||
)
|
||
help_text = (result.stdout or "") + "\n" + (result.stderr or "")
|
||
# Split into per-flag blocks: each --flag line plus its
|
||
# indented continuation lines, so the "argument has been
|
||
# removed" description sits with its flag.
|
||
blocks: dict[str, str] = {}
|
||
current_flags: list[str] = []
|
||
current_desc: list[str] = []
|
||
for line in help_text.splitlines():
|
||
stripped = line.strip()
|
||
if stripped.startswith("-") and not line.startswith(" "):
|
||
# New flag line; flush previous.
|
||
if current_flags:
|
||
desc = " ".join(current_desc)
|
||
for f in current_flags:
|
||
blocks[f] = desc
|
||
current_flags = []
|
||
current_desc = [stripped]
|
||
# Extract long-form flag tokens from the DECLARATION
|
||
# prefix only (comma-separated aliases). Stop at the
|
||
# first token that isn't itself a flag, so flag
|
||
# references inside descriptions are ignored.
|
||
for tok in re.split(r"[,\s]+", stripped):
|
||
if tok.startswith("--") and re.match(
|
||
r"--[A-Za-z][A-Za-z0-9_-]*$", tok
|
||
):
|
||
current_flags.append(tok)
|
||
elif tok.startswith("-") and len(tok) > 1:
|
||
# short alias like -fa; keep scanning aliases.
|
||
continue
|
||
else:
|
||
# First non-flag token marks end of decl.
|
||
break
|
||
else:
|
||
current_desc.append(stripped)
|
||
if current_flags:
|
||
desc = " ".join(current_desc)
|
||
for f in current_flags:
|
||
blocks[f] = desc
|
||
|
||
def _is_real(flag: str) -> bool:
|
||
"""True if the flag exists AND is not a removal stub."""
|
||
desc = blocks.get(flag)
|
||
if desc is None:
|
||
return False
|
||
return "argument has been removed" not in desc
|
||
|
||
# MTP token detection from --spec-type line.
|
||
spec_line = ""
|
||
for line in help_text.splitlines():
|
||
if "--spec-type" in line:
|
||
spec_line = line
|
||
break
|
||
# PR #22673 used draft-mtp; later renamed to mtp.
|
||
if "draft-mtp" in spec_line:
|
||
mtp_token = "draft-mtp"
|
||
elif re.search(r"[|,\[]mtp[|,\]]", spec_line):
|
||
mtp_token = "mtp"
|
||
|
||
# ngram-mod flag flavor. Post-rename builds advertise both
|
||
# the new args (real) and the legacy ones (stubs); pre-rename
|
||
# builds only have the legacy ones as real.
|
||
new_ngram_real = (
|
||
_is_real("--spec-ngram-mod-n-match")
|
||
and _is_real("--spec-ngram-mod-n-min")
|
||
and _is_real("--spec-ngram-mod-n-max")
|
||
)
|
||
legacy_ngram_real = (
|
||
_is_real("--spec-ngram-size-n")
|
||
and _is_real("--draft-max")
|
||
and _is_real("--draft-min")
|
||
)
|
||
if new_ngram_real:
|
||
ngram_mod_flavor = "new"
|
||
elif legacy_ngram_real:
|
||
ngram_mod_flavor = "legacy"
|
||
|
||
# n_max flag: prefer post-rename, fall back to legacy.
|
||
if _is_real("--spec-draft-n-max"):
|
||
spec_draft_n_max_flag = "--spec-draft-n-max"
|
||
elif _is_real("--draft-max"):
|
||
spec_draft_n_max_flag = "--draft-max"
|
||
except (OSError, subprocess.SubprocessError) as exc:
|
||
logger.debug(f"llama-server --help probe failed: {exc}")
|
||
|
||
info = {
|
||
"found": True,
|
||
"mtp_token": mtp_token,
|
||
"supports_mtp": mtp_token is not None,
|
||
"ngram_mod_flavor": ngram_mod_flavor,
|
||
"supports_ngram_mod": ngram_mod_flavor is not None,
|
||
"spec_draft_n_max_flag": spec_draft_n_max_flag,
|
||
}
|
||
cls._capability_cache[cache_key] = info
|
||
return info
|
||
|
||
# ── GPU allocation ────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _get_gguf_size_bytes(model_path: str) -> int:
|
||
"""Get total GGUF size in bytes, including split shards."""
|
||
main = Path(model_path)
|
||
total = main.stat().st_size
|
||
|
||
# Check for split shards (e.g., model-00001-of-00003.gguf)
|
||
m = _SHARD_FULL_RE.match(main.name)
|
||
if m:
|
||
prefix, _, num_total = m.group(1), m.group(2), m.group(3)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(num_total)
|
||
+ r"\.gguf$"
|
||
)
|
||
for sibling in main.parent.iterdir():
|
||
if sibling != main and sibling_pat.match(sibling.name):
|
||
total += sibling.stat().st_size
|
||
|
||
return total
|
||
|
||
@staticmethod
|
||
def _amd_apu_wants_unified_memory() -> bool:
|
||
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
|
||
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
|
||
False for discrete AMD, NVIDIA, CPU and macOS (the env hurts discrete
|
||
GPUs). ROCm reuses torch.cuda.*; the gcnArchName suffix is stripped."""
|
||
try:
|
||
import torch
|
||
|
||
if getattr(torch.version, "hip", None) is None:
|
||
return False
|
||
if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
|
||
return False
|
||
for _i in range(torch.cuda.device_count()):
|
||
try:
|
||
_arch = (
|
||
getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "")
|
||
or ""
|
||
)
|
||
except Exception:
|
||
continue
|
||
if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
|
||
return True
|
||
except Exception:
|
||
return False
|
||
return False
|
||
|
||
@staticmethod
|
||
def _get_gpu_free_memory() -> list[tuple[int, int]]:
|
||
"""Query free memory per GPU.
|
||
|
||
Order:
|
||
1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects
|
||
``CUDA_VISIBLE_DEVICES``.
|
||
2. ``torch.cuda.mem_get_info`` -- universal fallback that
|
||
works on AMD ROCm too because the HIP runtime
|
||
reuses the entire ``torch.cuda.*`` namespace. Covers the
|
||
AMD case for issue #5106 (nvidia-smi-only probe silently
|
||
returned [] on AMD hosts) and also rescues NVIDIA hosts
|
||
where ``nvidia-smi`` is missing from PATH.
|
||
|
||
Returns list of (gpu_index, free_mib) sorted by index. Empty
|
||
list if no supported GPU is reachable.
|
||
"""
|
||
import os
|
||
|
||
# ── NVIDIA via nvidia-smi ────────────────────────────────────
|
||
try:
|
||
result = subprocess.run(
|
||
[
|
||
"nvidia-smi",
|
||
"--query-gpu=index,memory.free",
|
||
"--format=csv,noheader,nounits",
|
||
],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 10,
|
||
env = child_env_without_native_path_secret(),
|
||
**_windows_hidden_subprocess_kwargs(),
|
||
)
|
||
if result.returncode == 0:
|
||
allowed: Optional[set[int]] = None
|
||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
if cvd is not None:
|
||
try:
|
||
# `if x.strip()` filters trailing-comma masks like
|
||
# "0,1," which would otherwise raise ValueError on
|
||
# an empty token. An explicitly empty mask (CVD="")
|
||
# yields an empty `allowed` set so all GPUs are
|
||
# filtered out, matching the codebase convention.
|
||
allowed = set(
|
||
int(x.strip()) for x in cvd.split(",") if x.strip()
|
||
)
|
||
except ValueError:
|
||
pass
|
||
gpus: list[tuple[int, int]] = []
|
||
for line in result.stdout.strip().splitlines():
|
||
parts = line.split(",")
|
||
if len(parts) == 2:
|
||
idx = int(parts[0].strip())
|
||
free_mib = int(parts[1].strip())
|
||
if allowed is not None and idx not in allowed:
|
||
continue
|
||
gpus.append((idx, free_mib))
|
||
# Match the docstring's sort-by-id guarantee. nvidia-smi
|
||
# almost always returns sorted output, but driver order
|
||
# is not formally guaranteed.
|
||
gpus.sort(key = lambda g: g[0])
|
||
if gpus:
|
||
return gpus
|
||
except Exception as e:
|
||
logger.debug(f"nvidia-smi probe failed: {e}")
|
||
|
||
# ── Torch fallback (covers AMD ROCm and missing nvidia-smi) ──
|
||
try:
|
||
import torch
|
||
|
||
if not hasattr(torch, "cuda") or not torch.cuda.is_available():
|
||
return []
|
||
if not hasattr(torch.cuda, "mem_get_info"):
|
||
return []
|
||
# torch.cuda enumerates GPUs RELATIVE to the visibility mask.
|
||
# On NVIDIA builds the mask is CUDA_VISIBLE_DEVICES; on AMD
|
||
# ROCm builds it is HIP_VISIBLE_DEVICES (or ROCR_VISIBLE_DEVICES
|
||
# if HIP is unset). Downstream we feed these IDs back into the
|
||
# llama-server subprocess as CVD, so we must translate visible
|
||
# ordinals back to physical indices first; otherwise launching
|
||
# with ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to
|
||
# ``CUDA_VISIBLE_DEVICES=0,1`` and target the wrong GPUs.
|
||
physical_ids: Optional[list[int]] = None
|
||
# Match the codebase convention in
|
||
# ``utils/hardware/hardware.py::_get_parent_visible_gpu_spec``:
|
||
# treat an explicitly empty mask (``HIP_VISIBLE_DEVICES=""``)
|
||
# as "set to no GPUs" rather than falling through to the next
|
||
# var. ``or`` would coerce empty string to falsy and silently
|
||
# promote the wrong source.
|
||
if getattr(torch.version, "hip", None) is not None:
|
||
hip_v = os.environ.get("HIP_VISIBLE_DEVICES")
|
||
rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES")
|
||
cvd = (
|
||
hip_v
|
||
if hip_v is not None
|
||
else rocr_v
|
||
if rocr_v is not None
|
||
else os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
)
|
||
else:
|
||
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
|
||
if cvd is not None:
|
||
try:
|
||
# Empty mask (CVD="") yields an empty list so the
|
||
# below loop produces no GPUs, consistent with the
|
||
# nvidia-smi path and utils/hardware/hardware.py.
|
||
physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()]
|
||
except ValueError:
|
||
physical_ids = None
|
||
gpus = []
|
||
for ordinal in range(torch.cuda.device_count()):
|
||
free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal)
|
||
idx = (
|
||
physical_ids[ordinal]
|
||
if physical_ids is not None and ordinal < len(physical_ids)
|
||
else ordinal
|
||
)
|
||
gpus.append((idx, free_bytes // (1024 * 1024)))
|
||
# Match the nvidia-smi path's docstring guarantee of sorted-by-id.
|
||
return sorted(gpus, key = lambda g: g[0])
|
||
except Exception as e:
|
||
logger.debug(f"torch GPU probe failed: {e}")
|
||
return []
|
||
|
||
# Skip the wait when the last kill is older than this; the GPU
|
||
# driver has already reclaimed the prior process's allocations.
|
||
_VRAM_SETTLE_WINDOW_S: float = 15.0
|
||
|
||
@staticmethod
|
||
def _wait_for_vram_settle(
|
||
max_wait: float = 2.0,
|
||
interval: float = 0.25,
|
||
tolerance_mib: int = 256,
|
||
since_kill: float = 0.0,
|
||
) -> None:
|
||
"""Poll ``_get_gpu_free_memory`` until free VRAM stabilises.
|
||
|
||
The GPU driver reclaims a dead process's allocations
|
||
asynchronously, so sampling free memory in the kill-to-spawn
|
||
window reads artificially low and pushes ``_select_gpus`` /
|
||
``_fit_context_to_vram`` toward needless CPU offload -- on a
|
||
tight VRAM card this is the Apply-reload OOM that bare-shell
|
||
launches with the same flags never see.
|
||
|
||
Short-circuits on cold start (``since_kill`` zero) or stale
|
||
kill (older than ``_VRAM_SETTLE_WINDOW_S``); also on CPU-only
|
||
hosts (empty probe), probe exceptions, and GPU-set changes.
|
||
``max_wait`` is a wall-clock bound that includes probe time,
|
||
so a wedged ``nvidia-smi`` cannot extend the reload.
|
||
"""
|
||
now = time.monotonic()
|
||
if since_kill <= 0.0:
|
||
return
|
||
if now - since_kill > LlamaCppBackend._VRAM_SETTLE_WINDOW_S:
|
||
return
|
||
deadline = now + max_wait
|
||
|
||
def _probe_or_none():
|
||
if time.monotonic() >= deadline:
|
||
return None
|
||
try:
|
||
return LlamaCppBackend._get_gpu_free_memory()
|
||
except Exception:
|
||
return None
|
||
|
||
prev = _probe_or_none()
|
||
if prev is None or not prev:
|
||
return
|
||
while time.monotonic() < deadline:
|
||
remaining = deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
return
|
||
# Clip the nap so a near-zero ``max_wait`` is respected.
|
||
time.sleep(min(interval, remaining))
|
||
curr = _probe_or_none()
|
||
if curr is None or not curr or len(curr) != len(prev):
|
||
return
|
||
prev_map = dict(prev)
|
||
stable = True
|
||
for idx, free in curr:
|
||
if idx not in prev_map:
|
||
stable = False
|
||
break
|
||
prev_free = prev_map[idx]
|
||
# Adaptive: 2 % of the larger sample dominates the
|
||
# 256 MiB floor on large-VRAM cards.
|
||
per_gpu_tol = max(tolerance_mib, int(max(free, prev_free) * 0.02))
|
||
if abs(free - prev_free) >= per_gpu_tol:
|
||
stable = False
|
||
break
|
||
if stable:
|
||
return
|
||
prev = curr
|
||
|
||
# Free-VRAM fraction at which Studio pins the GPU directly instead
|
||
# of deferring to ``--fit on``. 5% headroom covers CUDA context +
|
||
# compute buffers; 0.90 was too conservative and dropped 91-94%
|
||
# fits to CPU offload (#5106). The fork's --fit on still catches
|
||
# the truly-too-large case.
|
||
_GPU_PIN_VRAM_FRACTION = 0.95
|
||
|
||
@staticmethod
|
||
def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]:
|
||
"""Return DLL dirs from pip-installed CUDA wheels under
|
||
``<prefix>/Lib/site-packages/`` so llama-server.exe can load
|
||
``cudart64_X.dll`` / ``cublas64_X.dll`` without a system CUDA
|
||
toolkit. Mirrors the Linux ``nvidia/cu*/lib`` LD_LIBRARY_PATH
|
||
block, with parity for the Windows-specific wheel layouts seen
|
||
in the wild. Covered patterns:
|
||
* ``nvidia/<pkg>/bin`` -- legacy modular wheels
|
||
(``nvidia-cuda-runtime-cu12``, ``nvidia-cublas-cu12``, etc.).
|
||
* ``nvidia/<pkg>/bin/x86_64`` and ``.../bin/x64`` -- current
|
||
CUDA 13 wheel layout used by the unsuffixed
|
||
``nvidia-cuda-runtime`` / ``nvidia-cublas`` packages, which
|
||
ship under ``nvidia/cu13/bin/x86_64/`` (#5106).
|
||
* ``nvidia/<pkg>/Library/bin`` (and arch subdirs) -- conda-
|
||
style wheel repacks.
|
||
* ``torch/lib`` -- PyTorch's own CUDA-bundled Windows wheel,
|
||
which can ship ``cudart64_*.dll`` directly here instead of
|
||
as separate ``nvidia-*`` wheels. The install-side helper
|
||
``python_runtime_dirs`` in ``install_llama_prebuilt.py``
|
||
covers this path for the same reason.
|
||
|
||
Walks the tree with ``Path.iterdir`` rather than ``glob.glob``
|
||
so the resolver is safe against Windows paths containing
|
||
``[`` or ``]`` (valid in usernames; would otherwise be
|
||
interpreted as a glob character class and silently miss
|
||
existing dirs)."""
|
||
site_packages = Path(prefix) / "Lib" / "site-packages"
|
||
out: list[str] = []
|
||
seen: set[str] = set()
|
||
|
||
def _add(path: Path) -> None:
|
||
if not path.is_dir():
|
||
return
|
||
key = os.path.normcase(os.path.abspath(str(path)))
|
||
if key in seen:
|
||
return
|
||
seen.add(key)
|
||
out.append(str(path))
|
||
|
||
nvidia_root = site_packages / "nvidia"
|
||
if nvidia_root.is_dir():
|
||
for pkg_dir in nvidia_root.iterdir():
|
||
if not pkg_dir.is_dir():
|
||
continue
|
||
# Order matters for PATH search: arch-specific subdirs
|
||
# first so the explicit cudart64_X.dll location wins
|
||
# over a sibling ``bin`` that might be empty.
|
||
for sub in (
|
||
pkg_dir / "bin" / "x86_64",
|
||
pkg_dir / "bin" / "x64",
|
||
pkg_dir / "bin",
|
||
pkg_dir / "Library" / "bin" / "x86_64",
|
||
pkg_dir / "Library" / "bin" / "x64",
|
||
pkg_dir / "Library" / "bin",
|
||
):
|
||
_add(sub)
|
||
_add(site_packages / "torch" / "lib")
|
||
return out
|
||
|
||
@staticmethod
|
||
def _build_windows_path_dirs(
|
||
binary_dir: str, prefix: str, cuda_path: str
|
||
) -> list[str]:
|
||
"""Ordered PATH entries the win32 branch of start_llama_server
|
||
prepends so llama-server.exe resolves cudart / cublas DLLs:
|
||
binary_dir, pip nvidia wheels, CUDA_PATH/bin, CUDA_PATH/bin/x64.
|
||
Extracted so test_windows_gpu_detection_mock asserts against
|
||
production logic, not a hand-copy. #5106."""
|
||
path_dirs = [binary_dir]
|
||
path_dirs.extend(LlamaCppBackend._windows_pip_nvidia_dll_dirs(prefix))
|
||
if cuda_path:
|
||
cuda_bin = os.path.join(cuda_path, "bin")
|
||
if os.path.isdir(cuda_bin):
|
||
path_dirs.append(cuda_bin)
|
||
cuda_bin_x64 = os.path.join(cuda_path, "bin", "x64")
|
||
if os.path.isdir(cuda_bin_x64):
|
||
path_dirs.append(cuda_bin_x64)
|
||
return path_dirs
|
||
|
||
@staticmethod
|
||
def _select_gpus(
|
||
model_size_bytes: int,
|
||
gpus: list[tuple[int, int]],
|
||
) -> tuple[Optional[list[int]], bool]:
|
||
"""Pick GPU(s) for a model based on estimated VRAM and free memory.
|
||
|
||
``model_size_bytes`` should include both model weights and estimated
|
||
KV cache. The ``_GPU_PIN_VRAM_FRACTION`` threshold provides headroom
|
||
for compute buffers, CUDA context, and other runtime overhead.
|
||
|
||
Returns (gpu_indices, use_fit):
|
||
- ([1], False) model fits on 1 GPU at the headroom threshold
|
||
- ([1, 2], False) model needs 2 GPUs
|
||
- (None, True) model too large, let --fit handle it
|
||
"""
|
||
if not gpus:
|
||
return None, True
|
||
|
||
model_size_mib = model_size_bytes / (1024 * 1024)
|
||
usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
|
||
|
||
# Sort GPUs by free memory descending
|
||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
|
||
# Try fitting on 1 GPU at the usable-VRAM threshold.
|
||
if ranked[0][1] * usable_fraction >= model_size_mib:
|
||
return [ranked[0][0]], False
|
||
|
||
# Try fitting on N GPUs (accumulate free memory from most-free)
|
||
cumulative = 0
|
||
selected = []
|
||
for idx, free_mib in ranked:
|
||
selected.append(idx)
|
||
cumulative += free_mib * usable_fraction
|
||
if cumulative >= model_size_mib:
|
||
return sorted(selected), False
|
||
|
||
# Model is too large even for all GPUs, let --fit handle it
|
||
logger.debug(
|
||
"Model does not fit in available GPU memory, falling back to --fit",
|
||
model_size_mib = round(model_size_mib, 2),
|
||
ranked_gpus = ranked,
|
||
)
|
||
return None, True
|
||
|
||
# ── KV cache VRAM estimation ─────────────────────────────────────
|
||
|
||
def _can_estimate_kv(self) -> bool:
|
||
"""True if we have enough GGUF metadata to estimate KV cache size."""
|
||
if self._n_layers is None:
|
||
return False
|
||
# MLA: kv_lora_rank is sufficient (K-only cache)
|
||
if self._kv_lora_rank is not None:
|
||
return True
|
||
# New-style: need both explicit key AND value dimensions
|
||
if self._kv_key_length is not None and self._kv_value_length is not None:
|
||
return True
|
||
# Legacy: need embedding_length + a head count (scalar or per-layer).
|
||
return self._embedding_length is not None and (
|
||
self._n_kv_heads is not None
|
||
or self._n_heads is not None
|
||
or self._n_kv_heads_by_layer is not None
|
||
)
|
||
|
||
def _kv_heads_for_layer(self, layer_idx: int, fallback: int) -> int:
|
||
if self._n_kv_heads_by_layer is not None and layer_idx < len(
|
||
self._n_kv_heads_by_layer
|
||
):
|
||
return self._n_kv_heads_by_layer[layer_idx]
|
||
return fallback
|
||
|
||
def _estimate_kv_cache_bytes(
|
||
self,
|
||
n_ctx: int,
|
||
cache_type_kv: Optional[str] = None,
|
||
*,
|
||
swa_full: bool = False,
|
||
n_parallel: int = 1,
|
||
kv_unified: bool = True,
|
||
ctx_checkpoints: int = 0,
|
||
) -> int:
|
||
"""Estimate KV cache VRAM for a given context length.
|
||
|
||
Uses 5-path architecture-aware estimation:
|
||
1. MLA -- compressed KV latent + RoPE, K-only (no separate V)
|
||
2. Hybrid -- only attention layers need KV (Mamba layers don't)
|
||
3. SWA -- sliding-window layers cache min(ctx, window) tokens
|
||
4. GQA -- standard full KV with explicit key/value dimensions
|
||
5. Legacy -- fallback using embed // n_heads
|
||
|
||
Server-flag knobs (mirror llama-server's CLI):
|
||
swa_full -- ``--swa-full``: force SWA layers to cache the
|
||
full ``n_ctx`` (collapses path 3 to path 4
|
||
sizing for the SWA layers).
|
||
n_parallel -- ``--parallel``: number of server slots.
|
||
Verified empirically against llama-server:
|
||
non-SWA layers stay constant (cells split
|
||
across slots), SWA layers scale linearly
|
||
(per-slot window).
|
||
kv_unified -- ``--kv-unified`` (default on): retained for
|
||
API forward-compat. Currently a no-op for
|
||
memory math because the unified buffer total
|
||
matches per-slot buffers in measured cases.
|
||
ctx_checkpoints -- ``--ctx-checkpoints``: SWA snapshot count per
|
||
slot (PR #15293). Each snapshot stores one
|
||
sliding-window of state per SWA layer.
|
||
|
||
Returns 0 if metadata is insufficient for estimation.
|
||
"""
|
||
if not self._can_estimate_kv() or n_ctx <= 0:
|
||
return 0
|
||
|
||
n_layers = self._n_layers # type: ignore[assignment]
|
||
# Gemma 3n / Gemma 4 reuse KV from earlier layers in the last
|
||
# ``shared_kv_layers`` blocks -- those don't allocate their own
|
||
# cache. Floor at 1 so a misconfigured GGUF can't zero out KV.
|
||
shared = self._shared_kv_layers or 0
|
||
n_layers_kv = max(1, n_layers - shared)
|
||
n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment]
|
||
|
||
# Bytes per element depends on KV cache quantization
|
||
bpe = {
|
||
"f32": 4.0,
|
||
"f16": 2.0,
|
||
"bf16": 2.0,
|
||
"q8_0": 34 / 32,
|
||
"q5_1": 0.75,
|
||
"q5_0": 0.6875,
|
||
"q4_1": 0.625,
|
||
"q4_0": 0.5625,
|
||
"iq4_nl": 0.5625,
|
||
}.get(cache_type_kv or "f16", 2.0)
|
||
|
||
slots = max(1, n_parallel)
|
||
|
||
# Path 1: MLA (DeepSeek-V2/V3, GLM-4.7, GLM-5, Kimi-K2.5)
|
||
# MLA stores one compressed KV latent per token/layer (shared across heads).
|
||
# V is reconstructed from the latent on the fly -- no separate V cache.
|
||
# key_length = kv_lora_rank + rope_dim (the full compressed representation).
|
||
# MLA GGUFs set head_count_kv=1; default to 1 if absent to avoid
|
||
# falling back to n_heads (e.g., 128 for DeepSeek-V3) which would 128x.
|
||
if self._kv_lora_rank is not None:
|
||
n_kv_mla = self._n_kv_heads or 1
|
||
rope_dim = self._key_length_mla or 64
|
||
key_len = self._kv_key_length or (self._kv_lora_rank + rope_dim)
|
||
return int(n_layers_kv * n_ctx * n_kv_mla * key_len * bpe)
|
||
|
||
key_len = self._kv_key_length
|
||
val_len = self._kv_value_length
|
||
|
||
# Path 2: Hybrid Mamba/Attention (Qwen3.5-27B, Qwen3.5-35B-A3B)
|
||
# Only 1 in N layers is attention; the rest are Mamba (no KV cache).
|
||
if (
|
||
self._ssm_inner_size is not None
|
||
and self._full_attention_interval is not None
|
||
):
|
||
fai = self._full_attention_interval
|
||
n_attn = -(-n_layers // fai) if fai > 0 else n_layers # ceiling division
|
||
if key_len is not None and val_len is not None:
|
||
return int(n_attn * n_ctx * n_kv * (key_len + val_len) * bpe)
|
||
head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
|
||
return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe)
|
||
|
||
# Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...).
|
||
# Pattern is filled in by the resolver at parse time; if absent,
|
||
# falls through to the legacy 1/4-global heuristic below.
|
||
# Per-layer-type ``--parallel N`` accounting (verified empirically
|
||
# against ``llama-server``):
|
||
# * non-SWA layers: total cells = n_ctx, partitioned across
|
||
# slots -> total memory CONSTANT in slots.
|
||
# * SWA layers: per-slot cells = 2 * sliding_window
|
||
# (capped at n_ctx and at per_slot_ctx
|
||
# when ctx is split among many slots) ->
|
||
# total memory grows LINEARLY in slots.
|
||
# ``--swa-full`` forces full n_ctx for SWA layers instead.
|
||
# ``--ctx-checkpoints N`` adds N snapshots per SWA layer per slot.
|
||
if (
|
||
self._sliding_window is not None
|
||
and self._sliding_window > 0
|
||
and key_len is not None
|
||
and val_len is not None
|
||
):
|
||
swa = self._sliding_window
|
||
per_slot_ctx = max(1, n_ctx // slots)
|
||
# ``--swa-full`` makes SWA layers cache the full context just
|
||
# like non-SWA: cells get partitioned across slots, so per-slot
|
||
# cells = per_slot_ctx and the slots*per-slot product collapses
|
||
# back to the constant ``n_ctx`` total. Otherwise SWA caches
|
||
# 2*sliding_window per slot, clamped at the per-slot ctx.
|
||
swa_cells_per_slot = (
|
||
per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx)
|
||
)
|
||
key_len_swa = self._kv_key_length_swa or key_len
|
||
val_len_swa = self._kv_value_length_swa or val_len
|
||
if self._sliding_window_pattern is not None:
|
||
global_bytes = 0.0 # constant across slots
|
||
swa_bytes_per_slot = 0.0 # multiplied by slots
|
||
checkpoint_extra_per_slot = 0.0
|
||
# Iterate only over layers that allocate their own KV;
|
||
# the trailing ``shared`` layers reuse earlier caches.
|
||
for layer_idx in range(n_layers_kv):
|
||
layer_n_kv = self._kv_heads_for_layer(layer_idx, n_kv)
|
||
is_swa = (
|
||
layer_idx < len(self._sliding_window_pattern)
|
||
and self._sliding_window_pattern[layer_idx]
|
||
)
|
||
if is_swa:
|
||
swa_bytes_per_slot += (
|
||
swa_cells_per_slot
|
||
* layer_n_kv
|
||
* (key_len_swa + val_len_swa)
|
||
* bpe
|
||
)
|
||
if ctx_checkpoints > 0 and not swa_full:
|
||
checkpoint_extra_per_slot += (
|
||
ctx_checkpoints
|
||
* swa
|
||
* layer_n_kv
|
||
* (key_len_swa + val_len_swa)
|
||
* bpe
|
||
)
|
||
else:
|
||
global_bytes += n_ctx * layer_n_kv * (key_len + val_len) * bpe
|
||
return int(
|
||
global_bytes
|
||
+ slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)
|
||
)
|
||
n_global = max(1, n_layers_kv // 4)
|
||
n_swa = n_layers_kv - n_global
|
||
kv_per_token = n_kv * (key_len + val_len) * bpe
|
||
kv_per_token_swa = n_kv * (key_len_swa + val_len_swa) * bpe
|
||
global_bytes = n_global * n_ctx * kv_per_token
|
||
swa_bytes_per_slot = n_swa * swa_cells_per_slot * kv_per_token_swa
|
||
checkpoint_extra_per_slot = (
|
||
ctx_checkpoints * n_swa * swa * kv_per_token_swa
|
||
if ctx_checkpoints > 0 and not swa_full
|
||
else 0.0
|
||
)
|
||
return int(
|
||
global_bytes + slots * (swa_bytes_per_slot + checkpoint_extra_per_slot)
|
||
)
|
||
|
||
# Path 4: Standard GQA with explicit key/value dimensions
|
||
if key_len is not None and val_len is not None:
|
||
return int(n_layers_kv * n_ctx * n_kv * (key_len + val_len) * bpe)
|
||
|
||
# Path 5: Legacy fallback (old GGUFs without explicit dimensions)
|
||
head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator]
|
||
return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe)
|
||
|
||
def _fit_context_to_vram(
|
||
self,
|
||
requested_ctx: int,
|
||
available_mib: int,
|
||
model_size_bytes: int,
|
||
cache_type_kv: Optional[str] = None,
|
||
min_ctx: int = 4096,
|
||
*,
|
||
swa_full: bool = False,
|
||
n_parallel: int = 1,
|
||
kv_unified: bool = True,
|
||
ctx_checkpoints: int = 0,
|
||
kv_on_gpu: bool = True,
|
||
mtp_engaged: bool = False,
|
||
) -> int:
|
||
"""Return the largest context length that fits in GPU VRAM.
|
||
|
||
Uses 90% of available VRAM as the ctx-fit budget. Tighter than
|
||
``_GPU_PIN_VRAM_FRACTION`` on purpose: over-promising context
|
||
OOMs at runtime, while pinning conservatively just defers to
|
||
--fit on. If the weights alone don't fit, returns
|
||
``requested_ctx`` unchanged.
|
||
|
||
``kv_on_gpu`` mirrors ``--kv-offload`` (default on). When False
|
||
the KV cache lives in CPU RAM and doesn't compete with weights
|
||
for VRAM; the requested context is honored verbatim. The other
|
||
keyword args mirror ``_estimate_kv_cache_bytes``.
|
||
|
||
``mtp_engaged`` reserves extra VRAM for the MTP draft model's
|
||
KV cache + compute graph buffers. llama.cpp's MTP path keeps a
|
||
secondary cache sized off the target's KV; on tight VRAM tiers
|
||
(e.g. 32 GB) auto-fit at native context would otherwise spill
|
||
and force llama-server into a slower partial-offload path.
|
||
"""
|
||
if not self._can_estimate_kv():
|
||
logger.debug(
|
||
"Skipping context fit because KV cache metadata is unavailable",
|
||
requested_ctx = requested_ctx,
|
||
available_mib = available_mib,
|
||
)
|
||
return requested_ctx
|
||
|
||
# KV lives off-GPU: no VRAM accounting needed for the cache itself.
|
||
if not kv_on_gpu:
|
||
return requested_ctx
|
||
|
||
kv_kwargs = dict(
|
||
swa_full = swa_full,
|
||
n_parallel = n_parallel,
|
||
kv_unified = kv_unified,
|
||
ctx_checkpoints = ctx_checkpoints,
|
||
)
|
||
|
||
# MTP needs a tighter budget; drop from 0.90 to 0.85.
|
||
budget_frac = 0.85 if mtp_engaged else 0.90
|
||
budget_bytes = available_mib * 1024 * 1024 * budget_frac
|
||
model_footprint = model_size_bytes
|
||
|
||
# Check if requested context already fits
|
||
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
|
||
if model_footprint + kv <= budget_bytes:
|
||
return requested_ctx
|
||
|
||
# Model weights alone exceed budget -- can't help by reducing ctx.
|
||
# Return requested_ctx unchanged; --fit will handle VRAM management.
|
||
if model_footprint >= budget_bytes:
|
||
logger.debug(
|
||
"Model footprint exceeds GPU budget before KV cache",
|
||
requested_ctx = requested_ctx,
|
||
available_mib = available_mib,
|
||
model_size_gb = round(model_footprint / (1024**3), 2),
|
||
)
|
||
return requested_ctx
|
||
|
||
# Binary search for max context that fits
|
||
remaining = budget_bytes - model_footprint
|
||
effective_min = min(min_ctx, requested_ctx)
|
||
lo, hi = effective_min, requested_ctx
|
||
best = effective_min
|
||
while lo <= hi:
|
||
mid = (lo + hi) // 2
|
||
kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs)
|
||
if kv <= remaining:
|
||
best = mid
|
||
lo = mid + 1
|
||
else:
|
||
hi = mid - 1
|
||
|
||
# Round down to nearest 256 for alignment, but never exceed requested_ctx
|
||
best = (best // 256) * 256
|
||
best = max(effective_min, best)
|
||
best = min(best, requested_ctx)
|
||
return best
|
||
|
||
# ── Variant fallback ────────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_smallest_fitting_variant(
|
||
hf_repo: str,
|
||
free_bytes: int,
|
||
hf_token: Optional[str] = None,
|
||
) -> Optional[tuple[str, int]]:
|
||
"""Find the smallest GGUF variant (including all shards) that fits.
|
||
|
||
Groups split shards by variant prefix and sums their sizes.
|
||
For example, UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total.
|
||
|
||
Returns (first_shard_filename, total_size_bytes) or None if nothing fits.
|
||
"""
|
||
try:
|
||
from huggingface_hub import get_paths_info, list_repo_files
|
||
|
||
files = list_repo_files(hf_repo, token = hf_token)
|
||
gguf_files = [
|
||
f for f in files if f.endswith(".gguf") and "mmproj" not in f.lower()
|
||
]
|
||
if not gguf_files:
|
||
return None
|
||
|
||
# Get sizes for all GGUF files
|
||
path_infos = list(get_paths_info(hf_repo, gguf_files, token = hf_token))
|
||
size_map = {p.path: (p.size or 0) for p in path_infos}
|
||
|
||
# Group files by variant: shards share a prefix before -NNNNN-of-NNNNN
|
||
variants: dict[str, list[str]] = {}
|
||
for f in gguf_files:
|
||
m = _SHARD_RE.match(f)
|
||
key = m.group(1) if m else f
|
||
variants.setdefault(key, []).append(f)
|
||
|
||
# Sum shard sizes per variant, track the first shard (for download)
|
||
variant_sizes: list[tuple[str, int, list[str]]] = []
|
||
for key, shard_files in variants.items():
|
||
total = sum(size_map.get(f, 0) for f in shard_files)
|
||
first = sorted(shard_files)[0]
|
||
variant_sizes.append((first, total, shard_files))
|
||
|
||
# Sort by total size ascending and pick the smallest that fits
|
||
variant_sizes.sort(key = lambda x: x[1])
|
||
for first_file, total_size, _ in variant_sizes:
|
||
if total_size > 0 and total_size <= free_bytes:
|
||
return first_file, total_size
|
||
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
# ── Port allocation ───────────────────────────────────────────
|
||
|
||
@staticmethod
|
||
def _find_free_port() -> int:
|
||
"""Find an available TCP port."""
|
||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||
s.bind(("127.0.0.1", 0))
|
||
return s.getsockname()[1]
|
||
|
||
# ── Stdout drain (prevents pipe deadlock on Windows) ─────────
|
||
|
||
def _drain_stdout(self):
|
||
"""
|
||
Read lines from the subprocess stdout in a background thread.
|
||
|
||
This prevents a pipe-buffer deadlock on Windows where the default
|
||
pipe buffer is only ~4 KB. Without draining, llama-server blocks
|
||
on writes and never becomes healthy.
|
||
|
||
Each line is also teed to ``self._llama_log_fh`` when set so a
|
||
post-mortem (especially in CI) has the full subprocess output
|
||
even if the crash predates the drain-thread join in
|
||
``_wait_for_health``.
|
||
"""
|
||
try:
|
||
for line in self._process.stdout:
|
||
line = line.rstrip()
|
||
if line:
|
||
self._stdout_lines.append(line)
|
||
logger.debug(f"[llama-server] {line}")
|
||
fh = getattr(self, "_llama_log_fh", None)
|
||
if fh is not None:
|
||
try:
|
||
fh.write(line + "\n")
|
||
fh.flush()
|
||
except (ValueError, OSError):
|
||
# Log file closed under us; tee silently.
|
||
pass
|
||
except (ValueError, OSError):
|
||
# Pipe closed — process is terminating
|
||
pass
|
||
|
||
# GGUF KV type sizes for fast skipping
|
||
_GGUF_TYPE_SIZE = {
|
||
0: 1,
|
||
1: 1,
|
||
2: 2,
|
||
3: 2,
|
||
4: 4,
|
||
5: 4,
|
||
6: 4,
|
||
7: 1,
|
||
10: 8,
|
||
11: 8,
|
||
12: 8,
|
||
}
|
||
|
||
@staticmethod
|
||
def _gguf_skip_value(f, vtype: int) -> None:
|
||
"""Skip a GGUF KV value without reading it."""
|
||
sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(vtype)
|
||
if sz is not None:
|
||
f.seek(sz, 1)
|
||
elif vtype == 8: # STRING
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
f.seek(slen, 1)
|
||
elif vtype == 9: # ARRAY
|
||
atype = struct.unpack("<I", f.read(4))[0]
|
||
alen = struct.unpack("<Q", f.read(8))[0]
|
||
elem_sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(atype)
|
||
if elem_sz is not None:
|
||
f.seek(elem_sz * alen, 1)
|
||
elif atype == 8:
|
||
for _ in range(alen):
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
f.seek(slen, 1)
|
||
else:
|
||
for _ in range(alen):
|
||
LlamaCppBackend._gguf_skip_value(f, atype)
|
||
|
||
@staticmethod
|
||
def _gguf_read_array_value(f, atype: int, alen: int) -> Optional[list]:
|
||
if atype == 4: # UINT32
|
||
return [struct.unpack("<I", f.read(4))[0] for _ in range(alen)]
|
||
if atype == 5: # INT32
|
||
return [struct.unpack("<i", f.read(4))[0] for _ in range(alen)]
|
||
if atype == 7: # BOOL
|
||
return [struct.unpack("<?", f.read(1))[0] for _ in range(alen)]
|
||
|
||
for _ in range(alen):
|
||
LlamaCppBackend._gguf_skip_value(f, atype)
|
||
return None
|
||
|
||
def _read_gguf_metadata(self, gguf_path: str) -> None:
|
||
"""Read context_length, architecture params, and chat_template from a GGUF header.
|
||
|
||
Parses only the KV pairs we need (~30ms even for multi-GB files).
|
||
For split GGUFs, metadata is always in shard 1.
|
||
"""
|
||
# Reset metadata from any previously loaded model so stale flags
|
||
# (eg _supports_reasoning) do not carry over when switching models.
|
||
self._context_length = None
|
||
self._chat_template = None
|
||
self._supports_reasoning = False
|
||
self._reasoning_always_on = False
|
||
self._reasoning_style = "enable_thinking"
|
||
self._reasoning_default = True
|
||
self._supports_preserve_thinking = False
|
||
self._supports_tools = False
|
||
self._n_layers = None
|
||
self._n_kv_heads = None
|
||
self._n_kv_heads_by_layer = None
|
||
self._n_heads = None
|
||
self._embedding_length = None
|
||
self._kv_key_length = None
|
||
self._kv_value_length = None
|
||
self._sliding_window = None
|
||
self._sliding_window_pattern = None
|
||
self._full_attention_interval = None
|
||
self._kv_lora_rank = None
|
||
self._key_length_mla = None
|
||
self._kv_key_length_swa = None
|
||
self._kv_value_length_swa = None
|
||
self._ssm_inner_size = None
|
||
self._ssm_state_size = None
|
||
self._shared_kv_layers = None
|
||
self._nextn_predict_layers = None
|
||
|
||
try:
|
||
WANTED = {
|
||
"general.architecture",
|
||
"tokenizer.chat_template",
|
||
# Source-repo hints for the SWA resolver's HF fallback.
|
||
"general.source.huggingface.repository",
|
||
"general.source.url",
|
||
"general.source.repo_url",
|
||
"general.base_model.0.repo_url",
|
||
"general.base_model.0.organization",
|
||
"general.base_model.0.name",
|
||
"general.basename",
|
||
"general.organization",
|
||
"general.size_label",
|
||
"general.finetune",
|
||
}
|
||
# Additional arch-specific keys are added dynamically once
|
||
# we know the architecture name.
|
||
arch_keys: dict[str, str] = {} # gguf_key -> attribute name
|
||
arch = None
|
||
sliding_window_pattern_period: Optional[int] = None
|
||
general: dict[str, str] = {}
|
||
|
||
with open(gguf_path, "rb") as f:
|
||
magic = struct.unpack("<I", f.read(4))[0]
|
||
if magic != 0x46554747: # b"GGUF" as little-endian u32
|
||
return
|
||
_version = struct.unpack("<I", f.read(4))[0]
|
||
_tensor_count, kv_count = struct.unpack("<QQ", f.read(16))
|
||
|
||
for _ in range(kv_count):
|
||
# Tolerate truncated input (e.g., a partial header
|
||
# fetched via HTTP byte-range): bail out gracefully
|
||
# so the resolver fallback still runs on whatever
|
||
# we did manage to parse.
|
||
try:
|
||
key_len_bytes = f.read(8)
|
||
if len(key_len_bytes) < 8:
|
||
break
|
||
key_len = struct.unpack("<Q", key_len_bytes)[0]
|
||
key_bytes = f.read(key_len)
|
||
if len(key_bytes) < key_len:
|
||
break
|
||
key = key_bytes.decode("utf-8")
|
||
vtype_bytes = f.read(4)
|
||
if len(vtype_bytes) < 4:
|
||
break
|
||
vtype = struct.unpack("<I", vtype_bytes)[0]
|
||
except (struct.error, UnicodeDecodeError):
|
||
break
|
||
|
||
try:
|
||
if key in WANTED or key in arch_keys:
|
||
if vtype == 8: # STRING
|
||
slen = struct.unpack("<Q", f.read(8))[0]
|
||
val_s = f.read(slen).decode("utf-8")
|
||
if (
|
||
key.startswith("general.")
|
||
and key != "general.architecture"
|
||
):
|
||
general[key] = val_s
|
||
if key == "general.architecture":
|
||
arch = val_s
|
||
arch_keys = {
|
||
f"{arch}.context_length": "context_length",
|
||
f"{arch}.block_count": "n_layers",
|
||
f"{arch}.attention.head_count_kv": "n_kv_heads",
|
||
f"{arch}.attention.head_count": "n_heads",
|
||
f"{arch}.embedding_length": "embedding_length",
|
||
f"{arch}.attention.key_length": "kv_key_length",
|
||
f"{arch}.attention.value_length": "kv_value_length",
|
||
f"{arch}.attention.sliding_window": "sliding_window",
|
||
f"{arch}.attention.sliding_window_pattern": "sliding_window_pattern",
|
||
f"{arch}.full_attention_interval": "full_attention_interval",
|
||
f"{arch}.attention.kv_lora_rank": "kv_lora_rank",
|
||
f"{arch}.attention.key_length_mla": "key_length_mla",
|
||
f"{arch}.attention.key_length_swa": "kv_key_length_swa",
|
||
f"{arch}.attention.value_length_swa": "kv_value_length_swa",
|
||
f"{arch}.attention.shared_kv_layers": "shared_kv_layers",
|
||
f"{arch}.ssm.inner_size": "ssm_inner_size",
|
||
f"{arch}.ssm.state_size": "ssm_state_size",
|
||
f"{arch}.nextn_predict_layers": "nextn_predict_layers",
|
||
}
|
||
elif key == "tokenizer.chat_template":
|
||
self._chat_template = val_s
|
||
elif vtype in (4, 10): # UINT32 or UINT64
|
||
val_i = (
|
||
struct.unpack("<I", f.read(4))[0]
|
||
if vtype == 4
|
||
else struct.unpack("<Q", f.read(8))[0]
|
||
)
|
||
attr = arch_keys.get(key)
|
||
if attr:
|
||
if attr == "sliding_window_pattern":
|
||
sliding_window_pattern_period = val_i
|
||
else:
|
||
setattr(self, f"_{attr}", val_i)
|
||
elif vtype == 9: # ARRAY
|
||
atype = struct.unpack("<I", f.read(4))[0]
|
||
alen = struct.unpack("<Q", f.read(8))[0]
|
||
val_a = self._gguf_read_array_value(f, atype, alen)
|
||
attr = arch_keys.get(key)
|
||
if attr == "n_kv_heads" and val_a is not None:
|
||
self._n_kv_heads_by_layer = [int(x) for x in val_a]
|
||
if self._n_kv_heads is None and val_a:
|
||
self._n_kv_heads = max(int(x) for x in val_a)
|
||
elif (
|
||
attr == "sliding_window_pattern"
|
||
and val_a is not None
|
||
):
|
||
self._sliding_window_pattern = [
|
||
bool(x) for x in val_a
|
||
]
|
||
sliding_window_pattern_period = None
|
||
else:
|
||
self._gguf_skip_value(f, vtype)
|
||
else:
|
||
self._gguf_skip_value(f, vtype)
|
||
except (struct.error, UnicodeDecodeError):
|
||
# Truncated input (e.g., HTTP byte-range fetch
|
||
# of just the GGUF header); break so the
|
||
# resolver fallback still runs on what we have.
|
||
break
|
||
|
||
# Expand a scalar period straight from the GGUF first.
|
||
if (
|
||
self._sliding_window_pattern is None
|
||
and sliding_window_pattern_period
|
||
and self._n_layers
|
||
):
|
||
self._sliding_window_pattern = [
|
||
(i + 1) % sliding_window_pattern_period != 0
|
||
for i in range(self._n_layers)
|
||
]
|
||
|
||
# Otherwise hand off to the resolver (cache / bootstrap /
|
||
# transformers / HF). See `_resolve_swa_pattern`.
|
||
if (
|
||
self._sliding_window_pattern is None
|
||
and self._sliding_window
|
||
and self._n_layers
|
||
):
|
||
hf_repo_candidates = (
|
||
general.get("general.source.huggingface.repository"),
|
||
_hf_repo_from_url(general.get("general.source.url")),
|
||
_hf_repo_from_url(general.get("general.source.repo_url")),
|
||
_hf_repo_from_url(general.get("general.base_model.0.repo_url")),
|
||
(
|
||
f"{general['general.base_model.0.organization']}/"
|
||
f"{general['general.base_model.0.name']}".replace(" ", "-")
|
||
if general.get("general.base_model.0.organization")
|
||
and general.get("general.base_model.0.name")
|
||
else None
|
||
),
|
||
(
|
||
f"{general['general.organization']}/"
|
||
f"{general['general.basename']}".replace(" ", "-")
|
||
if general.get("general.organization")
|
||
and general.get("general.basename")
|
||
else None
|
||
),
|
||
)
|
||
self._sliding_window_pattern = _resolve_swa_pattern(
|
||
arch,
|
||
self._n_layers,
|
||
hf_repo_candidates,
|
||
)
|
||
|
||
if self._context_length:
|
||
logger.info(f"GGUF metadata: context_length={self._context_length}")
|
||
if self._chat_template:
|
||
logger.info(
|
||
f"GGUF metadata: chat_template={len(self._chat_template)} chars"
|
||
)
|
||
# Detect thinking/reasoning support from chat template
|
||
flags = detect_reasoning_flags(
|
||
self._chat_template,
|
||
self._model_identifier,
|
||
log_source = "GGUF metadata",
|
||
)
|
||
self._supports_reasoning = flags["supports_reasoning"]
|
||
self._reasoning_style = flags["reasoning_style"]
|
||
self._reasoning_always_on = flags["reasoning_always_on"]
|
||
self._supports_preserve_thinking = flags["supports_preserve_thinking"]
|
||
self._supports_tools = flags["supports_tools"]
|
||
except Exception as e:
|
||
logger.warning(f"Failed to read GGUF metadata: {e}")
|
||
|
||
# ── HF download (no lock held) ───────────────────────────────
|
||
|
||
def _download_gguf(
|
||
self,
|
||
*,
|
||
hf_repo: str,
|
||
hf_variant: Optional[str] = None,
|
||
hf_token: Optional[str] = None,
|
||
) -> str:
|
||
"""Download GGUF file(s) from HuggingFace. Returns local path.
|
||
|
||
Runs WITHOUT self._lock so that unload_model() can set
|
||
_cancel_event at any time. Checks _cancel_event between
|
||
each shard download.
|
||
"""
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
except ImportError:
|
||
raise RuntimeError(
|
||
"huggingface_hub is required for HF model loading. "
|
||
"Install it with: pip install huggingface_hub"
|
||
)
|
||
|
||
# Determine the filename from the variant
|
||
gguf_filename = None
|
||
gguf_extra_shards: list[str] = []
|
||
if hf_variant:
|
||
try:
|
||
from huggingface_hub import list_repo_files
|
||
|
||
files = list_repo_files(hf_repo, token = hf_token)
|
||
variant_lower = hf_variant.lower()
|
||
boundary = re.compile(
|
||
r"(?<![a-zA-Z0-9])" + re.escape(variant_lower) + r"(?![a-zA-Z0-9])"
|
||
)
|
||
gguf_files = sorted(
|
||
f
|
||
for f in files
|
||
if f.endswith(".gguf") and boundary.search(f.lower())
|
||
)
|
||
if gguf_files:
|
||
gguf_filename = gguf_files[0]
|
||
m = _SHARD_FULL_RE.match(gguf_filename)
|
||
if m:
|
||
prefix = m.group(1)
|
||
total = m.group(3)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(total)
|
||
+ r"\.gguf$"
|
||
)
|
||
gguf_extra_shards = [
|
||
f for f in gguf_files[1:] if sibling_pat.match(f)
|
||
]
|
||
except Exception as e:
|
||
logger.warning(f"Could not list repo files: {e}")
|
||
|
||
# Offline: resolve variant -> filename from the local HF cache.
|
||
# The heuristic below assumes filenames echo the repo name,
|
||
# which breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file).
|
||
# Match against the rel path (not just basename) so subdir
|
||
# layouts like ``BF16/foo.gguf`` are findable.
|
||
if not gguf_filename:
|
||
try:
|
||
from utils.models.model_config import _iter_hf_cache_snapshots
|
||
|
||
boundary = re.compile(
|
||
r"(?<![a-zA-Z0-9])"
|
||
+ re.escape(hf_variant.lower())
|
||
+ r"(?![a-zA-Z0-9])"
|
||
)
|
||
for snap in _iter_hf_cache_snapshots(hf_repo):
|
||
matches = sorted(
|
||
p.relative_to(snap).as_posix()
|
||
for p in snap.rglob("*.gguf")
|
||
if "mmproj" not in p.name.lower()
|
||
and boundary.search(p.relative_to(snap).as_posix().lower())
|
||
)
|
||
if not matches:
|
||
continue
|
||
gguf_filename = matches[0]
|
||
m = _SHARD_FULL_RE.match(Path(gguf_filename).name)
|
||
if m:
|
||
prefix = m.group(1)
|
||
total = m.group(3)
|
||
sibling_pat = re.compile(
|
||
r"^"
|
||
+ re.escape(prefix)
|
||
+ r"-\d{5}-of-"
|
||
+ re.escape(total)
|
||
+ r"\.gguf$"
|
||
)
|
||
gguf_extra_shards = [
|
||
f
|
||
for f in matches[1:]
|
||
if sibling_pat.match(Path(f).name)
|
||
]
|
||
logger.info(
|
||
"Resolved variant %s -> %s from local HF cache",
|
||
hf_variant,
|
||
gguf_filename,
|
||
)
|
||
break
|
||
except Exception as e:
|
||
logger.debug(f"Offline cache lookup for variant failed: {e}")
|
||
|
||
if not gguf_filename:
|
||
repo_name = hf_repo.split("/")[-1].replace("-GGUF", "")
|
||
gguf_filename = f"{repo_name}-{hf_variant}.gguf"
|
||
|
||
# Check disk space and fall back to a smaller variant if needed
|
||
all_gguf_files = [gguf_filename] + gguf_extra_shards
|
||
try:
|
||
from huggingface_hub import get_paths_info, try_to_load_from_cache
|
||
|
||
path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token))
|
||
total_bytes = sum((p.size or 0) for p in path_infos)
|
||
|
||
# Subtract bytes already present in the HF cache so we only
|
||
# preflight against what we actually have to download. Without
|
||
# this, re-loading a cached large model (e.g. MiniMax-M2.7-GGUF
|
||
# at 131 GB) fails cold whenever free disk is below the full
|
||
# weight footprint, even though nothing needs downloading.
|
||
already_cached_bytes = 0
|
||
for p in path_infos:
|
||
if not p.size:
|
||
continue
|
||
try:
|
||
cached_path = try_to_load_from_cache(hf_repo, p.path)
|
||
except Exception:
|
||
cached_path = None
|
||
if isinstance(cached_path, str) and os.path.exists(cached_path):
|
||
try:
|
||
on_disk = os.path.getsize(cached_path)
|
||
except OSError:
|
||
on_disk = 0
|
||
# Count as satisfied only when the full blob is present.
|
||
if on_disk >= p.size:
|
||
already_cached_bytes += p.size
|
||
|
||
total_download_bytes = max(0, total_bytes - already_cached_bytes)
|
||
|
||
if total_download_bytes > 0:
|
||
cache_dir = os.environ.get(
|
||
"HF_HUB_CACHE",
|
||
str(Path.home() / ".cache" / "huggingface" / "hub"),
|
||
)
|
||
Path(cache_dir).mkdir(parents = True, exist_ok = True)
|
||
free_bytes = shutil.disk_usage(cache_dir).free
|
||
|
||
total_gb = total_download_bytes / (1024**3)
|
||
free_gb = free_bytes / (1024**3)
|
||
cached_gb = already_cached_bytes / (1024**3)
|
||
|
||
logger.info(
|
||
f"GGUF download: {total_gb:.1f} GB needed "
|
||
f"({cached_gb:.1f} GB already cached), "
|
||
f"{free_gb:.1f} GB free on disk"
|
||
)
|
||
|
||
if total_download_bytes > free_bytes:
|
||
smaller = self._find_smallest_fitting_variant(
|
||
hf_repo,
|
||
free_bytes,
|
||
hf_token,
|
||
)
|
||
if smaller:
|
||
fallback_file, fallback_size = smaller
|
||
logger.info(
|
||
f"Selected variant too large ({total_gb:.1f} GB), "
|
||
f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)"
|
||
)
|
||
gguf_filename = fallback_file
|
||
_m = _SHARD_RE.match(gguf_filename)
|
||
_prefix = _m.group(1) if _m else None
|
||
if _prefix:
|
||
gguf_extra_shards = sorted(
|
||
f
|
||
for f in all_gguf_files
|
||
if f.startswith(_prefix)
|
||
and f != gguf_filename
|
||
and "mmproj" not in f.lower()
|
||
)
|
||
else:
|
||
gguf_extra_shards = []
|
||
else:
|
||
raise RuntimeError(
|
||
f"Not enough disk space to download any variant. "
|
||
f"Only {free_gb:.1f} GB free in {cache_dir}"
|
||
)
|
||
except RuntimeError:
|
||
raise
|
||
except Exception as e:
|
||
logger.warning(f"Could not check disk space: {e}")
|
||
|
||
gguf_label = f"{hf_repo}/{gguf_filename}" + (
|
||
f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else ""
|
||
)
|
||
logger.info(f"Resolving GGUF: {gguf_label}")
|
||
try:
|
||
if self._cancel_event.is_set():
|
||
raise RuntimeError("Cancelled")
|
||
dl_start = time.monotonic()
|
||
local_path = hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = gguf_filename,
|
||
token = hf_token,
|
||
)
|
||
for shard in gguf_extra_shards:
|
||
if self._cancel_event.is_set():
|
||
raise RuntimeError("Cancelled")
|
||
logger.info(f"Resolving GGUF shard: {shard}")
|
||
hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = shard,
|
||
token = hf_token,
|
||
)
|
||
except RuntimeError as e:
|
||
if "Cancelled" in str(e):
|
||
raise
|
||
raise RuntimeError(
|
||
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
|
||
)
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}"
|
||
)
|
||
|
||
dl_elapsed = time.monotonic() - dl_start
|
||
if dl_elapsed < 2.0:
|
||
logger.info(f"GGUF resolved from cache: {local_path}")
|
||
else:
|
||
logger.info(f"GGUF downloaded in {dl_elapsed:.1f}s: {local_path}")
|
||
return local_path
|
||
|
||
def _download_mmproj(
|
||
self,
|
||
*,
|
||
hf_repo: str,
|
||
hf_token: Optional[str] = None,
|
||
) -> Optional[str]:
|
||
"""Download the mmproj (vision projection) file from a GGUF repo.
|
||
|
||
Prefers mmproj-F16.gguf, falls back to any mmproj*.gguf file.
|
||
Returns the local path, or None if no mmproj file exists.
|
||
"""
|
||
|
||
def _pick_mmproj(candidates: list[str]) -> Optional[str]:
|
||
mmproj_files = sorted(
|
||
f
|
||
for f in candidates
|
||
if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower()
|
||
)
|
||
if not mmproj_files:
|
||
return None
|
||
for f in mmproj_files:
|
||
if f.lower().endswith("-f16.gguf"):
|
||
return f
|
||
return mmproj_files[0]
|
||
|
||
target: Optional[str] = None
|
||
try:
|
||
from huggingface_hub import list_repo_files
|
||
|
||
target = _pick_mmproj(list_repo_files(hf_repo, token = hf_token))
|
||
except Exception as e:
|
||
logger.debug(f"Could not list repo files for mmproj: {e}")
|
||
|
||
# Offline: resolve mmproj from the local HF cache snapshot, same
|
||
# shape as _download_gguf's offline fallback above.
|
||
if target is None:
|
||
try:
|
||
from utils.models.model_config import _iter_hf_cache_snapshots
|
||
|
||
for snap in _iter_hf_cache_snapshots(hf_repo):
|
||
rel_files = [
|
||
p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")
|
||
]
|
||
target = _pick_mmproj(rel_files)
|
||
if target is not None:
|
||
logger.info("Resolved mmproj %s from local HF cache", target)
|
||
break
|
||
except Exception as e:
|
||
logger.debug(f"Offline cache lookup for mmproj failed: {e}")
|
||
|
||
if target is None:
|
||
return None
|
||
|
||
try:
|
||
from huggingface_hub import hf_hub_download
|
||
|
||
logger.info(f"Downloading mmproj: {hf_repo}/{target}")
|
||
local_path = hf_hub_download(
|
||
repo_id = hf_repo,
|
||
filename = target,
|
||
token = hf_token,
|
||
)
|
||
return local_path
|
||
except Exception as e:
|
||
logger.warning(f"Could not download mmproj: {e}")
|
||
return None
|
||
|
||
def _resolve_launch_mmproj_path(
|
||
self,
|
||
*,
|
||
model_path: str,
|
||
mmproj_path: Optional[str],
|
||
) -> Optional[str]:
|
||
"""Return mmproj_path iff it exists on disk AND matches the model family.
|
||
|
||
Returns None if mmproj_path is None, missing on disk, or family-mismatched.
|
||
"""
|
||
if not mmproj_path:
|
||
return None
|
||
|
||
mmproj = Path(mmproj_path)
|
||
if not mmproj.is_file():
|
||
logger.warning(f"mmproj file not found: {mmproj_path}")
|
||
return None
|
||
|
||
from utils.models.model_config import mmproj_matches_model_family
|
||
|
||
if not mmproj_matches_model_family(model_path, str(mmproj)):
|
||
logger.warning(
|
||
f"mmproj does not match model family: model={Path(model_path).name} "
|
||
f"mmproj={mmproj.name}"
|
||
)
|
||
return None
|
||
|
||
return str(mmproj)
|
||
|
||
# ── Lifecycle ─────────────────────────────────────────────────
|
||
|
||
def load_model(
|
||
self,
|
||
*,
|
||
# Local mode: pass a path to a .gguf file
|
||
gguf_path: Optional[str] = None,
|
||
# Vision projection (mmproj) for local vision models
|
||
mmproj_path: Optional[str] = None,
|
||
# HF mode: let llama-server download via -hf "repo:quant"
|
||
hf_repo: Optional[str] = None,
|
||
hf_variant: Optional[str] = None,
|
||
hf_token: Optional[str] = None,
|
||
# Common
|
||
model_identifier: str,
|
||
is_vision: bool = False,
|
||
n_ctx: int = 4096,
|
||
chat_template_override: Optional[str] = None,
|
||
cache_type_kv: Optional[str] = None,
|
||
speculative_type: Optional[str] = None,
|
||
spec_draft_n_max: Optional[int] = None,
|
||
n_threads: Optional[int] = None,
|
||
n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused
|
||
n_parallel: int = 1,
|
||
extra_args: Optional[List[str]] = None,
|
||
) -> bool:
|
||
"""
|
||
Start llama-server with a GGUF model.
|
||
|
||
Two modes:
|
||
- Local: ``gguf_path="/path/to/model.gguf"`` → uses ``-m``
|
||
- HF: ``hf_repo="unsloth/gemma-3-4b-it-GGUF", hf_variant="Q4_K_M"`` → uses ``-hf``
|
||
|
||
In HF mode, llama-server handles downloading, caching, and
|
||
auto-loading mmproj files for vision models.
|
||
|
||
Returns True if server started and health check passed.
|
||
"""
|
||
# Serialise the whole load so concurrent /load calls never
|
||
# leave two llama-server processes alive (#5401 / #5161). Does
|
||
# not block /unload, /status, /load-progress.
|
||
with self._serial_load_lock:
|
||
# Duplicate /load that raced past the route-level check
|
||
# (the first one hadn't published _healthy=True yet). If the
|
||
# live server already satisfies this request, do nothing.
|
||
if self._already_in_target_state(
|
||
gguf_path = gguf_path,
|
||
model_identifier = model_identifier,
|
||
hf_variant = hf_variant,
|
||
n_ctx = n_ctx,
|
||
cache_type_kv = cache_type_kv,
|
||
speculative_type = speculative_type,
|
||
spec_draft_n_max = spec_draft_n_max,
|
||
chat_template_override = chat_template_override,
|
||
extra_args = extra_args,
|
||
is_vision = is_vision,
|
||
):
|
||
logger.info(
|
||
f"load_model: backend already in target state for "
|
||
f"'{model_identifier}', skipping reload"
|
||
)
|
||
# Retry probe only if a prior attempt didn't complete.
|
||
if not self._audio_probed:
|
||
try:
|
||
detected = self._detect_audio_type_strict()
|
||
self._audio_probed = True
|
||
except Exception as exc:
|
||
logger.debug("Fast-path audio probe failed: %s", exc)
|
||
detected = None
|
||
if detected in ("snac", "bicodec", "dac"):
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
try:
|
||
self.init_audio_codec(detected)
|
||
self._is_audio = True
|
||
self._audio_type = detected
|
||
except Exception as exc:
|
||
logger.warning(
|
||
"Failed to init audio codec '%s': %s",
|
||
detected,
|
||
exc,
|
||
)
|
||
self._audio_probed = False
|
||
return False
|
||
elif detected:
|
||
# csm / whisper / audio_vlm: track type but keep
|
||
# _is_audio False -- GGUF TTS routing only fires
|
||
# for snac/bicodec/dac.
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
self._audio_type = detected
|
||
if not self._healthy:
|
||
return False
|
||
return True
|
||
|
||
self._cancel_event.clear()
|
||
|
||
# ── Phase 1: kill old process (under lock, fast) ──────────
|
||
with self._lock:
|
||
self._kill_process()
|
||
|
||
binary = self._find_llama_server_binary()
|
||
if not binary:
|
||
raise RuntimeError(
|
||
"llama-server binary not found. "
|
||
"Run setup.sh to build it, install llama.cpp, "
|
||
"or set LLAMA_SERVER_PATH environment variable."
|
||
)
|
||
|
||
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
|
||
# Scope HF_HUB_OFFLINE to the download block only when DNS is
|
||
# dead; cleanup runs even on exception so a transient hiccup
|
||
# at the start of one load cannot quarantine future loads.
|
||
if hf_repo:
|
||
with _hf_offline_if_dns_dead():
|
||
model_path = self._download_gguf(
|
||
hf_repo = hf_repo,
|
||
hf_variant = hf_variant,
|
||
hf_token = hf_token,
|
||
)
|
||
# Auto-download mmproj for vision models
|
||
if is_vision and not mmproj_path:
|
||
mmproj_path = self._download_mmproj(
|
||
hf_repo = hf_repo,
|
||
hf_token = hf_token,
|
||
)
|
||
elif gguf_path:
|
||
if not Path(gguf_path).is_file():
|
||
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
|
||
model_path = gguf_path
|
||
else:
|
||
raise ValueError("Either gguf_path or hf_repo must be provided")
|
||
|
||
# Set identifier early so _read_gguf_metadata can use it for DeepSeek detection
|
||
self._model_identifier = model_identifier
|
||
|
||
# Read GGUF metadata (context_length, chat_template) -- fast, header only
|
||
self._read_gguf_metadata(model_path)
|
||
|
||
# Check cancel after download
|
||
if self._cancel_event.is_set():
|
||
logger.info("Load cancelled after download phase")
|
||
return False
|
||
|
||
# Outside ``self._lock`` so /unload, /cancel, /status are
|
||
# not blocked. ``unload_model`` also records the kill, so
|
||
# the frontend /unload+/load Apply path engages the wait
|
||
# here even though no in-process kill happened.
|
||
self._wait_for_vram_settle(since_kill = self._last_kill_monotonic)
|
||
|
||
# ── Phase 3: start llama-server (under lock) ──────────────
|
||
with self._lock:
|
||
# Re-check cancel inside lock
|
||
if self._cancel_event.is_set():
|
||
logger.info("Load cancelled before server start")
|
||
return False
|
||
|
||
self._port = self._find_free_port()
|
||
|
||
# Select GPU(s) based on model size + estimated KV cache.
|
||
# Seed safe defaults before GPU probing so the except path
|
||
# still has valid state to publish.
|
||
ctx_override = parse_ctx_override(extra_args)
|
||
requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
|
||
cache_override = parse_cache_override(extra_args)
|
||
cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv)
|
||
if ctx_override is not None and ctx_override > 0:
|
||
logger.info(
|
||
f"User --ctx-size {ctx_override} honored; "
|
||
"skipping auto-reduce"
|
||
)
|
||
if cache_override is not None:
|
||
logger.info(
|
||
f"User --cache-type-k/-v {cache_override} "
|
||
"honored for KV estimate"
|
||
)
|
||
effective_ctx = (
|
||
requested_ctx if requested_ctx > 0 else (self._context_length or 0)
|
||
)
|
||
max_available_ctx = self._context_length or effective_ctx
|
||
gpus: list[tuple[int, int]] = []
|
||
try:
|
||
model_size = self._get_gguf_size_bytes(model_path)
|
||
gpus = self._get_gpu_free_memory()
|
||
|
||
# Resolve effective context: 0 means let llama-server use the
|
||
# model's native length. Only expand to a known native length
|
||
# if metadata is available; otherwise preserve 0 as a sentinel.
|
||
if requested_ctx > 0:
|
||
effective_ctx = requested_ctx
|
||
elif self._context_length is not None:
|
||
effective_ctx = self._context_length
|
||
else:
|
||
effective_ctx = 0
|
||
original_ctx = effective_ctx
|
||
# Default UI ceiling to the model's native context length.
|
||
# GPU/VRAM-fit logic below may shrink this if hardware is limited.
|
||
max_available_ctx = self._context_length or effective_ctx
|
||
|
||
# Will MTP engage on this load? If so, the auto-fit
|
||
# budget needs to reserve extra VRAM for the draft
|
||
# model's KV cache + compute graph. Mirrors the
|
||
# canonical-mode resolver in _build_speculative_flags:
|
||
# forced mtp / mtp+ngram always engage; auto only
|
||
# engages on an MTP GGUF >= 3B (sub-3B auto falls
|
||
# back to ngram-mod which doesn't need headroom);
|
||
# ngram / ngram-simple / off never engage MTP.
|
||
_mtp_canonical = _canonicalize_spec_mode(speculative_type)
|
||
_mtp_effective = _mtp_canonical or "auto"
|
||
_mtp_size_for_fit = _extract_model_size_b(model_identifier)
|
||
_mtp_sub_3b_for_fit = (
|
||
_mtp_size_for_fit is not None and _mtp_size_for_fit < 3.0
|
||
)
|
||
_mtp_will_engage = bool(
|
||
not _extra_args_set_spec_type(extra_args)
|
||
and (
|
||
_mtp_effective in ("mtp", "mtp+ngram")
|
||
or (
|
||
_mtp_effective == "auto"
|
||
and (
|
||
bool(self._nextn_predict_layers)
|
||
or _is_mtp_model_name(model_identifier, model_path)
|
||
)
|
||
and not _mtp_sub_3b_for_fit
|
||
)
|
||
)
|
||
)
|
||
|
||
# Auto-cap context to fit in GPU VRAM and select GPUs.
|
||
#
|
||
# Two policies depending on whether the user set n_ctx:
|
||
#
|
||
# Explicit n_ctx (user chose a context length):
|
||
# Honor it. Try the full requested context with _select_gpus
|
||
# (which uses as many GPUs as needed). Only cap if it doesn't
|
||
# fit on any GPU combination.
|
||
#
|
||
# Auto n_ctx=0 (model's native context):
|
||
# Prefer fewer GPUs with reduced context over more GPUs,
|
||
# since multi-GPU is slower and the user didn't ask for a
|
||
# specific context length.
|
||
gpu_indices, use_fit = None, True
|
||
explicit_ctx = requested_ctx > 0
|
||
|
||
if gpus and self._can_estimate_kv() and effective_ctx > 0:
|
||
# Compute the largest hardware-aware cap from the model's
|
||
# native context across all usable GPU subsets (for UI
|
||
# bounds), independent of the currently requested context.
|
||
native_ctx_for_cap = self._context_length or effective_ctx
|
||
if native_ctx_for_cap > 0:
|
||
ranked_for_cap = sorted(
|
||
gpus, key = lambda g: g[1], reverse = True
|
||
)
|
||
best_cap = 0
|
||
for n_gpus in range(1, len(ranked_for_cap) + 1):
|
||
subset = ranked_for_cap[:n_gpus]
|
||
pool_mib = sum(free for _, free in subset)
|
||
capped = self._fit_context_to_vram(
|
||
native_ctx_for_cap,
|
||
pool_mib,
|
||
model_size,
|
||
cache_type_kv,
|
||
n_parallel = n_parallel,
|
||
mtp_engaged = _mtp_will_engage,
|
||
)
|
||
kv = self._estimate_kv_cache_bytes(
|
||
capped, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
total_mib = (model_size + kv) / (1024 * 1024)
|
||
if total_mib <= pool_mib * 0.90:
|
||
best_cap = max(best_cap, capped)
|
||
if best_cap > 0:
|
||
max_available_ctx = best_cap
|
||
else:
|
||
# Weights exceed 90% of every GPU subset's free
|
||
# memory, so there is no fitting context. Anchor
|
||
# the UI's "safe zone" threshold at 4096 (the
|
||
# spec's default when the model cannot fit) so
|
||
# the ctx slider shows the "might be slower"
|
||
# warning as soon as the user drags above the
|
||
# fallback default instead of never.
|
||
max_available_ctx = min(4096, native_ctx_for_cap)
|
||
|
||
if explicit_ctx:
|
||
# Honor the user's requested context verbatim. If it
|
||
# fits, pin GPUs and skip --fit; if it doesn't, ship
|
||
# -c <user_ctx> --fit on and let llama-server flex
|
||
# -ngl (CPU layer offload). The UI is expected to
|
||
# have surfaced the "might be slower" warning before
|
||
# the user submitted a ctx above the fit ceiling.
|
||
requested_total = (
|
||
model_size
|
||
+ self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
)
|
||
gpu_indices, use_fit = self._select_gpus(
|
||
requested_total, gpus
|
||
)
|
||
# No silent shrink: effective_ctx stays == requested_ctx.
|
||
else:
|
||
# Auto context: prefer fewer GPUs, cap context
|
||
# to fit. Same headroom threshold as
|
||
# _select_gpus (#5106).
|
||
ranked = sorted(gpus, key = lambda g: g[1], reverse = True)
|
||
pin_fraction = self._GPU_PIN_VRAM_FRACTION
|
||
for n_gpus in range(1, len(ranked) + 1):
|
||
subset = ranked[:n_gpus]
|
||
pool_mib = sum(free for _, free in subset)
|
||
capped = self._fit_context_to_vram(
|
||
effective_ctx,
|
||
pool_mib,
|
||
model_size,
|
||
cache_type_kv,
|
||
n_parallel = n_parallel,
|
||
mtp_engaged = _mtp_will_engage,
|
||
)
|
||
kv = self._estimate_kv_cache_bytes(
|
||
capped, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
total_mib = (model_size + kv) / (1024 * 1024)
|
||
if total_mib <= pool_mib * pin_fraction:
|
||
effective_ctx = capped
|
||
gpu_indices = sorted(idx for idx, _ in subset)
|
||
use_fit = False
|
||
break
|
||
else:
|
||
# Native ctx doesn't fit. Drop to 4096 and
|
||
# re-check before deferring to --fit on:
|
||
# a model that overflows at 131k may pin
|
||
# comfortably with a 4096 KV cache (#5106).
|
||
effective_ctx = min(4096, effective_ctx)
|
||
if effective_ctx > 0:
|
||
for n_gpus in range(1, len(ranked) + 1):
|
||
subset = ranked[:n_gpus]
|
||
pool_mib = sum(free for _, free in subset)
|
||
kv = self._estimate_kv_cache_bytes(
|
||
effective_ctx,
|
||
cache_type_kv,
|
||
n_parallel = n_parallel,
|
||
)
|
||
total_mib = (model_size + kv) / (1024 * 1024)
|
||
if total_mib <= pool_mib * pin_fraction:
|
||
gpu_indices = sorted(
|
||
idx for idx, _ in subset
|
||
)
|
||
use_fit = False
|
||
break
|
||
|
||
elif gpus:
|
||
# Can't estimate KV -- fall back to file-size-only check.
|
||
# Without KV estimation we cannot prove a hardware cap, so
|
||
# keep the ceiling at the native context (already the default).
|
||
logger.debug(
|
||
"Falling back to file-size-only GPU selection",
|
||
model_size_gb = round(model_size / (1024**3), 2),
|
||
)
|
||
gpu_indices, use_fit = self._select_gpus(model_size, gpus)
|
||
if use_fit and not explicit_ctx:
|
||
# Weights don't fit on any subset. Default the UI to
|
||
# 4096 so the slider doesn't land on an unusable native
|
||
# context. --fit on will flex -ngl at runtime.
|
||
effective_ctx = (
|
||
min(4096, effective_ctx) if effective_ctx > 0 else 4096
|
||
)
|
||
|
||
if effective_ctx < original_ctx:
|
||
kv_est = self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
logger.info(
|
||
f"Context auto-reduced: {original_ctx} -> {effective_ctx} "
|
||
f"(model: {model_size / (1024**3):.1f} GB, "
|
||
f"est. KV cache: {kv_est / (1024**3):.1f} GB)"
|
||
)
|
||
|
||
kv_cache_bytes = self._estimate_kv_cache_bytes(
|
||
effective_ctx, cache_type_kv, n_parallel = n_parallel
|
||
)
|
||
logger.info(
|
||
f"GGUF size: {model_size / (1024**3):.1f} GB, "
|
||
f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, "
|
||
f"context: {effective_ctx}, "
|
||
f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}"
|
||
)
|
||
except Exception as e:
|
||
logger.warning(f"GPU selection failed ({e}), using --fit on")
|
||
gpu_indices, use_fit = None, True
|
||
effective_ctx = requested_ctx # fall back to original
|
||
|
||
launch_mmproj_path = self._resolve_launch_mmproj_path(
|
||
model_path = model_path,
|
||
mmproj_path = mmproj_path,
|
||
)
|
||
# Need both a resolved mmproj AND the config vision flag; a stray
|
||
# mmproj passing the family-name heuristic must not flip a non-VLM
|
||
# GGUF into vision mode.
|
||
effective_is_vision = bool(launch_mmproj_path) and bool(is_vision)
|
||
if is_vision and not effective_is_vision:
|
||
logger.warning(
|
||
"Vision-capable GGUF loaded without a usable mmproj; "
|
||
"image input will be disabled for this session"
|
||
)
|
||
|
||
cmd = [
|
||
binary,
|
||
"-m",
|
||
model_path,
|
||
"--port",
|
||
str(self._port),
|
||
"-c",
|
||
str(effective_ctx) if effective_ctx > 0 else "0",
|
||
"--parallel",
|
||
str(n_parallel),
|
||
"--flash-attn",
|
||
"on", # Force flash attention for speed
|
||
# Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length".
|
||
"--no-context-shift",
|
||
]
|
||
|
||
if use_fit:
|
||
cmd.extend(["--fit", "on"])
|
||
elif gpu_indices is not None:
|
||
# Model fits on selected GPU(s) -- offload all layers
|
||
cmd.extend(["-ngl", "-1"])
|
||
|
||
# -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we
|
||
# do not inherit llama-server's internal default, which has historically
|
||
# varied (hardware concurrency incl. hyperthreads on some builds).
|
||
cmd.extend(
|
||
["--threads", str(n_threads if n_threads is not None else -1)]
|
||
)
|
||
|
||
# Always enable Jinja chat template rendering for proper template support
|
||
cmd.extend(["--jinja"])
|
||
|
||
# KV cache data type
|
||
_valid_cache_types = {
|
||
"f16",
|
||
"bf16",
|
||
"q8_0",
|
||
"q4_0",
|
||
"q4_1",
|
||
"q5_0",
|
||
"q5_1",
|
||
"iq4_nl",
|
||
"f32",
|
||
}
|
||
if cache_type_kv and cache_type_kv in _valid_cache_types:
|
||
cmd.extend(
|
||
[
|
||
"--cache-type-k",
|
||
cache_type_kv,
|
||
"--cache-type-v",
|
||
cache_type_kv,
|
||
]
|
||
)
|
||
self._cache_type_kv = cache_type_kv
|
||
logger.info(f"KV cache type: {cache_type_kv}")
|
||
else:
|
||
self._cache_type_kv = None
|
||
|
||
# Speculative decoding (n-gram self-speculation, zero VRAM cost)
|
||
# ngram-mod: ~16 MB shared hash pool, constant memory/complexity,
|
||
# variable draft lengths. Helps most when the model repeats
|
||
# existing text (code refactoring, summarization, reasoning).
|
||
# For general chat with low repetition, overhead is ~5 ms.
|
||
#
|
||
# Benchmarks from upstream llama.cpp speculative-decoding PRs:
|
||
# Scenario | Without | With | Speedup
|
||
# gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x
|
||
# Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x
|
||
# gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
|
||
#
|
||
# Params from llama.cpp server README:
|
||
# --spec-ngram-mod-n-match 24 (lookup length)
|
||
# --spec-ngram-mod-n-min 48 --spec-ngram-mod-n-max 64
|
||
# (MoEs need long drafts; dense models can reduce these)
|
||
# ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md
|
||
# ref: https://github.com/ggml-org/llama.cpp/pull/19164
|
||
# ref: https://github.com/ggml-org/llama.cpp/pull/18471
|
||
# draft-mtp: MTP heads on Unsloth's *-MTP GGUFs
|
||
# (llama.cpp #22673). Auto-enabled via nextn_predict_layers,
|
||
# fallback to -MTP in name. GPU: MTP-only. CPU/Mac: chain
|
||
# with ngram-mod. See unsloth.ai/docs/models/qwen3.6#mtp-guide.
|
||
spec_flags = self._build_speculative_flags(
|
||
speculative_type = speculative_type,
|
||
spec_draft_n_max = spec_draft_n_max,
|
||
extra_args = extra_args,
|
||
model_identifier = model_identifier,
|
||
model_path = model_path,
|
||
gpus = bool(gpus),
|
||
binary = binary,
|
||
)
|
||
cmd.extend(spec_flags)
|
||
|
||
# Apply custom chat template override if provided
|
||
self._chat_template_override = chat_template_override
|
||
if chat_template_override:
|
||
import tempfile
|
||
|
||
flags = detect_reasoning_flags(
|
||
chat_template_override,
|
||
self._model_identifier,
|
||
log_source = "GGUF chat template override",
|
||
)
|
||
self._supports_reasoning = flags["supports_reasoning"]
|
||
self._reasoning_style = flags["reasoning_style"]
|
||
self._reasoning_always_on = flags["reasoning_always_on"]
|
||
self._supports_preserve_thinking = flags[
|
||
"supports_preserve_thinking"
|
||
]
|
||
self._supports_tools = flags["supports_tools"]
|
||
|
||
self._chat_template_file = tempfile.NamedTemporaryFile(
|
||
mode = "w",
|
||
suffix = ".jinja",
|
||
delete = False,
|
||
prefix = "unsloth_chat_template_",
|
||
)
|
||
self._chat_template_file.write(chat_template_override)
|
||
self._chat_template_file.close()
|
||
cmd.extend(["--chat-template-file", self._chat_template_file.name])
|
||
logger.info(
|
||
f"Using custom chat template file: {self._chat_template_file.name}"
|
||
)
|
||
|
||
# For reasoning models, set default thinking mode.
|
||
# Qwen3.5/3.6 models below 9B (0.8B, 2B, 4B) disable thinking by default.
|
||
# Only 9B and larger enable thinking.
|
||
# Always-on templates ignore the kwarg entirely, so skip.
|
||
if self._supports_reasoning and not self._reasoning_always_on:
|
||
thinking_default = True
|
||
mid = (model_identifier or "").lower()
|
||
if "qwen3.5" in mid or "qwen3.6" in mid:
|
||
size_val = _extract_model_size_b(mid)
|
||
if size_val is not None and size_val < 9:
|
||
thinking_default = False
|
||
self._reasoning_default = thinking_default
|
||
reasoning_kw = self._reasoning_kwargs(thinking_default)
|
||
cmd.extend(
|
||
[
|
||
"--chat-template-kwargs",
|
||
json.dumps(reasoning_kw),
|
||
]
|
||
)
|
||
logger.info(f"Reasoning model: {reasoning_kw} by default")
|
||
|
||
if launch_mmproj_path and effective_is_vision:
|
||
cmd.extend(["--mmproj", launch_mmproj_path])
|
||
logger.info(f"Using mmproj for vision: {launch_mmproj_path}")
|
||
|
||
# Option C: add --api-key for direct client access when enabled
|
||
import os as _os
|
||
import secrets as _secrets
|
||
|
||
if _os.getenv("UNSLOTH_DIRECT_STREAM", "0") == "1":
|
||
self._api_key = _secrets.token_urlsafe(32)
|
||
cmd.extend(["--api-key", self._api_key])
|
||
logger.info(
|
||
"llama-server started with --api-key for direct streaming"
|
||
)
|
||
else:
|
||
self._api_key = None
|
||
|
||
# User-supplied pass-through args go last so llama.cpp's
|
||
# last-wins flag parsing lets the user override Studio's
|
||
# auto-set tier-2 flags (e.g. --cache-type-k, --spec-type).
|
||
# The route layer has already validated this list against
|
||
# the managed-flag denylist via validate_extra_args().
|
||
if extra_args:
|
||
cmd.extend(str(a) for a in extra_args)
|
||
logger.info(
|
||
f"Appending user extra args to llama-server: {list(extra_args)}"
|
||
)
|
||
|
||
_log_cmd = list(cmd)
|
||
if "--api-key" in _log_cmd:
|
||
_ki = _log_cmd.index("--api-key") + 1
|
||
if _ki < len(_log_cmd):
|
||
_log_cmd[_ki] = "<redacted>"
|
||
logger.info(f"Starting llama-server: {' '.join(_log_cmd)}")
|
||
|
||
# Set library paths so llama-server can find its shared libs and CUDA DLLs
|
||
import os
|
||
import sys
|
||
|
||
env = child_env_without_native_path_secret()
|
||
binary_dir = str(Path(binary).parent)
|
||
|
||
# AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
|
||
# shared system RAM. setdefault so a user value wins.
|
||
if self._amd_apu_wants_unified_memory():
|
||
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
|
||
logger.info(
|
||
"AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1"
|
||
)
|
||
|
||
if sys.platform == "win32":
|
||
# See _build_windows_path_dirs for ordering. #5106.
|
||
path_dirs = self._build_windows_path_dirs(
|
||
binary_dir,
|
||
sys.prefix,
|
||
os.environ.get("CUDA_PATH", ""),
|
||
)
|
||
existing_path = env.get("PATH", "")
|
||
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
|
||
|
||
# ROCm: the llama.cpp prebuilt bundles its own rocblas.dll
|
||
# but NOT the Tensile kernel library files it needs
|
||
# (rocblas/library/TensileLibrary*.dat + *.hsaco). The
|
||
# bundled DLL searches relative to its own location by
|
||
# default (i.e. <binary_dir>/rocblas/library/) which does
|
||
# not exist, causing a silent crash on the first GEMM.
|
||
# ROCBLAS_TENSILE_LIBPATH overrides that search to point at
|
||
# the ROCm installation where the kernel files actually are.
|
||
_hip_path = os.environ.get(
|
||
"HIP_PATH", os.environ.get("ROCM_PATH", "")
|
||
)
|
||
if _hip_path:
|
||
_rocblas_lib = os.path.join(
|
||
_hip_path, "bin", "rocblas", "library"
|
||
)
|
||
if os.path.isdir(_rocblas_lib):
|
||
env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
|
||
else:
|
||
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
|
||
# and CUDA runtime libs (libcudart, libcublas, etc.)
|
||
import platform
|
||
|
||
lib_dirs = [binary_dir]
|
||
_arch = platform.machine() # x86_64, aarch64, etc.
|
||
|
||
# Pip-installed nvidia CUDA runtime libs (e.g. torch's
|
||
# bundled cuda-bindings). The prebuilt llama.cpp binary
|
||
# links against libcudart.so.13 / libcublas.so.13 which
|
||
# live here, not in /usr/local/cuda.
|
||
import glob as _glob
|
||
|
||
for _nv_pattern in [
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"cu*",
|
||
"lib",
|
||
),
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"cudnn",
|
||
"lib",
|
||
),
|
||
os.path.join(
|
||
sys.prefix,
|
||
"lib",
|
||
"python*",
|
||
"site-packages",
|
||
"nvidia",
|
||
"nvjitlink",
|
||
"lib",
|
||
),
|
||
]:
|
||
for _nv_dir in _glob.glob(_nv_pattern):
|
||
if os.path.isdir(_nv_dir):
|
||
lib_dirs.append(_nv_dir)
|
||
|
||
for cuda_lib in [
|
||
"/usr/local/cuda/lib64",
|
||
f"/usr/local/cuda/targets/{_arch}-linux/lib",
|
||
# Fallback CUDA compat paths (e.g. binary built with
|
||
# CUDA 12 on a system where default /usr/local/cuda
|
||
# points to CUDA 13+).
|
||
"/usr/local/cuda-12/lib64",
|
||
"/usr/local/cuda-12.8/lib64",
|
||
f"/usr/local/cuda-12/targets/{_arch}-linux/lib",
|
||
f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib",
|
||
]:
|
||
if os.path.isdir(cuda_lib):
|
||
lib_dirs.append(cuda_lib)
|
||
existing_ld = env.get("LD_LIBRARY_PATH", "")
|
||
new_ld = ":".join(lib_dirs)
|
||
env["LD_LIBRARY_PATH"] = (
|
||
f"{new_ld}:{existing_ld}" if existing_ld else new_ld
|
||
)
|
||
|
||
# Pin to selected GPU(s). On ROCm, llama-server (and any torch
|
||
# in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES;
|
||
# narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing
|
||
# the full HIP/ROCR set the parent inherited.
|
||
if gpu_indices is not None:
|
||
pinned = ",".join(str(i) for i in gpu_indices)
|
||
env["CUDA_VISIBLE_DEVICES"] = pinned
|
||
try:
|
||
import torch as _torch
|
||
|
||
if getattr(_torch.version, "hip", None) is not None:
|
||
env["HIP_VISIBLE_DEVICES"] = pinned
|
||
env["ROCR_VISIBLE_DEVICES"] = pinned
|
||
except Exception as e:
|
||
logger.debug(
|
||
"Failed to set ROCm visibility env vars for child: %s", e
|
||
)
|
||
|
||
# Defensive kill: if a concurrent load slipped past Phase 1
|
||
# (because its `self._process` was None at the time) and
|
||
# already stored a Popen handle here, drop that orphan
|
||
# before we overwrite the reference. See issue #5161.
|
||
self._kill_process()
|
||
|
||
self._stdout_lines = []
|
||
# Tee llama-server output to a dedicated log file so a
|
||
# post-mortem in CI (or after a remote-debug session)
|
||
# has the full subprocess trail even when the parent
|
||
# only stored the last 50 lines. Path lives under the
|
||
# studio home so it ships in the same place all other
|
||
# Studio logs live.
|
||
self._llama_log_fh = None
|
||
try:
|
||
log_dir = _swa_cache_path().parent / "logs" / "llama-server"
|
||
log_dir.mkdir(parents = True, exist_ok = True)
|
||
self._llama_log_path = (
|
||
log_dir / f"llama-{int(time.time())}-port-{self._port}.log"
|
||
)
|
||
self._llama_log_fh = open(
|
||
self._llama_log_path,
|
||
"w",
|
||
encoding = "utf-8",
|
||
buffering = 1,
|
||
)
|
||
logger.info(f"llama-server stdout/stderr -> {self._llama_log_path}")
|
||
except OSError as e:
|
||
# Best-effort; never block the load on logging.
|
||
logger.debug(f"Could not open llama-server log file: {e}")
|
||
self._llama_log_path = None
|
||
self._process = subprocess.Popen(
|
||
cmd,
|
||
stdout = subprocess.PIPE,
|
||
stderr = subprocess.STDOUT,
|
||
text = True,
|
||
env = env,
|
||
**_windows_hidden_subprocess_kwargs(),
|
||
)
|
||
|
||
# Start background thread to drain stdout and prevent pipe deadlock
|
||
self._stdout_thread = threading.Thread(
|
||
target = self._drain_stdout, daemon = True, name = "llama-stdout"
|
||
)
|
||
self._stdout_thread.start()
|
||
|
||
# Store the resolved on-disk path, not the caller's kwarg. In
|
||
# HF mode the caller passes gguf_path=None and the real path
|
||
# (``model_path``) is what llama-server is actually mmap'ing.
|
||
# Downstream consumers (load_progress, log lines, etc.) need
|
||
# the path that exists on disk.
|
||
self._gguf_path = model_path
|
||
self._hf_repo = hf_repo
|
||
# For local GGUF files, extract variant from filename if not provided
|
||
if hf_variant:
|
||
self._hf_variant = hf_variant
|
||
elif gguf_path:
|
||
try:
|
||
from utils.models.model_config import _extract_quant_label
|
||
|
||
self._hf_variant = _extract_quant_label(gguf_path)
|
||
except Exception:
|
||
self._hf_variant = None
|
||
else:
|
||
self._hf_variant = None
|
||
self._is_vision = effective_is_vision
|
||
self._model_identifier = model_identifier
|
||
|
||
# Store the effective (possibly capped) context separately.
|
||
# Do NOT overwrite _context_length -- it holds the model's native
|
||
# context length from GGUF metadata and is used for display/info.
|
||
self._effective_context_length = (
|
||
effective_ctx if effective_ctx > 0 else self._context_length
|
||
)
|
||
self._max_context_length = (
|
||
max_available_ctx
|
||
if max_available_ctx > 0
|
||
else self._effective_context_length
|
||
)
|
||
|
||
# Wait for llama-server to become healthy
|
||
if not self._wait_for_health(timeout = 600.0):
|
||
self._kill_process()
|
||
_gguf = gguf_path or ""
|
||
_is_ollama = (
|
||
".studio_links" in _gguf
|
||
or os.sep + "ollama_links" + os.sep in _gguf
|
||
or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf
|
||
or (self._model_identifier or "").startswith("ollama/")
|
||
)
|
||
# Only show the Ollama-specific message when the server
|
||
# output indicates a GGUF compatibility issue, not for
|
||
# unrelated failures like OOM or missing binaries.
|
||
if _is_ollama:
|
||
_output = "\n".join(self._stdout_lines[-50:]).lower()
|
||
_gguf_compat_hints = (
|
||
"key not found",
|
||
"unknown model architecture",
|
||
"failed to load model",
|
||
)
|
||
if any(h in _output for h in _gguf_compat_hints):
|
||
raise RuntimeError(
|
||
"Some Ollama models do not work with llama.cpp. "
|
||
"Try a different model, or use this model directly through Ollama instead."
|
||
)
|
||
raise RuntimeError(
|
||
"llama-server failed to start. "
|
||
"Check that the GGUF file is valid and you have enough memory."
|
||
)
|
||
|
||
self._healthy = True
|
||
|
||
# Commit caller intent only after _healthy=True so a
|
||
# failed startup can't poison the next inheritance check.
|
||
# None keeps prior, [] clears, list sets. Source records
|
||
# the caller's hf_variant (None for local files) so the
|
||
# route's same_source check stays symmetric.
|
||
if extra_args is not None:
|
||
self._extra_args = list(extra_args)
|
||
self._extra_args_source = (model_identifier, hf_variant)
|
||
self._requested_n_ctx = int(n_ctx)
|
||
|
||
# Catch silent CPU fallback when GPU was intended (#5106).
|
||
self._gpu_offload_active = self._classify_gpu_offload(
|
||
gpu_indices is not None or use_fit, gpus or []
|
||
)
|
||
if self._gpu_offload_active is False:
|
||
logger.warning(
|
||
"llama-server appears to have loaded the model entirely "
|
||
"on CPU even though Studio detected at least one GPU. "
|
||
"This usually means the prebuilt binary's GPU backend "
|
||
"failed to load -- on Windows, cudart64_X.dll / "
|
||
"cublas64_X.dll could not be resolved. Reinstall the "
|
||
"Studio llama.cpp prebuilt or install a matching CUDA "
|
||
"toolkit (issue unslothai/unsloth#5106).",
|
||
)
|
||
|
||
logger.info(
|
||
f"llama-server ready on port {self._port} "
|
||
f"for model '{model_identifier}'"
|
||
)
|
||
|
||
# Probe outside _lock (interruptible by /unload); init inside.
|
||
self._is_audio = False
|
||
self._audio_type = None
|
||
self._audio_probed = False
|
||
try:
|
||
detected = self._detect_audio_type_strict()
|
||
self._audio_probed = True
|
||
except Exception as exc:
|
||
logger.debug("Audio probe failed: %s", exc)
|
||
detected = None
|
||
if detected in ("snac", "bicodec", "dac"):
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
try:
|
||
self.init_audio_codec(detected)
|
||
self._is_audio = True
|
||
self._audio_type = detected
|
||
except Exception as exc:
|
||
# Surface as HTTP 500 -- matches pre-PR contract.
|
||
logger.warning(
|
||
"Failed to init audio codec '%s': %s",
|
||
detected,
|
||
exc,
|
||
)
|
||
self._audio_probed = False
|
||
return False
|
||
elif detected:
|
||
# csm / whisper / audio_vlm: track type but keep _is_audio
|
||
# False -- GGUF TTS routing only fires for snac/bicodec/dac.
|
||
with self._lock:
|
||
if not self._healthy:
|
||
return False
|
||
self._audio_type = detected
|
||
|
||
if not self._healthy:
|
||
return False
|
||
return True
|
||
|
||
def _build_speculative_flags(
|
||
self,
|
||
*,
|
||
speculative_type: Optional[str],
|
||
spec_draft_n_max: Optional[int],
|
||
extra_args: Optional[List[str]],
|
||
model_identifier: str,
|
||
model_path: Optional[str],
|
||
gpus: bool,
|
||
binary: Optional[str],
|
||
) -> List[str]:
|
||
"""Return the llama-server flag list for the requested spec mode.
|
||
|
||
Side effects: sets ``self._speculative_type`` (resolved internal
|
||
emit), ``self._requested_spec_mode`` (canonical UI mode for the
|
||
status round-trip), and ``self._spec_draft_n_max`` (user override
|
||
only; None when the platform default applies).
|
||
|
||
Speculative decoding (n-gram self-speculation, zero VRAM cost):
|
||
ngram-mod uses a ~16 MB shared hash pool, constant memory /
|
||
complexity, variable draft lengths. Helps most when the model
|
||
repeats existing text (code refactor, summarisation, reasoning).
|
||
For general chat with low repetition, overhead is ~5 ms.
|
||
|
||
Benchmarks from upstream llama.cpp speculative-decoding PRs:
|
||
Scenario | Without | With | Speedup
|
||
gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x
|
||
Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x
|
||
gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x
|
||
|
||
Sub-3B dense MTP regresses vs spec-off because the draft head's
|
||
per-token cost exceeds the acceptance savings at this scale.
|
||
Q4_K_XL clean bench (each prompt once after an unrelated warmup)
|
||
on B200 + x86 CPU:
|
||
0.8B GPU: draft-mtp n=2 = 0.58x vs OFF; ngram-only = 1.10x
|
||
2B GPU: draft-mtp n=2 = 0.82x vs OFF; OFF or ngram = 1.00x
|
||
0.8B CPU: chained n=2 = 0.86x vs OFF; ngram-only = 1.19x
|
||
2B CPU: chained n=2 = 0.83x vs OFF; ngram-only = 1.01x
|
||
4B+ GPU/CPU: spec on is a net win (1.08x-1.46x).
|
||
Auto falls back to ngram-mod (zero-VRAM, near-zero idle cost on
|
||
diverse content); forced MTP variants engage anyway and just log
|
||
a warning per the user's choice.
|
||
"""
|
||
flags: List[str] = []
|
||
# Reset; emit branches re-set on the resolved emission.
|
||
self._spec_draft_n_max = None
|
||
self._speculative_type = None
|
||
|
||
# Canonical UI-facing requested mode: auto / mtp / ngram /
|
||
# mtp+ngram / off / ngram-simple. Legacy values are mapped via
|
||
# _canonicalize_spec_mode (default->auto, draft-mtp->mtp,
|
||
# ngram-mod->ngram, "ngram-mod,draft-mtp"->mtp+ngram).
|
||
canonical_mode = _canonicalize_spec_mode(speculative_type)
|
||
is_mtp_model = bool(self._nextn_predict_layers) or (
|
||
_is_mtp_model_name(model_identifier, model_path)
|
||
)
|
||
user_owns_spec_type = _extra_args_set_spec_type(extra_args)
|
||
_mtp_size_b = _extract_model_size_b(model_identifier)
|
||
_mtp_too_small = _mtp_size_b is not None and _mtp_size_b < 3.0
|
||
|
||
if user_owns_spec_type:
|
||
# User --spec-type in extra_args wins outright; suppress
|
||
# auto-emit so we don't emit a duplicate / conflicting
|
||
# spec block. Record requested mode as None.
|
||
self._requested_spec_mode = None
|
||
return flags
|
||
|
||
effective_mode = canonical_mode or "auto"
|
||
self._requested_spec_mode = effective_mode
|
||
|
||
def _resolved_draft_n_max() -> int:
|
||
# User override wins; else platform default (the B200 / x86
|
||
# clean-sweep sweet spot from PR #5582 is n=2 GPU, n=3 CPU;
|
||
# raising past 3 starts to regress on essay-style
|
||
# low-acceptance prompts).
|
||
if spec_draft_n_max is not None:
|
||
n = int(spec_draft_n_max)
|
||
self._spec_draft_n_max = n
|
||
return n
|
||
return 2 if gpus else 3
|
||
|
||
def _emit_mtp(*, chain_ngram: bool) -> bool:
|
||
"""Append --spec-type mtp[/draft-mtp][,ngram-mod] + n-max."""
|
||
caps = self.probe_server_capabilities(binary)
|
||
mtp_token = caps.get("mtp_token") if caps else None
|
||
if not mtp_token:
|
||
logger.warning(
|
||
"Requested MTP speculative decoding but "
|
||
"llama-server lacks --spec-type mtp/draft-mtp; "
|
||
"run `unsloth studio update`. Loading without "
|
||
"speculative decoding."
|
||
)
|
||
return False
|
||
draft_n_max = _resolved_draft_n_max()
|
||
n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max"
|
||
if chain_ngram:
|
||
ngram_knobs = _build_ngram_mod_flags(caps)
|
||
if ngram_knobs:
|
||
spec_value = f"ngram-mod,{mtp_token}"
|
||
else:
|
||
logger.warning(
|
||
"llama-server lacks ngram-mod tuning "
|
||
"flags; loading MTP only (no ngram chain)"
|
||
)
|
||
spec_value = mtp_token
|
||
flags.extend(
|
||
[
|
||
"--spec-type",
|
||
spec_value,
|
||
n_max_flag,
|
||
str(draft_n_max),
|
||
]
|
||
)
|
||
flags.extend(ngram_knobs)
|
||
else:
|
||
flags.extend(
|
||
[
|
||
"--spec-type",
|
||
mtp_token,
|
||
n_max_flag,
|
||
str(draft_n_max),
|
||
]
|
||
)
|
||
self._speculative_type = "draft-mtp"
|
||
chain_label = "chained ngram-mod" if chain_ngram else "MTP-only"
|
||
logger.info(f"Spec decoding: {mtp_token} ({chain_label})")
|
||
return True
|
||
|
||
def _emit_ngram_mod() -> bool:
|
||
"""Append --spec-type ngram-mod + flag-set knobs."""
|
||
ngram_caps = self.probe_server_capabilities(binary)
|
||
ngram_knobs = _build_ngram_mod_flags(ngram_caps)
|
||
flags.extend(["--spec-type", "ngram-mod"])
|
||
if not ngram_knobs:
|
||
logger.warning(
|
||
"llama-server lacks ngram-mod tuning "
|
||
"flags; loading without --spec-ngram-mod-* knobs"
|
||
)
|
||
flags.extend(ngram_knobs)
|
||
self._speculative_type = "ngram-mod"
|
||
logger.info("Spec decoding: ngram-mod")
|
||
return True
|
||
|
||
if effective_mode == "off":
|
||
return flags # nothing to emit
|
||
if effective_mode == "ngram-simple":
|
||
flags.extend(["--spec-type", "ngram-simple"])
|
||
self._speculative_type = "ngram-simple"
|
||
return flags
|
||
if effective_mode == "ngram":
|
||
_emit_ngram_mod()
|
||
return flags
|
||
if effective_mode == "mtp":
|
||
if _mtp_too_small:
|
||
logger.warning(
|
||
f"Forcing MTP on a {_mtp_size_b:.1f}B model; "
|
||
"the bench shows draft-mtp regresses below 3B. "
|
||
"Engaging anyway (user override)."
|
||
)
|
||
elif not is_mtp_model:
|
||
logger.warning(
|
||
"Forcing MTP on a non-MTP GGUF; llama-server may "
|
||
"fall back to spec-off if no nextn head is present. "
|
||
"Engaging anyway (user override)."
|
||
)
|
||
_emit_mtp(chain_ngram = False)
|
||
return flags
|
||
if effective_mode == "mtp+ngram":
|
||
if _mtp_too_small:
|
||
logger.warning(
|
||
f"Forcing MTP+Ngram on a {_mtp_size_b:.1f}B model; "
|
||
"the bench shows the chain regresses below 3B. "
|
||
"Engaging anyway (user override)."
|
||
)
|
||
elif not is_mtp_model:
|
||
logger.warning(
|
||
"Forcing MTP+Ngram on a non-MTP GGUF; llama-server "
|
||
"may fall back to ngram-only if no nextn head is "
|
||
"present. Engaging anyway (user override)."
|
||
)
|
||
_emit_mtp(chain_ngram = True)
|
||
return flags
|
||
|
||
# effective_mode == "auto": today's promotion path. llama.cpp
|
||
# #22673: MTP is compatible with mmproj, so there's no vision gate.
|
||
if is_mtp_model and not _mtp_too_small:
|
||
# GPU: MTP-only. CPU/Mac: chain ngram-mod + MTP.
|
||
_emit_mtp(chain_ngram = not gpus)
|
||
elif is_mtp_model and _mtp_too_small:
|
||
# Sub-3B fallback: drop the MTP draft head, keep ngram-mod
|
||
# when the binary supports it.
|
||
_small_caps = self.probe_server_capabilities(binary)
|
||
if _small_caps.get("supports_ngram_mod"):
|
||
logger.info(
|
||
f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
|
||
"is below the 3B speedup threshold; using ngram-mod "
|
||
"only (zero-VRAM, no draft head). Override via "
|
||
"--spec-type or the Studio Speculative Decoding "
|
||
"dropdown."
|
||
)
|
||
_emit_ngram_mod()
|
||
else:
|
||
logger.info(
|
||
f"MTP GGUF detected but model size {_mtp_size_b:.1f}B "
|
||
"is below the 3B speedup threshold and the bundled "
|
||
"llama-server does not advertise ngram-mod; "
|
||
"auto-disabling speculative decoding."
|
||
)
|
||
else:
|
||
# Non-MTP model: let llama-server choose its default strategy.
|
||
flags.append("--spec-default")
|
||
self._speculative_type = "default"
|
||
return flags
|
||
|
||
def _already_in_target_state(
|
||
self,
|
||
*,
|
||
model_identifier: str,
|
||
hf_variant: Optional[str],
|
||
n_ctx: int,
|
||
cache_type_kv: Optional[str],
|
||
speculative_type: Optional[str],
|
||
chat_template_override: Optional[str],
|
||
extra_args: Optional[List[str]],
|
||
is_vision: bool,
|
||
gguf_path: Optional[str] = None,
|
||
spec_draft_n_max: Optional[int] = None,
|
||
) -> bool:
|
||
"""True iff the live server already satisfies these load kwargs.
|
||
|
||
Mirrors ``routes/inference.py:_request_matches_loaded_settings``
|
||
but compares raw kwargs so ``load_model`` can short-circuit a
|
||
duplicate /load that raced past the route-level check (#5401).
|
||
"""
|
||
if not self.is_loaded:
|
||
return False
|
||
if (self._model_identifier or "").lower() != (model_identifier or "").lower():
|
||
return False
|
||
# Direct-file loads pass hf_variant=None while the backend
|
||
# stores an extracted filename label; compare paths instead
|
||
# to keep the guard symmetric.
|
||
if gguf_path is not None and self._gguf_path:
|
||
try:
|
||
if Path(self._gguf_path).resolve() != Path(gguf_path).resolve():
|
||
return False
|
||
except OSError:
|
||
return False
|
||
elif (self._hf_variant or "").lower() != (hf_variant or "").lower():
|
||
return False
|
||
if self._requested_n_ctx != int(n_ctx):
|
||
return False
|
||
|
||
def _norm(value):
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, str):
|
||
stripped = value.strip().lower()
|
||
return stripped or None
|
||
return value
|
||
|
||
if _norm(self._cache_type_kv) != _norm(cache_type_kv):
|
||
return False
|
||
|
||
# Compare on the canonical UI-facing mode the user requested.
|
||
# When extra_args carries --spec-type, the route-layer code paths
|
||
# bypass the dropdown anyway and the backend stores
|
||
# _requested_spec_mode = None; the request mirrors that by
|
||
# canonicalising to None.
|
||
if _extra_args_set_spec_type(extra_args):
|
||
req_mode = None
|
||
else:
|
||
req_mode = _canonicalize_spec_mode(speculative_type) or "auto"
|
||
backend_mode = self._requested_spec_mode
|
||
if req_mode != backend_mode:
|
||
return False
|
||
|
||
# spec_draft_n_max only matters when an MTP variant is actually
|
||
# engaged. Compare on the resolved spec rather than the requested
|
||
# mode so an Auto request that auto-promoted to draft-mtp under
|
||
# the hood still bounces a reload when the user changes n_max.
|
||
if (
|
||
self._speculative_type == "draft-mtp"
|
||
and spec_draft_n_max is not None
|
||
and int(spec_draft_n_max) != (self._spec_draft_n_max or 0)
|
||
):
|
||
return False
|
||
|
||
if (self._chat_template_override or None) != (chat_template_override or None):
|
||
return False
|
||
|
||
# extra_args=None means "no opinion" (inherit semantics handled
|
||
# at the route layer); only an explicit list forces equality.
|
||
if extra_args is not None:
|
||
current = list(self._extra_args) if self._extra_args is not None else []
|
||
if list(extra_args) != current:
|
||
return False
|
||
return True
|
||
|
||
def _classify_gpu_offload(
|
||
self,
|
||
expected_gpu: bool,
|
||
detected_gpus: list[tuple[int, int]],
|
||
) -> Optional[bool]:
|
||
"""True if a GPU model buffer was allocated, False if only CPU
|
||
buffers landed despite GPU intent, None when there's no signal
|
||
(no GPU detected, no buffer-size lines, etc.)."""
|
||
if not detected_gpus or not expected_gpu:
|
||
return None
|
||
# llama-server logs one ``... model buffer size = N MiB`` line
|
||
# per backend buffer; CUDA0 / ROCm0 / Metal / Vulkan0 /
|
||
# OpenCL0 / SYCL0 are GPU, CPU / CPU_Mapped are not.
|
||
gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL")
|
||
saw_buffer_line = False
|
||
saw_gpu_buffer = False
|
||
for line in self._stdout_lines:
|
||
if "model buffer size" not in line:
|
||
continue
|
||
saw_buffer_line = True
|
||
if any(marker in line for marker in gpu_markers):
|
||
saw_gpu_buffer = True
|
||
break
|
||
if not saw_buffer_line:
|
||
return None
|
||
return saw_gpu_buffer
|
||
|
||
def unload_model(self) -> bool:
|
||
"""Terminate the llama-server subprocess and cancel any in-flight download."""
|
||
self._cancel_event.set()
|
||
with self._lock:
|
||
self._kill_process()
|
||
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
|
||
self._model_identifier = None
|
||
self._gguf_path = None
|
||
self._hf_repo = None
|
||
self._hf_variant = None
|
||
self._is_vision = False
|
||
self._is_audio = False
|
||
self._audio_type = None
|
||
self._audio_probed = False
|
||
self._port = None
|
||
self._healthy = False
|
||
self._context_length = None
|
||
self._effective_context_length = None
|
||
self._max_context_length = None
|
||
self._chat_template = None
|
||
self._chat_template_override = None
|
||
self._supports_reasoning = False
|
||
self._reasoning_always_on = False
|
||
self._reasoning_style = "enable_thinking"
|
||
self._reasoning_default = True
|
||
self._supports_preserve_thinking = False
|
||
self._supports_tools = False
|
||
self._cache_type_kv = None
|
||
self._speculative_type = None
|
||
self._requested_spec_mode = None
|
||
self._spec_draft_n_max = None
|
||
self._n_layers = None
|
||
self._n_kv_heads = None
|
||
self._n_kv_heads_by_layer = None
|
||
self._n_heads = None
|
||
self._embedding_length = None
|
||
self._kv_key_length = None
|
||
self._kv_value_length = None
|
||
self._sliding_window = None
|
||
self._sliding_window_pattern = None
|
||
self._full_attention_interval = None
|
||
self._kv_lora_rank = None
|
||
self._key_length_mla = None
|
||
self._kv_key_length_swa = None
|
||
self._kv_value_length_swa = None
|
||
self._ssm_inner_size = None
|
||
self._ssm_state_size = None
|
||
self._shared_kv_layers = None
|
||
self._nextn_predict_layers = None
|
||
# Clean up temp chat template file
|
||
if hasattr(self, "_chat_template_file") and self._chat_template_file:
|
||
try:
|
||
import os
|
||
|
||
os.unlink(self._chat_template_file.name)
|
||
except Exception:
|
||
pass
|
||
self._chat_template_file = None
|
||
# Free audio codec GPU memory
|
||
if LlamaCppBackend._codec_mgr is not None:
|
||
LlamaCppBackend._codec_mgr.unload()
|
||
LlamaCppBackend._codec_mgr = None
|
||
import torch
|
||
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
return True
|
||
|
||
def _kill_process(self):
|
||
"""Terminate the subprocess if running."""
|
||
if self._process is None:
|
||
return
|
||
try:
|
||
self._process.terminate()
|
||
self._process.wait(timeout = 5)
|
||
except subprocess.TimeoutExpired:
|
||
logger.warning("llama-server did not exit on SIGTERM, sending SIGKILL")
|
||
self._process.kill()
|
||
self._process.wait(timeout = 5)
|
||
except Exception as e:
|
||
logger.warning(f"Error killing llama-server process: {e}")
|
||
finally:
|
||
self._process = None
|
||
# Clear healthy so a /load arriving during the replacement
|
||
# server's warm-up window cannot short-circuit against the
|
||
# previous server's health (#5401).
|
||
self._healthy = False
|
||
# Drives _wait_for_vram_settle in the next load_model;
|
||
# set in finally so both in-process and frontend
|
||
# /unload+/load Apply paths record the kill.
|
||
self._last_kill_monotonic = time.monotonic()
|
||
if self._stdout_thread is not None:
|
||
self._stdout_thread.join(timeout = 2)
|
||
self._stdout_thread = None
|
||
fh = getattr(self, "_llama_log_fh", None)
|
||
if fh is not None:
|
||
try:
|
||
fh.close()
|
||
except Exception:
|
||
pass
|
||
self._llama_log_fh = None
|
||
|
||
@staticmethod
|
||
def _kill_orphaned_servers():
|
||
"""Kill orphaned llama-server processes started by studio.
|
||
|
||
Only kills processes whose resolved binary lives under a known
|
||
Studio install directory (or matches an exact env-var override)
|
||
to avoid terminating unrelated llama-server instances.
|
||
|
||
Mirrors every location that _find_llama_server_binary() can
|
||
return from so that orphans from any supported install path
|
||
are still cleaned up.
|
||
|
||
Uses psutil for cross-platform support (Linux, macOS, Windows).
|
||
Falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
|
||
not installed.
|
||
"""
|
||
import os
|
||
import signal
|
||
import sys
|
||
|
||
try:
|
||
# -- Build the ownership allowlist --------------------------------
|
||
# Two kinds of matches:
|
||
# exact_binaries -- env var overrides (exact path match only)
|
||
# install_roots -- directory trees that are Studio-owned
|
||
# (binary must be *under* one of these)
|
||
install_roots: list[Path] = []
|
||
|
||
# Env-mode custom root (mirrors _find_llama_server_binary).
|
||
_is_custom_root = False
|
||
try:
|
||
from utils.paths.storage_roots import studio_root as _sr # noqa: WPS433
|
||
|
||
_resolved_sr = _sr()
|
||
_legacy_studio = Path.home() / ".unsloth" / "studio"
|
||
try:
|
||
_is_custom_root = _resolved_sr.resolve() != _legacy_studio.resolve()
|
||
except (OSError, ValueError):
|
||
_is_custom_root = _resolved_sr != _legacy_studio
|
||
if _is_custom_root:
|
||
install_roots.append(_resolved_sr / "llama.cpp")
|
||
except (ImportError, OSError, ValueError):
|
||
pass
|
||
|
||
# Primary install dir (default mode only). Env-mode skips this so
|
||
# a custom-root Studio cannot kill a concurrent default-install
|
||
# Studio's llama-server (same OS user, different install).
|
||
if not _is_custom_root:
|
||
install_roots.append(Path.home() / ".unsloth" / "llama.cpp")
|
||
|
||
# Legacy in-tree build dirs (older setup.sh versions)
|
||
project_root = Path(__file__).resolve().parents[4]
|
||
install_roots.append(project_root / "llama.cpp")
|
||
|
||
# Legacy: extracted binary
|
||
install_roots.append(project_root / "bin")
|
||
|
||
# UNSLOTH_LLAMA_CPP_PATH env var (custom install dir)
|
||
custom_dir = os.environ.get("UNSLOTH_LLAMA_CPP_PATH")
|
||
if custom_dir:
|
||
install_roots.append(Path(custom_dir))
|
||
|
||
# LLAMA_SERVER_PATH env var (exact binary path)
|
||
exact_binaries: list[Path] = []
|
||
env_binary = os.environ.get("LLAMA_SERVER_PATH")
|
||
if env_binary:
|
||
try:
|
||
exact_binaries.append(Path(env_binary).resolve())
|
||
except OSError:
|
||
pass
|
||
|
||
# Resolve all roots so is_relative_to works reliably
|
||
resolved_roots: list[Path] = []
|
||
for root in install_roots:
|
||
try:
|
||
resolved_roots.append(root.resolve())
|
||
except OSError:
|
||
pass
|
||
|
||
my_pid = os.getpid()
|
||
|
||
# -- Enumerate processes -------------------------------------------
|
||
# Prefer psutil (cross-platform). Fall back to pgrep + /proc on
|
||
# Linux when psutil is not installed.
|
||
try:
|
||
import psutil
|
||
|
||
has_psutil = True
|
||
except ImportError:
|
||
has_psutil = False
|
||
|
||
if has_psutil:
|
||
for proc in psutil.process_iter(["pid", "name", "exe"]):
|
||
try:
|
||
if proc.info["pid"] == my_pid:
|
||
continue
|
||
|
||
name = proc.info.get("name") or ""
|
||
if not name.lower().startswith("llama-server"):
|
||
continue
|
||
|
||
exe = proc.info.get("exe")
|
||
if not exe:
|
||
continue
|
||
|
||
exe_path = Path(exe).resolve()
|
||
|
||
# Check ownership: exact binary match OR binary is
|
||
# under a known install root (proper ancestry, not
|
||
# substring).
|
||
is_ours = exe_path in exact_binaries or any(
|
||
exe_path.is_relative_to(root) for root in resolved_roots
|
||
)
|
||
if not is_ours:
|
||
continue
|
||
|
||
proc.kill()
|
||
logger.info(
|
||
f"Killed orphaned llama-server process "
|
||
f"(pid={proc.info['pid']})"
|
||
)
|
||
except (
|
||
psutil.NoSuchProcess,
|
||
psutil.AccessDenied,
|
||
psutil.ZombieProcess,
|
||
):
|
||
pass
|
||
else:
|
||
# -- Fallback: pgrep + /proc/<pid>/exe (Linux only) -----------
|
||
if sys.platform != "linux":
|
||
return
|
||
result = subprocess.run(
|
||
["pgrep", "-a", "-f", "llama-server"],
|
||
capture_output = True,
|
||
text = True,
|
||
timeout = 5,
|
||
env = child_env_without_native_path_secret(),
|
||
)
|
||
if result.returncode != 0:
|
||
return
|
||
|
||
for line in result.stdout.strip().splitlines():
|
||
parts = line.strip().split(None, 1)
|
||
if len(parts) < 2:
|
||
continue
|
||
pid = int(parts[0])
|
||
if pid == my_pid:
|
||
continue
|
||
|
||
# Resolve the actual executable. /proc/<pid>/exe is a
|
||
# symlink to the real binary and avoids all cmdline-
|
||
# parsing ambiguities (spaces in paths, argv rewriting).
|
||
# Fall back to the first cmdline token when /proc is
|
||
# unavailable.
|
||
proc_exe = Path(f"/proc/{pid}/exe")
|
||
try:
|
||
binary = proc_exe.resolve(strict = True)
|
||
except (OSError, ValueError):
|
||
cmdline = parts[1]
|
||
token = cmdline.split()[0] if cmdline.strip() else ""
|
||
if not token:
|
||
continue
|
||
binary = Path(token).resolve(strict = False)
|
||
|
||
owned = binary in exact_binaries or any(
|
||
binary.is_relative_to(root) for root in resolved_roots
|
||
)
|
||
if not owned:
|
||
continue
|
||
|
||
try:
|
||
os.kill(pid, signal.SIGKILL)
|
||
logger.info(f"Killed orphaned llama-server process (pid={pid})")
|
||
except ProcessLookupError:
|
||
pass
|
||
except PermissionError:
|
||
pass
|
||
except Exception:
|
||
logger.warning("Error during orphan server cleanup", exc_info = True)
|
||
|
||
def _cleanup(self):
|
||
"""atexit handler to ensure llama-server is terminated."""
|
||
self._kill_process()
|
||
|
||
def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool:
|
||
"""
|
||
Poll llama-server's /health endpoint until it responds 200.
|
||
|
||
Also monitors subprocess for early exit/crash.
|
||
"""
|
||
deadline = time.monotonic() + timeout
|
||
url = f"http://127.0.0.1:{self._port}/health"
|
||
|
||
while time.monotonic() < deadline:
|
||
# Check if process crashed
|
||
if self._process.poll() is not None:
|
||
# Give the drain thread a moment to collect final output
|
||
if self._stdout_thread is not None:
|
||
self._stdout_thread.join(timeout = 2)
|
||
output = "\n".join(self._stdout_lines[-50:])
|
||
logger.error(
|
||
f"llama-server exited with code {self._process.returncode}. "
|
||
f"Output: {output[:2000]}"
|
||
)
|
||
return False
|
||
|
||
try:
|
||
resp = httpx.get(url, timeout = 2.0)
|
||
if resp.status_code == 200:
|
||
return True
|
||
except (
|
||
httpx.ConnectError,
|
||
httpx.TimeoutException,
|
||
# ReadError covers TCP RST mid-read while llama-server is
|
||
# still binding the port (Windows: WinError 10054). The
|
||
# crash-detection branch above catches a real exit; this
|
||
# one keeps a transient socket close from masking it.
|
||
httpx.ReadError,
|
||
httpx.RemoteProtocolError,
|
||
httpx.WriteError,
|
||
):
|
||
pass
|
||
|
||
time.sleep(interval)
|
||
|
||
logger.error(f"llama-server health check timed out after {timeout}s")
|
||
return False
|
||
|
||
# ── Message building (OpenAI format) ──────────────────────────
|
||
|
||
@staticmethod
|
||
def _parse_tool_calls_from_text(content: str) -> list[dict]:
|
||
"""Thin wrapper around the shared parser in tool_call_parser
|
||
so safetensors and llama_cpp pick up the same fixes."""
|
||
return _shared_parse_tool_calls_from_text(content)
|
||
|
||
@staticmethod
|
||
def _build_openai_messages(
|
||
messages: list[dict],
|
||
image_b64: Optional[str] = None,
|
||
) -> list[dict]:
|
||
"""
|
||
Build OpenAI-format messages, optionally injecting an image_url
|
||
content part into the last user message for vision models.
|
||
|
||
If no image is provided, returns messages as-is.
|
||
"""
|
||
if not image_b64:
|
||
return messages
|
||
|
||
# Find the last user message and convert to multimodal content parts
|
||
result = [msg.copy() for msg in messages]
|
||
last_user_idx = None
|
||
for i, msg in enumerate(result):
|
||
if msg["role"] == "user":
|
||
last_user_idx = i
|
||
|
||
if last_user_idx is not None:
|
||
text_content = result[last_user_idx].get("content", "")
|
||
result[last_user_idx]["content"] = [
|
||
{"type": "text", "text": text_content},
|
||
{
|
||
"type": "image_url",
|
||
"image_url": {
|
||
"url": f"data:image/png;base64,{image_b64}",
|
||
},
|
||
},
|
||
]
|
||
|
||
return result
|
||
|
||
# ── Generation (proxy to llama-server) ────────────────────────
|
||
|
||
@staticmethod
|
||
def _iter_text_cancellable(
|
||
response: "httpx.Response",
|
||
cancel_event: Optional[threading.Event] = None,
|
||
) -> Generator[str, None, None]:
|
||
"""Iterate over an httpx streaming response with cancel support.
|
||
|
||
Checks cancel_event between chunks and on ReadTimeout. The
|
||
cancel watcher in _stream_with_retry also calls response.close()
|
||
on cancel, which unblocks iter_text() once the response exists.
|
||
During normal streaming llama-server sends tokens frequently,
|
||
so the cancel check between chunks is the primary mechanism.
|
||
"""
|
||
text_iter = response.iter_text()
|
||
while True:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
response.close()
|
||
return
|
||
try:
|
||
chunk = next(text_iter)
|
||
yield chunk
|
||
except StopIteration:
|
||
return
|
||
except httpx.ReadTimeout:
|
||
# No data within the timeout window -- just loop back
|
||
# and re-check cancel_event.
|
||
continue
|
||
|
||
@staticmethod
|
||
@contextlib.contextmanager
|
||
def _stream_with_retry(
|
||
client: "httpx.Client",
|
||
url: str,
|
||
payload: dict,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
headers: Optional[dict] = None,
|
||
):
|
||
"""Open an httpx streaming POST with cancel support.
|
||
|
||
Sends the request once with a long read timeout (120 s) so
|
||
prompt processing (prefill) can finish without triggering a
|
||
retry storm. The previous 0.5 s timeout caused duplicate POST
|
||
requests every half second, forcing llama-server to restart
|
||
processing each time.
|
||
|
||
A background watcher thread provides cancel by closing the
|
||
response when cancel_event is set. Limitation: httpx does not
|
||
allow interrupting a blocked read from another thread before
|
||
the response object exists, so cancel during the initial
|
||
header wait (prefill phase) only takes effect once headers
|
||
arrive. After that, response.close() unblocks reads promptly.
|
||
In practice llama-server prefill is 1-5 s for typical prompts,
|
||
during which cancel is deferred -- still much better than the
|
||
old retry storm which made prefill slower.
|
||
"""
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
|
||
# Background watcher: close the response if cancel is requested.
|
||
# Only effective after response headers arrive (httpx limitation).
|
||
_cancel_closed = threading.Event()
|
||
_response_ref: list = [None]
|
||
|
||
def _cancel_watcher():
|
||
while not _cancel_closed.is_set():
|
||
if cancel_event.wait(timeout = 0.3):
|
||
# Cancel requested. Keep polling until the response object
|
||
# exists so we can close it, or until the main thread
|
||
# finishes on its own (_cancel_closed is set in finally).
|
||
while not _cancel_closed.is_set():
|
||
r = _response_ref[0]
|
||
if r is not None:
|
||
try:
|
||
r.close()
|
||
return
|
||
except Exception as e:
|
||
logger.debug(
|
||
f"Error closing response in cancel watcher: {e}"
|
||
)
|
||
# Response not created yet -- wait briefly and retry
|
||
_cancel_closed.wait(timeout = 0.1)
|
||
return
|
||
|
||
watcher = None
|
||
if cancel_event is not None:
|
||
watcher = threading.Thread(
|
||
target = _cancel_watcher, daemon = True, name = "prefill-cancel"
|
||
)
|
||
watcher.start()
|
||
|
||
try:
|
||
# Long read timeout so prefill (prompt processing) can finish
|
||
# without triggering a retry storm. Cancel during both
|
||
# prefill and streaming is handled by the watcher thread
|
||
# which closes the response, unblocking any httpx read.
|
||
prefill_timeout = httpx.Timeout(
|
||
connect = 30,
|
||
read = 120.0,
|
||
write = 10,
|
||
pool = 10,
|
||
)
|
||
with client.stream(
|
||
"POST",
|
||
url,
|
||
json = payload,
|
||
timeout = prefill_timeout,
|
||
headers = headers,
|
||
) as response:
|
||
_response_ref[0] = response
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
yield response
|
||
return
|
||
except (httpx.ReadError, httpx.RemoteProtocolError, httpx.CloseError):
|
||
# Response was closed by the cancel watcher
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
raise GeneratorExit
|
||
raise
|
||
finally:
|
||
_cancel_closed.set()
|
||
|
||
def generate_chat_completion(
|
||
self,
|
||
messages: list[dict],
|
||
image_b64: Optional[str] = None,
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 20,
|
||
min_p: float = 0.01,
|
||
max_tokens: Optional[int] = None,
|
||
repetition_penalty: float = 1.0,
|
||
presence_penalty: float = 0.0,
|
||
stop: Optional[list[str]] = None,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
enable_thinking: Optional[bool] = None,
|
||
reasoning_effort: Optional[str] = None,
|
||
preserve_thinking: Optional[bool] = None,
|
||
) -> Generator[str | dict, None, None]:
|
||
"""
|
||
Send a chat completion request to llama-server and stream tokens back.
|
||
|
||
Uses /v1/chat/completions — llama-server handles chat template
|
||
application and vision (multimodal image_url parts) natively.
|
||
|
||
Yields cumulative text (matching InferenceBackend's convention).
|
||
"""
|
||
if not self.is_loaded:
|
||
raise RuntimeError("llama-server is not loaded")
|
||
|
||
openai_messages = self._build_openai_messages(messages, image_b64)
|
||
|
||
payload = {
|
||
"messages": openai_messages,
|
||
"stream": True,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
}
|
||
# Pass enable_thinking / reasoning_effort / preserve_thinking per-request
|
||
_reasoning_kw = self._request_reasoning_kwargs(
|
||
enable_thinking, reasoning_effort, preserve_thinking
|
||
)
|
||
if _reasoning_kw is not None:
|
||
payload["chat_template_kwargs"] = _reasoning_kw
|
||
# Default cap to the model's effective context length when known,
|
||
# otherwise the conservative floor. The wall-clock backstop below
|
||
# keeps a stuck model from running indefinitely either way.
|
||
payload["max_tokens"] = (
|
||
max_tokens
|
||
if max_tokens is not None
|
||
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
||
)
|
||
payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
|
||
if stop:
|
||
payload["stop"] = stop
|
||
payload["stream_options"] = {"include_usage": True}
|
||
|
||
url = f"{self.base_url}/v1/chat/completions"
|
||
cumulative = ""
|
||
in_thinking = False
|
||
_stream_done = False
|
||
_metadata_usage = None
|
||
_metadata_timings = None
|
||
|
||
try:
|
||
# _stream_with_retry uses a 120 s read timeout so prefill
|
||
# can finish. Cancel during streaming is handled by the
|
||
# watcher thread (closes the response on cancel_event).
|
||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(
|
||
timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
|
||
) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: {error_body}"
|
||
)
|
||
|
||
buffer = ""
|
||
has_content_tokens = False
|
||
reasoning_text = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response, cancel_event
|
||
):
|
||
buffer += raw_chunk
|
||
while "\n" in buffer:
|
||
line, buffer = buffer.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
if in_thinking:
|
||
if has_content_tokens:
|
||
# Real thinking + content: close the tag
|
||
cumulative += "</think>"
|
||
yield cumulative
|
||
else:
|
||
# Only reasoning_content, no content tokens:
|
||
# the model put its entire reply in reasoning
|
||
# (e.g. Qwen3 always-think mode). Show it
|
||
# as the main response, not as a thinking block.
|
||
cumulative = reasoning_text
|
||
yield cumulative
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
data = json.loads(line[6:])
|
||
# Capture server timings/usage from final chunks
|
||
_chunk_timings = data.get("timings")
|
||
if _chunk_timings:
|
||
_metadata_timings = _chunk_timings
|
||
_chunk_usage = data.get("usage")
|
||
if _chunk_usage:
|
||
_metadata_usage = _chunk_usage
|
||
choices = data.get("choices", [])
|
||
if choices:
|
||
delta = choices[0].get("delta", {})
|
||
|
||
# Handle reasoning/thinking tokens
|
||
# llama-server sends these as "reasoning_content"
|
||
# Wrap in <think> tags for the frontend parser
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_text += reasoning
|
||
if not in_thinking:
|
||
cumulative += "<think>"
|
||
in_thinking = True
|
||
cumulative += reasoning
|
||
yield cumulative
|
||
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
if in_thinking:
|
||
cumulative += "</think>"
|
||
in_thinking = False
|
||
cumulative += token
|
||
yield cumulative
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
if _metadata_usage or _metadata_timings:
|
||
_metadata_usage = _backfill_usage_from_timings(
|
||
_metadata_usage, _metadata_timings
|
||
)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": _metadata_usage,
|
||
"timings": _metadata_timings,
|
||
}
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── Tool-calling agentic loop ──────────────────────────────
|
||
|
||
def generate_chat_completion_with_tools(
|
||
self,
|
||
messages: list[dict],
|
||
tools: list[dict],
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 20,
|
||
min_p: float = 0.01,
|
||
max_tokens: Optional[int] = None,
|
||
repetition_penalty: float = 1.0,
|
||
presence_penalty: float = 0.0,
|
||
stop: Optional[list[str]] = None,
|
||
cancel_event: Optional[threading.Event] = None,
|
||
enable_thinking: Optional[bool] = None,
|
||
reasoning_effort: Optional[str] = None,
|
||
preserve_thinking: Optional[bool] = None,
|
||
max_tool_iterations: int = 25,
|
||
auto_heal_tool_calls: bool = True,
|
||
tool_call_timeout: int = 300,
|
||
session_id: Optional[str] = None,
|
||
) -> Generator[dict, None, None]:
|
||
"""
|
||
Agentic loop: let the model call tools, execute them, and continue.
|
||
|
||
Yields dicts with:
|
||
{"type": "status", "text": "Searching: ..."/"Reading: ..."} -- tool status updates
|
||
{"type": "content", "text": "token"} -- streamed content tokens (cumulative)
|
||
{"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative)
|
||
"""
|
||
from core.inference.tools import execute_tool
|
||
|
||
if not self.is_loaded:
|
||
raise RuntimeError("llama-server is not loaded")
|
||
|
||
conversation = list(messages)
|
||
url = f"{self.base_url}/v1/chat/completions"
|
||
_accumulated_completion_tokens = 0
|
||
_accumulated_predicted_ms = 0.0
|
||
_accumulated_predicted_n = 0
|
||
|
||
def _strip_tool_markup(text: str, *, final: bool = False) -> str:
|
||
if not auto_heal_tool_calls:
|
||
return text
|
||
return strip_tool_call_markup(text, final = final)
|
||
|
||
# XML prefixes that signal a tool call in content.
|
||
# Empty when auto_heal is disabled so the buffer never
|
||
# speculatively holds content for XML detection.
|
||
_TOOL_XML_SIGNALS = (
|
||
("<tool_call>", "<function=") if auto_heal_tool_calls else ()
|
||
)
|
||
_MAX_BUFFER_CHARS = 32
|
||
|
||
# ── Duplicate tool-call detection ────────────────────────
|
||
# Track recent (tool_name, arguments) hashes to detect loops
|
||
# where the model repeats the exact same call. Retries after
|
||
# a transient failure are allowed (only block when the previous
|
||
# identical call succeeded).
|
||
_tool_call_history: list[tuple[str, bool]] = [] # (key, failed)
|
||
|
||
# ── Re-prompt on plan-without-action ─────────────────
|
||
# When the model describes what it intends to do (forward-looking
|
||
# language) without actually calling a tool, re-prompt once.
|
||
# Only triggers on responses that signal intent/planning -- a
|
||
# direct answer like "4" or "Hello!" will not match.
|
||
# Pattern is compiled once at module level (_INTENT_SIGNAL).
|
||
_reprompt_count = 0
|
||
|
||
# Reserve extra iterations for re-prompts so they don't
|
||
# consume the caller's tool-call budget. Only add the
|
||
# extra slot when tool iterations are actually allowed.
|
||
_extra = _MAX_REPROMPTS if max_tool_iterations > 0 else 0
|
||
for iteration in range(max_tool_iterations + _extra):
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
|
||
# Build payload -- stream: True so we detect tool signals
|
||
# in the first 1-2 chunks without a non-streaming penalty.
|
||
payload = {
|
||
"messages": conversation,
|
||
"stream": True,
|
||
"stream_options": {"include_usage": True},
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
"tools": tools,
|
||
"tool_choice": "auto",
|
||
}
|
||
_reasoning_kw = self._request_reasoning_kwargs(
|
||
enable_thinking, reasoning_effort, preserve_thinking
|
||
)
|
||
if _reasoning_kw is not None:
|
||
payload["chat_template_kwargs"] = _reasoning_kw
|
||
payload["max_tokens"] = (
|
||
max_tokens
|
||
if max_tokens is not None
|
||
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
||
)
|
||
payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
|
||
if stop:
|
||
payload["stop"] = stop
|
||
|
||
try:
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"}
|
||
if self._api_key
|
||
else None
|
||
)
|
||
|
||
# ── Speculative buffer state machine ──────────────────
|
||
# BUFFERING: accumulating content, checking for tool signals
|
||
# STREAMING: no tool detected, yielding tokens to caller
|
||
# DRAINING: tool signal found, silently consuming rest
|
||
_S_BUFFERING = 0
|
||
_S_STREAMING = 1
|
||
_S_DRAINING = 2
|
||
|
||
detect_state = _S_BUFFERING
|
||
content_buffer = "" # Raw content held during BUFFERING
|
||
content_accum = "" # All content tokens (for tool parsing)
|
||
reasoning_accum = ""
|
||
cumulative_display = "" # Cumulative text yielded (with <think>)
|
||
in_thinking = False
|
||
has_content_tokens = False
|
||
tool_calls_acc = {} # Structured delta.tool_calls fragments
|
||
has_structured_tc = False
|
||
_iter_usage = None
|
||
_iter_timings = None
|
||
_stream_done = False
|
||
_last_emitted = ""
|
||
|
||
stream_timeout = httpx.Timeout(
|
||
connect = 10,
|
||
read = 0.5,
|
||
write = 10,
|
||
pool = 10,
|
||
)
|
||
with httpx.Client(
|
||
timeout = stream_timeout,
|
||
limits = httpx.Limits(max_keepalive_connections = 0),
|
||
) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: "
|
||
f"{error_body}"
|
||
)
|
||
|
||
raw_buf = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response,
|
||
cancel_event,
|
||
):
|
||
raw_buf += raw_chunk
|
||
while "\n" in raw_buf:
|
||
line, raw_buf = raw_buf.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
# Flush thinking state for STREAMING
|
||
if detect_state == _S_STREAMING and in_thinking:
|
||
if has_content_tokens:
|
||
cumulative_display += "</think>"
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative_display,
|
||
final = True,
|
||
),
|
||
}
|
||
else:
|
||
cumulative_display = reasoning_accum
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
chunk_data = json.loads(line[6:])
|
||
_ct = chunk_data.get("timings")
|
||
if _ct:
|
||
_iter_timings = _ct
|
||
_cu = chunk_data.get("usage")
|
||
if _cu:
|
||
_iter_usage = _cu
|
||
|
||
choices = chunk_data.get("choices", [])
|
||
if not choices:
|
||
continue
|
||
|
||
delta = choices[0].get("delta", {})
|
||
|
||
# ── Structured tool_calls ──
|
||
tc_deltas = delta.get("tool_calls")
|
||
if tc_deltas:
|
||
# Once visible content has been
|
||
# emitted, do not reclassify this
|
||
# turn as a tool call.
|
||
if _last_emitted:
|
||
continue
|
||
has_structured_tc = True
|
||
detect_state = _S_DRAINING
|
||
for tc_d in tc_deltas:
|
||
idx = tc_d.get("index", 0)
|
||
if idx not in tool_calls_acc:
|
||
tool_calls_acc[idx] = {
|
||
"id": tc_d.get("id", f"call_{idx}"),
|
||
"type": "function",
|
||
"function": {
|
||
"name": "",
|
||
"arguments": "",
|
||
},
|
||
}
|
||
elif tc_d.get("id"):
|
||
# Update ID if real one
|
||
# arrives on a later delta
|
||
tool_calls_acc[idx]["id"] = tc_d["id"]
|
||
func = tc_d.get("function", {})
|
||
if func.get("name"):
|
||
tool_calls_acc[idx]["function"][
|
||
"name"
|
||
] += func["name"]
|
||
if func.get("arguments"):
|
||
tool_calls_acc[idx]["function"][
|
||
"arguments"
|
||
] += func["arguments"]
|
||
continue
|
||
|
||
# ── Reasoning tokens ──
|
||
# Only yield in STREAMING state. In BUFFERING
|
||
# and DRAINING, accumulate silently so we don't
|
||
# corrupt the consumer's prev_text tracker
|
||
# (routes/inference.py never resets prev_text
|
||
# between tool iterations).
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_accum += reasoning
|
||
if detect_state == _S_STREAMING:
|
||
if not in_thinking:
|
||
cumulative_display += "<think>"
|
||
in_thinking = True
|
||
cumulative_display += reasoning
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
|
||
# ── Content tokens ──
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
content_accum += token
|
||
|
||
if detect_state == _S_DRAINING:
|
||
pass # accumulate silently
|
||
|
||
elif detect_state == _S_STREAMING:
|
||
if in_thinking:
|
||
cumulative_display += "</think>"
|
||
in_thinking = False
|
||
cumulative_display += token
|
||
cleaned = _strip_tool_markup(
|
||
cumulative_display,
|
||
)
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {
|
||
"type": "content",
|
||
"text": cleaned,
|
||
}
|
||
|
||
elif detect_state == _S_BUFFERING:
|
||
content_buffer += token
|
||
stripped_buf = content_buffer.lstrip()
|
||
if not stripped_buf:
|
||
continue
|
||
|
||
# Check tool signal prefixes
|
||
is_prefix = False
|
||
is_match = False
|
||
for sig in _TOOL_XML_SIGNALS:
|
||
if stripped_buf.startswith(sig):
|
||
is_match = True
|
||
break
|
||
if sig.startswith(stripped_buf):
|
||
is_prefix = True
|
||
break
|
||
|
||
if is_match:
|
||
detect_state = _S_DRAINING
|
||
elif (
|
||
is_prefix
|
||
and len(stripped_buf)
|
||
< _MAX_BUFFER_CHARS
|
||
):
|
||
pass # keep buffering
|
||
else:
|
||
# Not a tool -- flush buffer
|
||
detect_state = _S_STREAMING
|
||
# Flush any reasoning accumulated
|
||
# during BUFFERING phase
|
||
if reasoning_accum:
|
||
cumulative_display += "<think>"
|
||
cumulative_display += (
|
||
reasoning_accum
|
||
)
|
||
cumulative_display += "</think>"
|
||
cumulative_display += content_buffer
|
||
cleaned = _strip_tool_markup(
|
||
cumulative_display,
|
||
)
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {
|
||
"type": "content",
|
||
"text": cleaned,
|
||
}
|
||
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
|
||
# ── Resolve BUFFERING at stream end ──
|
||
if detect_state == _S_BUFFERING:
|
||
stripped_buf = content_buffer.lstrip()
|
||
if (
|
||
stripped_buf
|
||
and auto_heal_tool_calls
|
||
and any(s in stripped_buf for s in _TOOL_XML_SIGNALS)
|
||
):
|
||
detect_state = _S_DRAINING
|
||
elif content_accum or reasoning_accum:
|
||
detect_state = _S_STREAMING
|
||
if content_buffer:
|
||
# Flush any reasoning accumulated first
|
||
if reasoning_accum:
|
||
cumulative_display += "<think>"
|
||
cumulative_display += reasoning_accum
|
||
cumulative_display += "</think>"
|
||
cumulative_display += content_buffer
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative_display,
|
||
final = True,
|
||
),
|
||
}
|
||
elif reasoning_accum and not has_content_tokens:
|
||
# Reasoning-only response (no content tokens):
|
||
# show reasoning as plain text, matching
|
||
# the final streaming pass behavior for
|
||
# models that put everything in reasoning.
|
||
cumulative_display = reasoning_accum
|
||
yield {
|
||
"type": "content",
|
||
"text": cumulative_display,
|
||
}
|
||
else:
|
||
return
|
||
|
||
# ── STREAMING path: no tool call ──
|
||
if detect_state == _S_STREAMING:
|
||
# Safety net: check for XML tool signals in content.
|
||
# The route layer resets prev_text on tool_start, so
|
||
# post-tool synthesis streams correctly even if
|
||
# content was already emitted before the tool XML.
|
||
_safety_tc = None
|
||
if auto_heal_tool_calls and any(
|
||
s in content_accum for s in _TOOL_XML_SIGNALS
|
||
):
|
||
_safety_tc = self._parse_tool_calls_from_text(
|
||
content_accum,
|
||
)
|
||
if not _safety_tc:
|
||
# ── Re-prompt on plan-without-action ──
|
||
# If the model described what it intends to do
|
||
# (forward-looking language) without calling any
|
||
# tool, nudge it to act. Only fires once per
|
||
# request and only on short responses that
|
||
# contain intent signals -- a direct answer
|
||
# like "4" or "Hello!" won't trigger this.
|
||
# Use content if available, otherwise fall back
|
||
# to reasoning text (reasoning-only stalls).
|
||
_stripped = content_accum.strip()
|
||
if not _stripped:
|
||
_stripped = reasoning_accum.strip()
|
||
if (
|
||
tools
|
||
and _reprompt_count < _MAX_REPROMPTS
|
||
and 0 < len(_stripped) < _REPROMPT_MAX_CHARS
|
||
and _INTENT_SIGNAL.search(_stripped)
|
||
):
|
||
_reprompt_count += 1
|
||
logger.info(
|
||
f"Re-prompt {_reprompt_count}/{_MAX_REPROMPTS}: "
|
||
f"model responded without calling tools "
|
||
f"({len(_stripped)} chars)"
|
||
)
|
||
conversation.append(
|
||
{
|
||
"role": "assistant",
|
||
"content": _stripped,
|
||
}
|
||
)
|
||
conversation.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"STOP. Do NOT write code or explain. "
|
||
"You MUST call a tool NOW. "
|
||
"Call web_search or python immediately."
|
||
),
|
||
}
|
||
)
|
||
# Accumulate tokens and timing from this iteration
|
||
_fu_r = (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings)
|
||
or {}
|
||
)
|
||
_accumulated_completion_tokens += _fu_r.get(
|
||
"completion_tokens", 0
|
||
)
|
||
_it_r = _iter_timings or {}
|
||
_accumulated_predicted_ms += _it_r.get("predicted_ms", 0)
|
||
_accumulated_predicted_n += _it_r.get("predicted_n", 0)
|
||
yield {"type": "status", "text": ""}
|
||
continue
|
||
|
||
# Content was already streamed. Yield metadata.
|
||
yield {"type": "status", "text": ""}
|
||
_fu = (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings)
|
||
or {}
|
||
)
|
||
_fc = _fu.get("completion_tokens", 0)
|
||
_fp = _fu.get("prompt_tokens", 0)
|
||
_tc = _fc + _accumulated_completion_tokens
|
||
if (
|
||
_iter_usage
|
||
or _iter_timings
|
||
or _accumulated_completion_tokens
|
||
):
|
||
_mt = dict(_iter_timings) if _iter_timings else {}
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_mt["predicted_ms"] = (
|
||
_mt.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_tn = (
|
||
_mt.get("predicted_n", 0) + _accumulated_predicted_n
|
||
)
|
||
_mt["predicted_n"] = _tn
|
||
_tms = _mt["predicted_ms"]
|
||
if _tms > 0:
|
||
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _fp,
|
||
"completion_tokens": _tc,
|
||
"total_tokens": _fp + _tc,
|
||
},
|
||
"timings": _mt,
|
||
}
|
||
return
|
||
|
||
# Safety net caught tool XML -- treat as tool call
|
||
tool_calls = _safety_tc
|
||
content_text = _strip_tool_markup(
|
||
content_accum,
|
||
final = True,
|
||
)
|
||
logger.info(
|
||
f"Safety net: parsed {len(tool_calls)} tool call(s) "
|
||
f"from streamed content"
|
||
)
|
||
else:
|
||
# ── DRAINING path: assemble tool_calls ──
|
||
tool_calls = None
|
||
content_text = content_accum
|
||
if has_structured_tc:
|
||
# Filter out incomplete fragments (e.g. from
|
||
# truncation by max_tokens or disconnect).
|
||
tool_calls = [
|
||
tool_calls_acc[i]
|
||
for i in sorted(tool_calls_acc)
|
||
if (
|
||
tool_calls_acc[i]
|
||
.get("function", {})
|
||
.get("name", "")
|
||
.strip()
|
||
)
|
||
] or None
|
||
if (
|
||
not tool_calls
|
||
and auto_heal_tool_calls
|
||
and any(s in content_accum for s in _TOOL_XML_SIGNALS)
|
||
):
|
||
tool_calls = self._parse_tool_calls_from_text(
|
||
content_accum,
|
||
)
|
||
if tool_calls and not has_structured_tc:
|
||
content_text = _strip_tool_markup(
|
||
content_text,
|
||
final = True,
|
||
)
|
||
if tool_calls:
|
||
logger.info(
|
||
f"Parsed {len(tool_calls)} tool call(s) from "
|
||
f"{'structured delta' if has_structured_tc else 'content text'}"
|
||
)
|
||
if not tool_calls:
|
||
# DRAINING but no tool calls (false positive).
|
||
# Merge accumulated metrics from prior tool
|
||
# iterations so they are not silently dropped.
|
||
yield {"type": "status", "text": ""}
|
||
if content_accum:
|
||
# Strip leaked tool-call XML before yielding
|
||
content_accum = _strip_tool_markup(
|
||
content_accum, final = True
|
||
)
|
||
if content_accum:
|
||
yield {"type": "content", "text": content_accum}
|
||
_fu = (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings)
|
||
or {}
|
||
)
|
||
_fc = _fu.get("completion_tokens", 0)
|
||
_fp = _fu.get("prompt_tokens", 0)
|
||
_tc = _fc + _accumulated_completion_tokens
|
||
if (
|
||
_iter_usage
|
||
or _iter_timings
|
||
or _accumulated_completion_tokens
|
||
):
|
||
_mt = dict(_iter_timings) if _iter_timings else {}
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_mt["predicted_ms"] = (
|
||
_mt.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_tn = (
|
||
_mt.get("predicted_n", 0) + _accumulated_predicted_n
|
||
)
|
||
_mt["predicted_n"] = _tn
|
||
_tms = _mt["predicted_ms"]
|
||
if _tms > 0:
|
||
_mt["predicted_per_second"] = _tn / (_tms / 1000.0)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _fp,
|
||
"completion_tokens": _tc,
|
||
"total_tokens": _fp + _tc,
|
||
},
|
||
"timings": _mt,
|
||
}
|
||
return
|
||
|
||
# ── Execute tool calls ──
|
||
_accumulated_completion_tokens += (
|
||
_backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
|
||
).get("completion_tokens", 0)
|
||
_it = _iter_timings or {}
|
||
_accumulated_predicted_ms += _it.get("predicted_ms", 0)
|
||
_accumulated_predicted_n += _it.get("predicted_n", 0)
|
||
|
||
assistant_msg = {"role": "assistant", "content": content_text}
|
||
if tool_calls:
|
||
assistant_msg["tool_calls"] = tool_calls
|
||
conversation.append(assistant_msg)
|
||
|
||
for tc in tool_calls or []:
|
||
func = tc.get("function", {})
|
||
tool_name = func.get("name", "")
|
||
raw_args = func.get("arguments", {})
|
||
|
||
if isinstance(raw_args, str):
|
||
try:
|
||
arguments = json.loads(raw_args)
|
||
except (json.JSONDecodeError, ValueError):
|
||
if auto_heal_tool_calls:
|
||
arguments = {"query": raw_args}
|
||
else:
|
||
arguments = {"raw": raw_args}
|
||
else:
|
||
arguments = raw_args
|
||
|
||
if tool_name == "web_search":
|
||
_ws_url = (arguments.get("url") or "").strip()
|
||
if _ws_url:
|
||
_parsed = urlparse(_ws_url)
|
||
if _parsed.scheme in ("http", "https") and _parsed.hostname:
|
||
_ws_host = _parsed.hostname
|
||
if _ws_host.startswith("www."):
|
||
_ws_host = _ws_host[4:]
|
||
status_text = f"Reading: {_ws_host}"
|
||
else:
|
||
status_text = "Reading page..."
|
||
else:
|
||
status_text = f"Searching: {arguments.get('query', '')}"
|
||
elif tool_name == "python":
|
||
preview = (
|
||
(arguments.get("code") or "").strip().split("\n")[0][:60]
|
||
)
|
||
status_text = (
|
||
f"Running Python: {preview}"
|
||
if preview
|
||
else "Running Python..."
|
||
)
|
||
elif tool_name == "terminal":
|
||
cmd_preview = (arguments.get("command") or "")[:60]
|
||
status_text = (
|
||
f"Running: {cmd_preview}"
|
||
if cmd_preview
|
||
else "Running command..."
|
||
)
|
||
else:
|
||
status_text = f"Calling: {tool_name}"
|
||
yield {"type": "status", "text": status_text}
|
||
|
||
yield {
|
||
"type": "tool_start",
|
||
"tool_name": tool_name,
|
||
"tool_call_id": tc.get("id", ""),
|
||
"arguments": arguments,
|
||
}
|
||
|
||
# ── Duplicate call detection ──────────────
|
||
# str(dict) is stable here: arguments always comes from
|
||
# json.loads on the same model output within one request,
|
||
# so insertion order is deterministic (Python 3.7+).
|
||
_tc_key = tool_name + str(arguments)
|
||
_prev = _tool_call_history[-1] if _tool_call_history else None
|
||
if _prev and _prev[0] == _tc_key and not _prev[1]:
|
||
result = (
|
||
"You already made this exact call. "
|
||
"Do not repeat the same tool call. "
|
||
"Try a different approach: fetch a URL "
|
||
"from previous results, use Python to "
|
||
"process data you already have, or "
|
||
"provide your final answer now."
|
||
)
|
||
else:
|
||
_effective_timeout = (
|
||
None if tool_call_timeout >= 9999 else tool_call_timeout
|
||
)
|
||
# Guard against the model emitting a tool not in the
|
||
# per-request advertised set: filtered MCP names, a
|
||
# built-in the caller opted out of, or a stale name
|
||
# from a prior turn. Mirrors the safetensors loop's
|
||
# allowed_tool_names check.
|
||
_allowed = {
|
||
(t.get("function") or {}).get("name")
|
||
for t in (tools or [])
|
||
if (t.get("function") or {}).get("name")
|
||
}
|
||
if _allowed and tool_name not in _allowed:
|
||
result = (
|
||
f"Error: tool '{tool_name}' is not enabled "
|
||
"for this request. Use one of the enabled "
|
||
"tools or provide a final answer."
|
||
)
|
||
else:
|
||
result = execute_tool(
|
||
tool_name,
|
||
arguments,
|
||
cancel_event = cancel_event,
|
||
timeout = _effective_timeout,
|
||
session_id = session_id,
|
||
)
|
||
|
||
yield {
|
||
"type": "tool_end",
|
||
"tool_name": tool_name,
|
||
"tool_call_id": tc.get("id", ""),
|
||
"result": result,
|
||
}
|
||
|
||
# Nudge model to try a different approach on errors
|
||
_error_prefixes = (
|
||
"Error",
|
||
"Search failed",
|
||
"Execution error",
|
||
"Blocked:",
|
||
"Exit code",
|
||
"Failed to fetch",
|
||
"Failed to resolve",
|
||
"No query provided",
|
||
)
|
||
_is_error = isinstance(result, str) and result.lstrip().startswith(
|
||
_error_prefixes
|
||
)
|
||
_tool_call_history.append((_tc_key, _is_error))
|
||
# Strip image sentinel before feeding result to the LLM
|
||
# (the full result with sentinel is still yielded via
|
||
# tool_end so the frontend can extract image paths).
|
||
_result_content = result
|
||
if "\n__IMAGES__:" in _result_content:
|
||
_result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0]
|
||
if _is_error:
|
||
_result_content = (
|
||
_result_content + "\n\nThe tool call encountered an issue. "
|
||
"Please try a different approach or rephrase your request."
|
||
)
|
||
|
||
tool_msg = {
|
||
"role": "tool",
|
||
"name": tool_name,
|
||
"content": _result_content,
|
||
}
|
||
tool_call_id = tc.get("id")
|
||
if tool_call_id:
|
||
tool_msg["tool_call_id"] = tool_call_id
|
||
conversation.append(tool_msg)
|
||
|
||
# Clear tool status badge before next generation iteration
|
||
yield {"type": "status", "text": ""}
|
||
# Continue the loop to let model respond with context
|
||
continue
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── Tool iteration cap reached -- synthesize final answer ──
|
||
# The model used all iterations without producing a final text
|
||
# response. Inject a nudge so the final streaming pass produces
|
||
# a useful answer instead of continuing to request tools.
|
||
if max_tool_iterations > 0:
|
||
conversation.append(
|
||
{
|
||
"role": "user",
|
||
"content": (
|
||
"You have used all available tool calls. Based on "
|
||
"everything you have found so far, provide your final "
|
||
"answer now. Do not call any more tools."
|
||
),
|
||
}
|
||
)
|
||
|
||
# Clear status
|
||
yield {"type": "status", "text": ""}
|
||
|
||
# Final streaming pass with the full conversation context
|
||
stream_payload = {
|
||
"messages": conversation,
|
||
"stream": True,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
"presence_penalty": presence_penalty,
|
||
}
|
||
_reasoning_kw = self._request_reasoning_kwargs(
|
||
enable_thinking, reasoning_effort, preserve_thinking
|
||
)
|
||
if _reasoning_kw is not None:
|
||
stream_payload["chat_template_kwargs"] = _reasoning_kw
|
||
stream_payload["max_tokens"] = (
|
||
max_tokens
|
||
if max_tokens is not None
|
||
else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR)
|
||
)
|
||
stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS
|
||
if stop:
|
||
stream_payload["stop"] = stop
|
||
stream_payload["stream_options"] = {"include_usage": True}
|
||
|
||
cumulative = ""
|
||
_last_emitted = ""
|
||
in_thinking = False
|
||
has_content_tokens = False
|
||
reasoning_text = ""
|
||
_metadata_usage = None
|
||
_metadata_timings = None
|
||
_stream_done = False
|
||
|
||
try:
|
||
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(
|
||
timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
|
||
) as client:
|
||
with self._stream_with_retry(
|
||
client,
|
||
url,
|
||
stream_payload,
|
||
cancel_event,
|
||
headers = _auth_headers,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
error_body = response.read().decode()
|
||
raise RuntimeError(
|
||
f"llama-server returned {response.status_code}: {error_body}"
|
||
)
|
||
|
||
buffer = ""
|
||
for raw_chunk in self._iter_text_cancellable(
|
||
response, cancel_event
|
||
):
|
||
buffer += raw_chunk
|
||
while "\n" in buffer:
|
||
line, buffer = buffer.split("\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
if line == "data: [DONE]":
|
||
if in_thinking:
|
||
if has_content_tokens:
|
||
cumulative += "</think>"
|
||
yield {
|
||
"type": "content",
|
||
"text": _strip_tool_markup(
|
||
cumulative, final = True
|
||
),
|
||
}
|
||
else:
|
||
cumulative = reasoning_text
|
||
yield {"type": "content", "text": cumulative}
|
||
_stream_done = True
|
||
break # exit inner while
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
try:
|
||
chunk_data = json.loads(line[6:])
|
||
# Capture server timings/usage from final chunks
|
||
_chunk_timings = chunk_data.get("timings")
|
||
if _chunk_timings:
|
||
_metadata_timings = _chunk_timings
|
||
_chunk_usage = chunk_data.get("usage")
|
||
if _chunk_usage:
|
||
_metadata_usage = _chunk_usage
|
||
choices = chunk_data.get("choices", [])
|
||
if choices:
|
||
delta = choices[0].get("delta", {})
|
||
|
||
reasoning = delta.get("reasoning_content", "")
|
||
if reasoning:
|
||
reasoning_text += reasoning
|
||
if not in_thinking:
|
||
cumulative += "<think>"
|
||
in_thinking = True
|
||
cumulative += reasoning
|
||
yield {"type": "content", "text": cumulative}
|
||
|
||
token = delta.get("content", "")
|
||
if token:
|
||
has_content_tokens = True
|
||
if in_thinking:
|
||
cumulative += "</think>"
|
||
in_thinking = False
|
||
cumulative += token
|
||
cleaned = _strip_tool_markup(cumulative)
|
||
# Only emit when cleaned text grows (monotonic).
|
||
if len(cleaned) > len(_last_emitted):
|
||
_last_emitted = cleaned
|
||
yield {"type": "content", "text": cleaned}
|
||
except json.JSONDecodeError:
|
||
logger.debug(
|
||
f"Skipping malformed SSE line: {line[:100]}"
|
||
)
|
||
if _stream_done:
|
||
break # exit outer for
|
||
_final_usage = _metadata_usage or {}
|
||
_final_completion = _final_usage.get("completion_tokens", 0)
|
||
_final_prompt = _final_usage.get("prompt_tokens", 0)
|
||
_total_completion = (
|
||
_final_completion + _accumulated_completion_tokens
|
||
)
|
||
if _metadata_usage or _metadata_timings:
|
||
_merged_timings = (
|
||
dict(_metadata_timings) if _metadata_timings else {}
|
||
)
|
||
if _accumulated_predicted_ms or _accumulated_predicted_n:
|
||
_merged_timings["predicted_ms"] = (
|
||
_merged_timings.get("predicted_ms", 0)
|
||
+ _accumulated_predicted_ms
|
||
)
|
||
_total_predicted_n = (
|
||
_merged_timings.get("predicted_n", 0)
|
||
+ _accumulated_predicted_n
|
||
)
|
||
_merged_timings["predicted_n"] = _total_predicted_n
|
||
_total_predicted_ms = _merged_timings["predicted_ms"]
|
||
if _total_predicted_ms > 0:
|
||
_merged_timings["predicted_per_second"] = (
|
||
_total_predicted_n / (_total_predicted_ms / 1000.0)
|
||
)
|
||
yield {
|
||
"type": "metadata",
|
||
"usage": {
|
||
"prompt_tokens": _final_prompt,
|
||
"completion_tokens": _total_completion,
|
||
"total_tokens": _final_prompt + _total_completion,
|
||
},
|
||
"timings": _merged_timings,
|
||
}
|
||
|
||
except httpx.ConnectError:
|
||
raise RuntimeError("Lost connection to llama-server")
|
||
except Exception as e:
|
||
if cancel_event is not None and cancel_event.is_set():
|
||
return
|
||
raise
|
||
|
||
# ── TTS support ────────────────────────────────────────────
|
||
|
||
def detect_audio_type(self) -> Optional[str]:
|
||
"""Detect audio/TTS codec; swallows errors (use _strict variant to distinguish)."""
|
||
try:
|
||
return self._detect_audio_type_strict()
|
||
except Exception as e:
|
||
logger.debug(f"Audio type detection failed: {e}")
|
||
return None
|
||
|
||
def _detect_audio_type_strict(self) -> Optional[str]:
|
||
"""Codec name on match, None on definitive non-audio, raises on transport/JSON errors."""
|
||
if not self.is_loaded:
|
||
return None
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(timeout = 10, headers = _auth_headers) as client:
|
||
|
||
def _detok(tid: int) -> str:
|
||
# Non-200 means "marker not in vocab" -- keep probing.
|
||
# Transport / JSON errors still raise.
|
||
r = client.post(f"{self.base_url}/detokenize", json = {"tokens": [tid]})
|
||
if r.status_code != 200:
|
||
return ""
|
||
return r.json().get("content", "")
|
||
|
||
def _tok(text: str) -> list[int]:
|
||
r = client.post(
|
||
f"{self.base_url}/tokenize",
|
||
json = {"content": text, "add_special": False},
|
||
)
|
||
if r.status_code != 200:
|
||
return []
|
||
return r.json().get("tokens", [])
|
||
|
||
# Check codec-specific tokens (not generic ones that may exist in non-audio models)
|
||
if "<custom_token_" in _detok(128258) and "<custom_token_" in _detok(
|
||
128259
|
||
):
|
||
return "snac"
|
||
if len(_tok("<|AUDIO|>")) == 1 and len(_tok("<|audio_eos|>")) == 1:
|
||
return "csm"
|
||
if len(_tok("<|startoftranscript|>")) == 1:
|
||
return "whisper"
|
||
if len(_tok("<audio_soft_token>")) == 1:
|
||
return "audio_vlm"
|
||
if (
|
||
len(_tok("<|bicodec_semantic_0|>")) == 1
|
||
and len(_tok("<|bicodec_global_0|>")) == 1
|
||
):
|
||
return "bicodec"
|
||
if len(_tok("<|c1_0|>")) == 1 and len(_tok("<|c2_0|>")) == 1:
|
||
return "dac"
|
||
return None
|
||
|
||
# Prompt format per codec: (template, stop_tokens, needs_token_ids)
|
||
# Matches prompts in InferenceBackend._generate_snac/bicodec/dac
|
||
_TTS_PROMPTS = {
|
||
"snac": (
|
||
"<custom_token_3>{text}<|eot_id|><custom_token_4>",
|
||
["<custom_token_2>"],
|
||
True,
|
||
),
|
||
"bicodec": (
|
||
"<|task_tts|><|start_content|>{text}<|end_content|><|start_global_token|>",
|
||
["<|im_end|>", "</s>"],
|
||
False,
|
||
),
|
||
"dac": (
|
||
"<|im_start|>\n<|text_start|>{text}<|text_end|>\n<|audio_start|><|global_features_start|>\n",
|
||
["<|im_end|>", "<|audio_end|>"],
|
||
False,
|
||
),
|
||
}
|
||
|
||
_codec_mgr = None # Shared AudioCodecManager instance
|
||
|
||
def init_audio_codec(self, audio_type: str) -> None:
|
||
"""Load the audio codec at model load time (mirrors non-GGUF path)."""
|
||
import torch
|
||
from core.inference.audio_codecs import AudioCodecManager
|
||
|
||
if LlamaCppBackend._codec_mgr is None:
|
||
LlamaCppBackend._codec_mgr = AudioCodecManager()
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
model_repo_path = None
|
||
|
||
# BiCodec needs a repo with BiCodec/ weights — download canonical SparkTTS
|
||
if audio_type == "bicodec":
|
||
from huggingface_hub import snapshot_download
|
||
import os
|
||
|
||
repo_path = snapshot_download(
|
||
"unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B"
|
||
)
|
||
model_repo_path = os.path.abspath(repo_path)
|
||
|
||
LlamaCppBackend._codec_mgr.load_codec(
|
||
audio_type, device, model_repo_path = model_repo_path
|
||
)
|
||
logger.info(f"Loaded audio codec for GGUF TTS: {audio_type}")
|
||
|
||
def generate_audio_response(
|
||
self,
|
||
text: str,
|
||
audio_type: str,
|
||
temperature: float = 0.6,
|
||
top_p: float = 0.95,
|
||
top_k: int = 50,
|
||
min_p: float = 0.0,
|
||
max_new_tokens: int = 2048,
|
||
repetition_penalty: float = 1.1,
|
||
) -> tuple:
|
||
"""
|
||
Generate TTS audio via llama-server /completion + codec decoding.
|
||
Returns (wav_bytes, sample_rate).
|
||
"""
|
||
if audio_type not in self._TTS_PROMPTS:
|
||
raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.")
|
||
|
||
tpl, stop, need_ids = self._TTS_PROMPTS[audio_type]
|
||
|
||
payload: dict = {
|
||
"prompt": tpl.format(text = text),
|
||
"stream": False,
|
||
"n_predict": max_new_tokens,
|
||
"temperature": temperature,
|
||
"top_p": top_p,
|
||
"top_k": top_k if top_k >= 0 else 0,
|
||
"min_p": min_p,
|
||
"repeat_penalty": repetition_penalty,
|
||
}
|
||
if stop:
|
||
payload["stop"] = stop
|
||
if need_ids:
|
||
payload["n_probs"] = 1
|
||
|
||
_auth_headers = (
|
||
{"Authorization": f"Bearer {self._api_key}"} if self._api_key else None
|
||
)
|
||
with httpx.Client(
|
||
timeout = httpx.Timeout(300, connect = 10), headers = _auth_headers
|
||
) as client:
|
||
resp = client.post(f"{self.base_url}/completion", json = payload)
|
||
if resp.status_code != 200:
|
||
raise RuntimeError(
|
||
f"llama-server returned {resp.status_code}: {resp.text}"
|
||
)
|
||
|
||
data = resp.json()
|
||
token_ids = (
|
||
[p["id"] for p in data.get("completion_probabilities", []) if "id" in p]
|
||
if need_ids
|
||
else None
|
||
)
|
||
|
||
import torch
|
||
|
||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
return LlamaCppBackend._codec_mgr.decode(
|
||
audio_type, device, token_ids = token_ids, text = data.get("content", "")
|
||
)
|