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 <noreply@anthropic.com>
2026-06-09 23:12:47 -07:00 · 2026-06-09 23:12:47 -07:00 · c31a6e876d
commit c31a6e876d
parent ca8c1b434d
5 changed files with 79 additions and 4 deletions
--- 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")
--- 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
--- 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) {
--- 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)
--- 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"