From c31a6e876d66c0dcc9cd5c26b10e9784e589fae7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 9 Jun 2026 23:12:47 -0700 Subject: [PATCH] fix(install,studio): address Codex round-4 review (6 of 7 comments real) - install.sh: gate the new aarch64 bitsandbytes block on SKIP_TORCH=false -- with --no-torch/UNSLOTH_NO_TORCH (GGUF-only install) it would have pulled torch back into the venv through bitsandbytes' dependencies. - studio worker: in the new Spark OOM-guard section, decide PYTORCH_CUDA_ALLOC_CONF (expandable_segments) BEFORE the guard's first CUDA touch -- get_device_properties initializes the CUDA allocator, after which the env var is ignored, and the later `import unsloth` (patch_dgx_spark_memory_config) is too late for the worker process. Uses the same CUDA-free nvidia-smi name sniff, append-don't-override, and UNSLOTH_NO_EXPANDABLE_SEGMENTS opt-out as the library patch. Live-verified on the N1X: env set while torch.cuda.is_initialized() is still False. - uninstall.ps1: only run `fuser -k 8888/tcp` in a probed WSL distro when an Unsloth install actually exists there (checked BEFORE the rm deletes the marker) -- an unrelated listener on 8888 (e.g. Jupyter) in a clean distro must survive a Windows-side uninstall. The Unsloth-specific pkills stay unconditional. - install.ps1 + uninstall.ps1: persist the chosen WSL distro to %LOCALAPPDATA%\Unsloth\wsl-distro.txt at install; uninstall reads it (before removing the directory) and prepends it to the cleanup candidates, so a custom UNSLOTH_WSL_DISTRO install is cleaned without the env var being set again at uninstall time. - provision_llama_cuda.sh: honor UNSLOTH_LLAMA_PR (numeric-validated, best-effort fetch of pull/N/head after clone) so a provisioned tree matches a PR pin the way setup.sh does; and require only llama-server in the main cmake build (mirroring setup.sh), building the helper targets (llama-cli/quantize/mtmd-cli/gguf-split) best-effort afterwards -- an older UNSLOTH_LLAMA_TAG pin lacking a newer helper target no longer fails the whole provision. Not changed: the "--tauri rejection doesn't restore the venv rollback" comment is incorrect -- the rejection returns through Exit-InstallFailure, which itself calls Restore-StudioVenvRollback. Co-Authored-By: Claude Opus 4.8 --- install.ps1 | 3 +++ install.sh | 5 ++++- scripts/uninstall.ps1 | 20 +++++++++++++++++- studio/backend/core/training/worker.py | 28 ++++++++++++++++++++++++++ studio/scripts/provision_llama_cuda.sh | 27 +++++++++++++++++++++++-- 5 files changed, 79 insertions(+), 4 deletions(-) diff --git a/install.ps1 b/install.ps1 index f8e0041496..8ae84f97bb 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1703,6 +1703,9 @@ shell.Run cmd, 0, False "wsl.exe -d $_distroArg -u root -- /root/.unsloth/studio/unsloth_studio/bin/unsloth %*" ) Set-Content -LiteralPath (Join-Path $shimDir "unsloth.cmd") -Value $shimLines -Encoding ASCII + # Record the distro for the uninstaller: a custom UNSLOTH_WSL_DISTRO install + # must be cleanable without the env var being set again at uninstall time. + try { Set-Content -LiteralPath (Join-Path (Split-Path $shimDir -Parent) "wsl-distro.txt") -Value $distro -Encoding ASCII } catch {} # A fresh Windows profile may have no HKCU 'Path' value at all -> $userPath is null # and $userPath.TrimEnd() would throw, losing the shim. Treat null as empty. $userPath = [Environment]::GetEnvironmentVariable("Path", "User") diff --git a/install.sh b/install.sh index 09e637aede..7595255a39 100755 --- a/install.sh +++ b/install.sh @@ -2379,7 +2379,10 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # out of the box. bitsandbytes ships working aarch64 manylinux wheels # (verified on sm_121 Blackwell via PTX JIT), so add it best-effort -- a # platform without a wheel just keeps 16-bit LoRA / full finetuning. - if { [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; } \ + # Gated on SKIP_TORCH: a --no-torch/UNSLOTH_NO_TORCH (GGUF-only) install must + # not have bitsandbytes drag torch back into the venv via its dependencies. + if [ "$SKIP_TORCH" = false ] \ + && { [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; } \ && command -v nvidia-smi >/dev/null 2>&1 \ && nvidia-smi -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}' \ && ! "$_VENV_PY" -c "import bitsandbytes" >/dev/null 2>&1; then diff --git a/scripts/uninstall.ps1 b/scripts/uninstall.ps1 index 46f9a87a77..2cdf01e5d8 100644 --- a/scripts/uninstall.ps1 +++ b/scripts/uninstall.ps1 @@ -344,6 +344,19 @@ function Uninstall-UnslothStudio { # %LOCALAPPDATA%\Unsloth (not "Unsloth Studio") with a PATH entry -- all missed by the cleanup above. _Step "Removing WSL-fallback artifacts (shim, launcher, PATH entry, WSL install)..." $unslothDir = if ($env:LOCALAPPDATA) { Join-Path $env:LOCALAPPDATA "Unsloth" } else { $null } + # The installer records its WSL distro in wsl-distro.txt so a custom + # UNSLOTH_WSL_DISTRO install is cleanable without the env var being set again + # at uninstall time. Read it BEFORE the directory is removed below. + $_recordedDistro = $null + if ($unslothDir) { + try { + $_distroFile = Join-Path $unslothDir "wsl-distro.txt" + if (Test-Path -LiteralPath $_distroFile) { + $_recordedDistro = (Get-Content -LiteralPath $_distroFile -ErrorAction SilentlyContinue | Select-Object -First 1) + if ($_recordedDistro) { $_recordedDistro = $_recordedDistro.Trim() } + } + } catch { } + } if ($unslothDir) { $shimDir = (Join-Path $unslothDir "bin").TrimEnd('\', '/') try { @@ -379,8 +392,13 @@ function Uninstall-UnslothStudio { # touching /home/*/.unsloth would erase an unrelated WSL user's own Unsloth/cache that this # installer never created. pkill patterns use the [x]-regex self-exclusion trick: '[u]nsloth_studio' # keeps the shell's own argv from matching while real processes still match. Same for '[l]lama-server'. - $_clean = 'rm -rf /root/.unsloth /root/llama-cuda /root/provision_llama_cuda.sh /root/llama_cuda_build.log 2>/dev/null; rm -f /root/.local/bin/unsloth 2>/dev/null; fuser -k 8888/tcp 2>/dev/null; pkill -9 -f ''[u]nsloth_studio'' 2>/dev/null; pkill -9 -f ''[l]lama-server'' 2>/dev/null; true' + # The port-8888 kill is gated on an Unsloth install actually existing in the + # distro (checked BEFORE the rm deletes the marker): a probed distro with an + # unrelated listener on 8888 (Jupyter etc.) must not lose it. The pkills are + # already Unsloth-specific, so they stay unconditional. + $_clean = '_had=0; if [ -d /root/.unsloth ] || [ -L /root/.local/bin/unsloth ]; then _had=1; fi; rm -rf /root/.unsloth /root/llama-cuda /root/provision_llama_cuda.sh /root/llama_cuda_build.log 2>/dev/null; rm -f /root/.local/bin/unsloth 2>/dev/null; if [ $_had -eq 1 ]; then fuser -k 8888/tcp 2>/dev/null; fi; pkill -9 -f ''[u]nsloth_studio'' 2>/dev/null; pkill -9 -f ''[l]lama-server'' 2>/dev/null; true' $_cands = @('', 'Ubuntu', 'Ubuntu-24.04', 'Ubuntu-22.04', 'Debian') + if ($_recordedDistro) { $_cands = @($_recordedDistro) + $_cands } if ($env:UNSLOTH_WSL_DISTRO) { $_cands = @($env:UNSLOTH_WSL_DISTRO) + $_cands } $_done = @{} foreach ($d in $_cands) { diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 0390e6745e..15f9e6ff51 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2236,6 +2236,34 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # training loop surfaces the resulting OutOfMemoryError with remediation. else: try: + # The Spark allocator config must be decided BEFORE this guard's first + # CUDA touch: get_device_properties below initializes the CUDA allocator, + # after which PYTORCH_CUDA_ALLOC_CONF changes are ignored -- and the later + # `import unsloth` (patch_dgx_spark_memory_config) would be too late for + # THIS worker process even though it is in time for a plain + # `import unsloth`. CUDA-free sniff via nvidia-smi device names (mirrors + # _is_dgx_spark_no_cuda_init), with the same append-don't-override and + # UNSLOTH_NO_EXPANDABLE_SEGMENTS opt-out semantics as the library patch. + try: + import platform as _plat + _spark_smi = False + if _plat.machine().lower() in ("aarch64", "arm64"): + _smi = _sp.run( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output = True, text = True, timeout = 5, + ) + _names_u = (_smi.stdout or "").upper() + _spark_smi = any( + t in _names_u for t in ("GB10", "GB110", "JMJWOA", "N1X", "DGX SPARK") + ) + if _spark_smi and os.environ.get("UNSLOTH_NO_EXPANDABLE_SEGMENTS") != "1": + _conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "") + if "expandable_segments" not in _conf: + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = ( + (_conf + "," if _conf else "") + "expandable_segments:True" + ) + except Exception: + pass import torch as _torch_mem if _torch_mem.cuda.is_available(): _props = _torch_mem.cuda.get_device_properties(0) diff --git a/studio/scripts/provision_llama_cuda.sh b/studio/scripts/provision_llama_cuda.sh index a15d69e1d2..d89e192dc4 100644 --- a/studio/scripts/provision_llama_cuda.sh +++ b/studio/scripts/provision_llama_cuda.sh @@ -143,6 +143,20 @@ if [ ! -d "$LLAMA_DIR/.git" ]; then _restore_prev exit 0 fi + # Honor a llama.cpp PR pin (UNSLOTH_LLAMA_PR, the same var setup.sh supports) + # so a provisioned tree matches the user's request instead of silently building + # the default branch. Best-effort: a failed fetch keeps the default branch. + case "${UNSLOTH_LLAMA_PR:-}" in + ''|*[!0-9]*) ;; + *) + if git -C "$LLAMA_DIR" fetch --depth 1 origin "pull/${UNSLOTH_LLAMA_PR}/head:_unsloth_pr_${UNSLOTH_LLAMA_PR}" >/dev/null 2>&1 \ + && git -C "$LLAMA_DIR" checkout "_unsloth_pr_${UNSLOTH_LLAMA_PR}" >/dev/null 2>&1; then + log "checked out llama.cpp PR #${UNSLOTH_LLAMA_PR} (UNSLOTH_LLAMA_PR)" + else + log "could not fetch llama.cpp PR #${UNSLOTH_LLAMA_PR}; building the default branch" + fi + ;; + esac fi cd "$LLAMA_DIR" || { _restore_prev; exit 0; } @@ -190,8 +204,16 @@ _NICE="" command -v nice >/dev/null 2>&1 && _NICE="nice -n 19" command -v ionice >/dev/null 2>&1 && _NICE="$_NICE ionice -c 3" _cmake_build() { - $_NICE cmake --build build -j"$JOBS" --target \ - llama-server llama-cli llama-quantize llama-mtmd-cli llama-gguf-split >/dev/null 2>&1 + # Only llama-server is REQUIRED (mirrors setup.sh's source path): an older + # UNSLOTH_LLAMA_TAG pin may predate newer helper targets (llama-mtmd-cli, + # llama-gguf-split), and those missing must not fail the whole provision. + $_NICE cmake --build build -j"$JOBS" --target llama-server >/dev/null 2>&1 +} +_cmake_build_extras() { + # Helper targets unsloth-zoo's GGUF exporter also uses -- best-effort each. + for _t in llama-cli llama-quantize llama-mtmd-cli llama-gguf-split; do + $_NICE cmake --build build -j"$JOBS" --target "$_t" >/dev/null 2>&1 || true + done } if ! _cmake_build; then # An interrupted build (e.g. a thermal/power shutdown mid-compile, which this @@ -203,6 +225,7 @@ if ! _cmake_build; then _cmake_configure || { log "cmake configure failed"; cd /; _restore_prev; exit 0; } _cmake_build || { log "cmake build failed"; cd /; _restore_prev; exit 0; } fi +_cmake_build_extras if is_cuda_server "$SERVER"; then log "CUDA llama-server ready: $SERVER"