diff --git a/.github/scripts/assert-llama-loads.sh b/.github/scripts/assert-llama-loads.sh new file mode 100755 index 0000000000..c2ffe27469 --- /dev/null +++ b/.github/scripts/assert-llama-loads.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +# +# Assert Studio installed a llama.cpp that loads and runs on THIS macOS. Tests +# the contract that matters (binaries load and their minimum-OS is <= this host) +# instead of the old "did install.sh fall back to a source build?" grep, since a +# source build with a correct deployment target is a valid outcome. +set -uo pipefail + +UNSLOTH_HOME="${STUDIO_HOME:-$HOME/.unsloth}" +LLAMA_DIR="${LLAMA_CPP_DIR:-$UNSLOTH_HOME/llama.cpp}" +BIN_DIR="$LLAMA_DIR/build/bin" + +fail() { + echo "::error::$*" + if [ -f logs/install.log ]; then + echo "---- install.log (llama.cpp lines) ----" + grep -E "llama-prebuilt|llama\.cpp|macos prebuilt|falling back" logs/install.log | tail -80 || true + fi + exit 1 +} + +SERVER="$(find "$LLAMA_DIR" -type f -name 'llama-server' 2>/dev/null | head -1)" +QUANT="$(find "$LLAMA_DIR" -type f -name 'llama-quantize' 2>/dev/null | head -1)" +[ -n "$SERVER" ] || fail "llama-server not found under $LLAMA_DIR after install" +[ -n "$QUANT" ] || fail "llama-quantize not found under $LLAMA_DIR after install" + +HOST_VER="$(sw_vers -productVersion 2>/dev/null || echo '0')" +HOST_MAJOR="${HOST_VER%%.*}" + +# Static minimum-OS check on every Mach-O we ship. vtool ships with the Xcode +# command line tools, which GitHub macOS runners always have; if it is somehow +# missing we skip the static check and rely on the runtime launch below. +if command -v vtool >/dev/null 2>&1; then + while IFS= read -r macho; do + [ -n "$macho" ] || continue + minos="$(vtool -show-build "$macho" 2>/dev/null | awk '/minos/{print $2; exit}')" + [ -n "$minos" ] || continue + min_major="${minos%%.*}" + if [ "$min_major" -gt "$HOST_MAJOR" ] 2>/dev/null; then + fail "$(basename "$macho") is built for macOS $minos but this runner is macOS $HOST_VER (prebuilt is newer than the host)" + fi + done < <(find "$BIN_DIR" -type f \( -name '*.dylib' -o -name 'llama-server' -o -name 'llama-quantize' \) 2>/dev/null) +fi + +# Runtime launch: --version forces dyld to load every linked dylib (including +# libggml-metal.dylib). A missing Metal symbol or too-new binary fails here. +if ! "$SERVER" --version >/tmp/llama-server-version.txt 2>&1; then + echo "---- llama-server --version output ----" + cat /tmp/llama-server-version.txt || true + fail "llama-server failed to launch on macOS $HOST_VER (dyld load / symbol error)" +fi + +echo "llama.cpp load validation passed on macOS $HOST_VER" +echo " server: $SERVER" +sed -n '1,4p' /tmp/llama-server-version.txt 2>/dev/null || true diff --git a/.github/workflows/lint-ci.yml b/.github/workflows/lint-ci.yml index 00e6e357e2..8cb3535075 100644 --- a/.github/workflows/lint-ci.yml +++ b/.github/workflows/lint-ci.yml @@ -79,6 +79,56 @@ jobs: run: | ruff check unsloth unsloth_cli studio tests cli.py unsloth-cli.py + - name: Import-hoist verifier self-test + # scripts/verify_import_hoist.py is a scope-aware (LEGB) AST + # resolver that gates import-hoisting / alias-rename refactors + # against two bugs ruff and pyflakes both miss: + # 1. dangling alias -- `from a import b as _b` hoisted to + # `from a import b` but a leftover `_b` reference now + # resolves to nothing (or to some other module-level `_b`). + # 2. rename clash -- `_b -> b` silently re-points at a + # different object already named `b` in that scope. + # This step runs the tool's 8 negative-control cases so a + # regression in the verifier itself fails before we trust it on + # a diff. Hermetic, stdlib-only, sub-second. Hard gate. + run: | + python scripts/verify_import_hoist.py --self-test + + - name: Import-hoist / alias-rename safety (changed Python files) + # Runs the verifier in compare mode on every in-place-modified + # .py in the PR: parses each file BEFORE (base branch) and AFTER + # (this diff), resolves every name load, and fails on a BLOCKER + # (dangling alias / rename clash / re-pointed import). INFO + # findings (a helper relocated to another file) do not fail. + # + # --diff-filter=M (in-place edits only) is deliberate: that is + # exactly where a hoist refactor lives, and it skips brand-new + # files whose re-export imports would otherwise look "unused". + # + # actions/checkout uses fetch-depth: 1, so the base branch is not + # present locally. Fetch the single base commit with an explicit + # refspec so origin/ is reliably created (a bare + # `git fetch origin ` only updates FETCH_HEAD in some + # configs). Two-dot diff avoids needing a merge-base on a shallow + # clone. + if: github.event_name == 'pull_request' + run: | + git fetch --no-tags --depth=1 origin \ + "${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}" + mapfile -t CHANGED < <( + git diff --name-only --diff-filter=M \ + "origin/${{ github.base_ref }}" HEAD -- '*.py' \ + | grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true + ) + if [ "${#CHANGED[@]}" -eq 0 ]; then + echo "no in-place-modified Python files to check" + exit 0 + fi + printf 'checking %d file(s):\n' "${#CHANGED[@]}" + printf ' %s\n' "${CHANGED[@]}" + python scripts/verify_import_hoist.py \ + --before "origin/${{ github.base_ref }}" --after HEAD "${CHANGED[@]}" + - name: No leftover debugger / pdb / breakpoint calls # Catches the "I'll just stick a breakpoint() here" mistake # before it ships. AST-based so commented-out debugger diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index a1e7b2efa6..33ac3b9bd8 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -72,6 +72,31 @@ concurrency: permissions: contents: read +# ────────────────────────────────────────────────────────────────────── +# Network-resilience knobs, applied to every job/step. These add retries +# and backoff ONLY; they do not relax a single integrity check. cargo +# still resolves against Cargo.lock (--locked), pip still verifies the +# wheels it downloads, npm still enforces package-lock integrity, the +# harden-runner egress allowlists below are unchanged, and every action +# stays SHA-pinned. The advisory-audit run on 2026-05-29 red-failed when +# one crates.io tarball fetch hit "Recv failure: Connection reset by +# peer" (curl 56); cargo's default of 3 retries over an HTTP/2-multiplexed +# connection did not recover. The settings below make that class of +# transient fault self-heal instead of failing the whole run. +env: + # pip: raise the built-in retry count and per-connection timeout. + PIP_RETRIES: "10" + PIP_DEFAULT_TIMEOUT: "60" + # cargo: retry network ops and disable HTTP/2 multiplexing -- the + # documented mitigation for the curl-56 connection resets above. + CARGO_NET_RETRY: "10" + CARGO_HTTP_MULTIPLEXING: "false" + CARGO_NET_GIT_FETCH_WITH_CLI: "true" + # npm: retry registry fetches with capped exponential backoff. + NPM_CONFIG_FETCH_RETRIES: "5" + NPM_CONFIG_FETCH_RETRY_MINTIMEOUT: "2000" + NPM_CONFIG_FETCH_RETRY_MAXTIMEOUT: "60000" + jobs: # ───────────────────────────────────────────────────────────────────── # Combined advisory-DB audit: pip-audit + npm audit + cargo audit @@ -140,7 +165,7 @@ jobs: - uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27 - - uses: swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2.9.1 + - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 with: workspaces: studio/src-tauri -> target @@ -153,8 +178,23 @@ jobs: # crashes with a TOML parse error on that file. # npm audit is bundled with the node toolchain, no install. run: | - python -m pip install --upgrade pip 'pip-audit>=2.7' - cargo install --locked --version '^0.22' cargo-audit + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + retry 5 python -m pip install --upgrade pip 'pip-audit>=2.7' + # --locked keeps the resolved tree identical to Cargo.lock; the + # CARGO_NET_* env above plus this outer loop survive transient + # crates.io connection resets without weakening that guarantee. + retry 5 cargo install --locked --version '^0.22' cargo-audit # ───────────────────────────────────────────────────────────── # Python: pip-audit @@ -330,32 +370,60 @@ jobs: # ───────────────────────────────────────────────────────────── # OSV-Scanner: cross-ecosystem advisory DB (PyPI + npm + cargo) # ───────────────────────────────────────────────────────────── + - name: Download + verify OSV-Scanner + # Split out from the scan below so binary integrity is a HARD gate: + # a checksum mismatch (swapped release asset, the Trivy-style pivot + # this workflow refuses) fails the job instead of being swallowed by + # the scan step's continue-on-error. A download still failing after + # retries is transient, so we skip the scan rather than red-fail. + # SHA-256 verified BEFORE chmod +x / exec. Bump OSV_SHA256 in lockstep + # with OSV_VERSION (value from the release's osv-scanner_SHA256SUMS). + run: | + set -euo pipefail + OSV_VERSION="v2.0.2" + OSV_SHA256="3abcfd7126c453a00421487e721b296e0cb68085bd431d6cef60872774170fc8" + if ! curl --proto '=https' --tlsv1.2 -fsSL \ + --retry 5 --retry-delay 3 --retry-connrefused --retry-all-errors \ + -o /tmp/osv-scanner \ + "https://github.com/google/osv-scanner/releases/download/${OSV_VERSION}/osv-scanner_linux_amd64"; then + echo "::warning::osv-scanner download failed after retries; skipping scan" >&2 + rm -f /tmp/osv-scanner + exit 0 # transient availability: do not red-fail the job + fi + if ! echo "${OSV_SHA256} /tmp/osv-scanner" | sha256sum -c -; then + echo "::error::osv-scanner checksum mismatch; refusing to execute" >&2 + rm -f /tmp/osv-scanner + exit 1 # integrity failure: hard-fail + fi + chmod +x /tmp/osv-scanner + /tmp/osv-scanner --version + - name: OSV-Scanner (PyPI + npm + cargo, cross-ecosystem advisories) # OSV's advisory feed is a superset of GitHub-Advisory + RustSec # + npm advisories; running it alongside the per-ecosystem audit # tools catches CVEs that haven't propagated to the per-ecosystem # DBs yet (e.g. langchain-core CVE-2025-68664 was on OSV before # GitHub Advisory). Single binary, one transitive resolver, all - # three lockfile types in one pass. Non-blocking until baselines - # close. + # three lockfile types in one pass. Binary is checksum-verified in + # the step above; only the advisory scan stays non-blocking until + # baselines close. continue-on-error: true run: | set +e - # OSV-Scanner ships a raw binary (no tarball) in v2.x. - curl -fsSL -o /tmp/osv-scanner \ - https://github.com/google/osv-scanner/releases/download/v2.0.2/osv-scanner_linux_amd64 - chmod +x /tmp/osv-scanner - /tmp/osv-scanner --version - /tmp/osv-scanner scan source \ - --lockfile=studio/frontend/package-lock.json \ - --lockfile=studio/src-tauri/Cargo.lock \ - --lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \ - --lockfile=requirements.txt:audit-reqs/studio.txt \ - --lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \ - --lockfile=requirements.txt:audit-reqs/overrides.txt \ - --lockfile=requirements.txt:audit-reqs/extras.txt \ - --lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \ - --format=table 2>&1 | tee logs-osv-scanner.txt + if [ ! -x /tmp/osv-scanner ]; then + echo "osv-scanner unavailable this run; skipping scan" | tee logs-osv-scanner.txt + else + /tmp/osv-scanner scan source \ + --lockfile=studio/frontend/package-lock.json \ + --lockfile=studio/src-tauri/Cargo.lock \ + --lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \ + --lockfile=requirements.txt:audit-reqs/studio.txt \ + --lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \ + --lockfile=requirements.txt:audit-reqs/overrides.txt \ + --lockfile=requirements.txt:audit-reqs/extras.txt \ + --lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \ + --format=table 2>&1 | tee logs-osv-scanner.txt + fi { echo "## OSV-Scanner (cross-ecosystem)" echo @@ -1075,7 +1143,23 @@ jobs: # new-install-script gate below protects against, and we must # not run any third-party hook to set up the audit. working-directory: studio/frontend - run: npm ci --ignore-scripts + run: | + retry() { # retry with exponential backoff + local max="$1"; shift + local n=1 delay=5 + until "$@"; do + if [ "$n" -ge "$max" ]; then + echo "::error::command failed after ${n} attempts: $*" >&2 + return 1 + fi + echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2 + sleep "$delay"; n=$((n + 1)); delay=$((delay * 2)) + done + } + # --ignore-scripts is mandatory here (no third-party hook runs); + # the retry only re-attempts the registry fetch, it never relaxes + # that flag or the package-lock integrity check npm ci enforces. + retry 5 npm ci --ignore-scripts - name: npm audit signatures (informational) # Surfaces unsigned / mis-signed packages from the npm diff --git a/.github/workflows/studio-mac-api-smoke.yml b/.github/workflows/studio-mac-api-smoke.yml index b4e274155e..8839b559fd 100644 --- a/.github/workflows/studio-mac-api-smoke.yml +++ b/.github/workflows/studio-mac-api-smoke.yml @@ -89,13 +89,8 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install pyjwt for the JWT-expiry forge test run: pip install 'pyjwt>=2.6' diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index fab0a36bd1..1096b1abb4 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -114,13 +114,8 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' @@ -369,13 +364,8 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Reset auth + boot Studio (API-only, default tool policy) # We deliberately use the API-only mode rather than @@ -760,13 +750,8 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install OpenAI + Anthropic Python SDKs run: pip install 'openai>=1.50' 'anthropic>=0.40' diff --git a/.github/workflows/studio-mac-install-matrix.yml b/.github/workflows/studio-mac-install-matrix.yml new file mode 100644 index 0000000000..4e2722d1cd --- /dev/null +++ b/.github/workflows/studio-mac-install-matrix.yml @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Proves Studio's llama.cpp install loads on every supported macOS. The heavy +# app smokes stay single-OS; this matrix covers the OS-version dimension cheaply +# (install.sh + binary-load assert). Regression guard for the macOS-version +# selection in studio/install_llama_prebuilt.py. + +name: Mac Studio Install Matrix CI + +on: + pull_request: + paths: + - 'studio/install_llama_prebuilt.py' + - 'studio/setup.sh' + - 'install.sh' + - '.github/scripts/assert-llama-loads.sh' + - '.github/workflows/studio-mac-install-matrix.yml' + push: + branches: [main, pip] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + install-load: + name: Install + load (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + timeout-minutes: 25 + continue-on-error: ${{ matrix.experimental }} + strategy: + fail-fast: false + matrix: + include: + - os: macos-14 # Apple Silicon, macOS 14 Sonoma + experimental: false + - os: macos-15 # Apple Silicon, macOS 15 Sequoia + experimental: false + - os: macos-26 # Apple Silicon, macOS 26 Tahoe + experimental: false + - os: macos-15-intel # Intel x86_64, macOS 15 (informational) + experimental: true + - os: macos-26-intel # Intel x86_64, macOS 26 (last Intel macOS) + experimental: true + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh + + - name: Upload install log + if: always() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: mac-install-matrix-${{ matrix.os }}-log + path: logs/install.log + retention-days: 7 diff --git a/.github/workflows/studio-mac-ui-smoke.yml b/.github/workflows/studio-mac-ui-smoke.yml index b353f0ec83..0176b0a168 100644 --- a/.github/workflows/studio-mac-ui-smoke.yml +++ b/.github/workflows/studio-mac-ui-smoke.yml @@ -89,13 +89,8 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: Install Playwright + Chromium # No --with-deps on Mac: that flag installs Linux apt packages. diff --git a/.github/workflows/studio-mac-update-smoke.yml b/.github/workflows/studio-mac-update-smoke.yml index b65439f174..1eca227096 100644 --- a/.github/workflows/studio-mac-update-smoke.yml +++ b/.github/workflows/studio-mac-update-smoke.yml @@ -67,21 +67,8 @@ jobs: set -o pipefail bash install.sh --local --no-torch 2>&1 | tee logs/install.log - - name: Assert install.sh used the Mac llama.cpp prebuilt - run: | - # Mac install must take the prebuilt path. Source-build - # fallback here is an Unsloth bug. - if grep -q "falling back to source build" logs/install.log; then - echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi - if ! grep -qE "prebuilt installed and validated|prebuilt up to date and validated|bin-macos-arm64" logs/install.log; then - echo "::error::no Mac prebuilt llama.cpp marker in install.log." - grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60 - exit 1 - fi - echo "install.sh installed the Mac prebuilt llama.cpp" + - name: Assert llama.cpp loads on this macOS + run: bash .github/scripts/assert-llama-loads.sh - name: First update should be a no-op (prebuilt already validated) env: diff --git a/README.md b/README.md index ecf0f8a7a3..948d84a789 100644 --- a/README.md +++ b/README.md @@ -202,7 +202,7 @@ unsloth studio -p 8888 #### Nightly: Windows: Run in Windows Powershell: -```bash +```powershell git clone https://github.com/unslothai/unsloth.git cd unsloth git checkout nightly diff --git a/install.ps1 b/install.ps1 index c4e8d8d522..cab66f5ae1 100644 --- a/install.ps1 +++ b/install.ps1 @@ -976,6 +976,7 @@ shell.Run cmd, 0, False # Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null. Write-TauriLog "STEP" "Installing Python" $DetectedPython = Find-CompatiblePython + if ($DetectedPython) { step "python" "Python $($DetectedPython.Version) already installed" } @@ -1239,11 +1240,196 @@ shell.Run cmd, 0, False } } } + # ── AMD ROCm detection (Windows) — mirrors setup.ps1 ── + $HasROCm = $false + $HipSdkInstalled = $false # HIP SDK binary found (independent of device accessibility) + $ROCmGpuLabel = $null + $ROCmVersion = $null + $ROCmGfxArch = $null + if (-not $HasNvidiaSmi) { + # hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution). + # AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type. + $hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue + if (-not $hipinfoExe) { + $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } + $hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } + if ($hipRoot) { + $hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe" + if (Test-Path $hipinfoCandidate) { + Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow + Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow + Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow + $hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate } + } else { + Write-Host " [WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" -ForegroundColor Yellow + Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow + Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow + } + } + } + if ($hipinfoExe) { + $HipSdkInstalled = $true # binary found → SDK is installed regardless of device state + try { + $hipOut = & $hipinfoExe.Source 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") { + $HasROCm = $true + $_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() }) + $_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } + if ($_hipAllArches.Count -gt 0) { + $ROCmGfxArch = if ($_hipVisIdx -lt $_hipAllArches.Count) { $_hipAllArches[$_hipVisIdx] } else { $_hipAllArches[0] } + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + } else { + $ROCmGpuLabel = "AMD ROCm" + } + } elseif ($LASTEXITCODE -ne 0) { + # hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected") + $firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1) + Write-Host " [WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" -ForegroundColor Yellow + Write-Host " $firstLine" -ForegroundColor Yellow + Write-Host " Ensure ROCm drivers are installed: https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow + } + } catch {} + } + if (-not $HasROCm) { + $amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue + if ($amdSmiExe) { + try { + $smiOut = & $amdSmiExe.Source list 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") { + $HasROCm = $true + # Mirror the hipinfo path: collect all gfx tokens in enumeration + # order and pick the runtime-visible one via HIP_VISIBLE_DEVICES. + $_smiVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 } + # Attempt 1: newer amd-smi versions embed the gfx arch in list output. + $_smiGfxTokens = @([regex]::Matches($smiOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() }) + if ($_smiGfxTokens.Count -gt 0) { + $ROCmGfxArch = if ($_smiVisIdx -lt $_smiGfxTokens.Count) { $_smiGfxTokens[$_smiVisIdx] } else { $_smiGfxTokens[0] } + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + } else { + # Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+, + # including the GFX target needed for wheel index selection. + $smiAsicOut = "" + try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {} + $_asicGfxTokens = @([regex]::Matches($smiAsicOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() }) + if ($_asicGfxTokens.Count -gt 0) { + $ROCmGfxArch = if ($_smiVisIdx -lt $_asicGfxTokens.Count) { $_asicGfxTokens[$_smiVisIdx] } else { $_asicGfxTokens[0] } + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + } elseif ($smiAsicOut -match "(?im)Market.?Name\s*[:\|]\s*([^\r\n]+)") { + $ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))" + } else { + $ROCmGpuLabel = "AMD ROCm" + } + } + } + } catch {} + } + } + if (-not $HasROCm) { + try { + $wmiGpu = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue | + Where-Object { $_.Name -match "AMD|Radeon" } | + Select-Object -First 1 + if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name } + } catch {} + } + # ── Arch resolution: env-var override → name inference ────────────── + # Covers users whose amd-smi is too old to report the GFX target and + # who don't have hipinfo (HIP-runtime-only, common on Strix Halo / iGPU). + if ($HasROCm -and -not $ROCmGfxArch) { + # 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running. + if ($env:UNSLOTH_ROCM_GFX_ARCH) { + $ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower() + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan" + } + # 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI). + elseif ($ROCmGpuLabel) { + $nameArchTable = @( + @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 + @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 + @{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail) + @{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point) + @{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop + @{ P = "RX 7600"; A = "gfx1102" } # RDNA 3 + @{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix) + ) + foreach ($row in $nameArchTable) { + if ($ROCmGpuLabel -match $row.P) { + $ROCmGfxArch = $row.A + $ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)" + substep "gfx arch inferred from GPU name: $ROCmGfxArch" "Cyan" + substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$ROCmGfxArch to skip inference next time" "Cyan" + break + } + } + } + } + # Capture ROCm version for wheel selection (hipconfig, then amd-smi). + # Run whenever the HIP SDK binary is present, not just when the device is accessible -- + # hipconfig --version works even when hipinfo reports no ROCm device (driver issue). + if ($HasROCm -or $HipSdkInstalled) { + $hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue + if (-not $hipConfigExe) { + $hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null } + if ($hipRoot) { + $hipConfigCandidate = Join-Path $hipRoot "bin\hipconfig.exe" + if (Test-Path $hipConfigCandidate) { + $hipConfigEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" } + Write-Host " [WARN] hipconfig not on PATH -- located via ${hipConfigEnvLabel}: $hipConfigCandidate" -ForegroundColor Yellow + $hipConfigExe = [PSCustomObject]@{ Source = $hipConfigCandidate } + } + } + } + if ($hipConfigExe) { + try { + $hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String + if ($LASTEXITCODE -eq 0) { + $hipVerLine = ($hipVerOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1).Trim() + if ($hipVerLine -match '(\d+\.\d+)') { + $ROCmVersion = $Matches[1] + $ROCmVersionFull = $hipVerLine + } + } + } catch {} + } + if (-not $ROCmVersion) { + $amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue + if ($amdSmiVer) { + try { + $smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String + if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') { + $ROCmVersion = $Matches[1] + } + } catch {} + } + } + } + } + if ($HasNvidiaSmi) { step "gpu" "NVIDIA GPU detected" + } elseif ($HasROCm) { + step "gpu" $ROCmGpuLabel + $hipSdkPath = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { "on system PATH" } + substep "HIP SDK: $hipSdkPath" + if ($ROCmVersionFull) { substep "hipconfig: $ROCmVersionFull" } + } elseif ($HipSdkInstalled -and $ROCmGpuLabel) { + # HIP SDK is installed but ROCm can't see the device (driver issue, not SDK issue) + $sdkVer = if ($ROCmVersionFull) { " (HIP $ROCmVersionFull)" } else { "" } + step "gpu" "AMD GPU detected -- not ROCm-accessible$sdkVer" "Yellow" + substep "Detected: $ROCmGpuLabel" "Yellow" + substep "[WARN] HIP SDK is installed but hipinfo reports no ROCm-capable device." "Yellow" + substep " This is a driver issue, not an SDK issue." "Yellow" + substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow" + substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" + } elseif ($ROCmGpuLabel) { + step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow" + substep "Detected: $ROCmGpuLabel" "Yellow" + substep "Install the HIP SDK for ROCm GPU inference:" "Yellow" + substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow" } else { step "gpu" "none (chat-only / GGUF)" "Yellow" - substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow" + substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow" } # ── Choose the correct PyTorch index URL based on driver CUDA version ── @@ -1270,14 +1456,73 @@ shell.Run cmd, 0, False return "$baseUrl/cu126" } $TorchIndexUrl = Get-TorchIndexUrl - $TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl + + # ── GPU arch → newest compatible Windows ROCm wheel release ── + # Wheels bundle their own ROCm runtime; the installed HIP SDK version does + # not constrain which release to use. Always picks the newest release that + # supports the GPU architecture. + # ── AMD Windows ROCm: arch-aware pip index (repo.amd.com) ── + # Wheels bundle their own ROCm runtime and support all Python versions. + # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. + $ROCmIndexUrl = $null + $ROCmTorchFloor = $null + if ($HasROCm -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { + $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } + $archFamilyMap = @{ + "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 + "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) + "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 + "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" + "gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100 + } + # gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in + # torch._C._grouped_mm on torch <2.11.0 (rocm7.12 and rocm7.1 respectively). + # TheRock issues #5284 and #3284. Force torch>=2.11.0 so pip never resolves + # to the broken 2.10.0 wheels even though they exist on the AMD index. + # The <2.12.0 ceiling matches the Linux install_python_stack.py constraint + # for the same arches: AMD actively publishes new versions on their index, + # so without a ceiling a future 2.12.0+rocmX.Y wheel would be pulled in + # automatically before it has been validated on these architectures. + # Bump the ceiling here (and in install_python_stack.py) when 2.12.x is + # confirmed working on gfx120X / Strix. + $torchFloorMap = @{ + "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" + "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" + } + $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } + if ($archFamily) { + $ROCmIndexUrl = "$amdIndexBase/$archFamily/" + $ROCmTorchFloor = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { $null } + $archLabel = if ($ROCmGfxArch) { $ROCmGfxArch } else { "AMD GPU" } + substep "$archLabel -- AMD repo.amd.com index selected" "Cyan" + if ($ROCmTorchFloor) { + substep " enforcing $ROCmTorchFloor (known _grouped_mm bug in older wheels)" "Cyan" + } + } elseif ($ROCmGfxArch) { + substep "AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow" + } else { + substep "AMD GPU detected but arch unknown -- falling back to CPU-only PyTorch" "Yellow" + } + } + + if ($ROCmIndexUrl) { + $TorchIndexFamily = "rocm" + } else { + $TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl + } $GpuBranch = Get-TauriGpuBranch $TorchIndexFamily Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version # ── Print CPU-only hint when no GPU detected ── - if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") { + if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") { Write-Host "" - substep "No NVIDIA GPU detected." "Yellow" + if ($HipSdkInstalled -and -not $HasROCm) { + substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow" + } elseif ($ROCmGpuLabel) { + substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow" + } else { + substep "No NVIDIA GPU detected." "Yellow" + } substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow" substep "re-run with --no-torch for a faster, lighter install:" "Yellow" substep ".\install.ps1 --no-torch" "Yellow" @@ -1355,9 +1600,18 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) } } - } elseif ($TorchIndexUrl) { + } elseif ($TorchIndexUrl -or $ROCmIndexUrl) { if ($SkipTorch) { substep "skipping PyTorch (--no-torch flag set)." "Yellow" + } elseif ($ROCmIndexUrl) { + Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" + substep "installing PyTorch from $ROCmIndexUrl..." + $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } + $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio } + if ($torchInstallExit -ne 0) { + Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit) + } } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." diff --git a/install.sh b/install.sh index a3b76a3011..f0af60e2d6 100755 --- a/install.sh +++ b/install.sh @@ -183,10 +183,21 @@ _install_bnb_rocm() { fi if [ -n "$_bnb_whl_url" ]; then substep "installing bitsandbytes for AMD ROCm (pre-release, PR #1887)..." - if run_install_cmd "$_label (pre-release)" "$_venv_py" -m pip install \ - --force-reinstall --no-cache-dir --no-deps "$_bnb_whl_url"; then + _bnb_log=$(mktemp) + if "$_venv_py" -m pip install \ + --disable-pip-version-check \ + --force-reinstall --no-cache-dir --no-deps \ + --retries 8 --timeout 90 \ + "$_bnb_whl_url" >"$_bnb_log" 2>&1; then + rm -f "$_bnb_log" return 0 fi + _bnb_rc=$? + if _is_verbose; then + cat "$_bnb_log" >&2 + fi + rm -f "$_bnb_log" + step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2 substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN" fi run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \ @@ -245,6 +256,9 @@ _tauri_torch_index_family() { rocm[0-9]*.[0-9]*) echo "$_diag_family" ;; *) echo "auto" ;; esac ;; + # AMD arch-specific index (e.g. repo.amd.com/rocm/whl/gfx1151/) -- + # used for Strix Halo/Point where torch 2.11+rocm7.13 has the real fix. + *repo.amd.com/rocm/whl/gfx*|*rocm/whl/gfx*) echo "rocm7.13" ;; "") echo "none" ;; *) echo "auto" ;; esac @@ -1568,16 +1582,19 @@ _find_no_torch_runtime() { } # ── AMD ROCm GPU detection helper ── -# Returns 0 (true) if an actual AMD GPU is present, 1 (false) otherwise. -# Checks rocminfo for gfx[1-9]* (excludes gfx000 CPU agent) and -# amd-smi list for GPU data rows (excludes header-only output). +# Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs +# KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices). _has_amd_rocm_gpu() { if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[0-9]/ && !/Name:[[:space:]]*gfx000/{found=1} END{exit !found}'; then + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then return 0 elif command -v amd-smi >/dev/null 2>&1 && \ amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then return 0 + elif [ -e /dev/kfd ] && \ + awk '/gpu_id/{ if ($2+0 > 0) found=1 } END{ exit !found }' \ + /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then + return 0 fi return 1 } @@ -1656,31 +1673,39 @@ get_torch_index_url() { if [ -n "$_rocm_tag" ]; then # Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older) case "$_rocm_tag" in - rocm[1-5].*) echo "$_base/cpu"; return ;; + rocm[1-5].*) + echo "[WARN] ROCm $_rocm_tag detected but PyTorch ROCm wheels require ROCm 6.0+ -- falling back to CPU-only PyTorch" >&2 + echo "[WARN] Upgrade ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo "$_base/cpu"; return ;; esac - # ROCm 7.2 only has torch 2.11.0 which exceeds current bounds - # (<2.11.0). Fall back to rocm7.1 index which has torch 2.10.0. - # Enumerate explicit versions rather than matching rocm6.* so - # a host on ROCm 6.5 or 6.6 (no PyTorch wheels published) is - # clipped down to the last supported 6.x (rocm6.4) instead of - # constructing https://download.pytorch.org/whl/rocm6.5 which - # returns HTTP 403. PyTorch only ships: rocm5.7, 6.0, 6.1, 6.2, - # 6.3, 6.4, 7.0, 7.1, 7.2 (and 5.7 is below our minimum). - # TODO: uncomment rocm7.2 when the torch upper bound is bumped - # to >=2.11.0. + # Supported tags; 6.5+ clips to rocm6.4, 7.3+ caps to rocm7.2. + # PyTorch publishes major.minor URLs only (no patch level), so + # rocm7.2.1 / rocm6.0.2 / etc. must normalise to rocm7.2 / rocm6.0. case "$_rocm_tag" in - rocm6.0|rocm6.0.*|rocm6.1|rocm6.1.*|rocm6.2|rocm6.2.*|rocm6.3|rocm6.3.*|rocm6.4|rocm6.4.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*) - echo "$_base/$_rocm_tag" ;; + rocm6.0|rocm6.0.*) echo "$_base/rocm6.0" ;; + rocm6.1|rocm6.1.*) echo "$_base/rocm6.1" ;; + rocm6.2|rocm6.2.*) echo "$_base/rocm6.2" ;; + rocm6.3|rocm6.3.*) echo "$_base/rocm6.3" ;; + rocm6.4|rocm6.4.*) echo "$_base/rocm6.4" ;; + rocm7.0|rocm7.0.*) echo "$_base/rocm7.0" ;; + rocm7.1|rocm7.1.*) echo "$_base/rocm7.1" ;; + rocm7.2|rocm7.2.*) echo "$_base/rocm7.2" ;; rocm6.*) # ROCm 6.5+ (no published PyTorch wheels): clip down # to the last supported 6.x wheel set. echo "$_base/rocm6.4" ;; *) - # ROCm 7.2+ (including future 10.x+): cap to rocm7.1 - echo "$_base/rocm7.1" ;; + # ROCm 7.3+ (future): cap to rocm7.2 (latest known) + echo "$_base/rocm7.2" ;; esac return fi + # AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be + # read from any source (amd-smi, /opt/rocm/.info/version, hipconfig, + # dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch. + echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2 + echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2 + echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 echo "$_base/cpu"; return fi # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P). @@ -1760,9 +1785,9 @@ print('cp{}{}'.format(sys.version_info.major, sys.version_info.minor)) } _pick_radeon_wheel() { - # Usage: _pick_radeon_wheel PACKAGE_NAME + # Usage: _pick_radeon_wheel PACKAGE_NAME [VERSION_PREFIX] # Scans $_RADEON_LISTING for the newest wheel whose filename starts exactly - # with PACKAGE_NAME- and matches _RADEON_PYTAG + linux_x86_64. + # with PACKAGE_NAME- (and optionally VERSION_PREFIX) and matches _RADEON_PYTAG + linux_x86_64. # Prints the full URL (resolving relative hrefs against _RADEON_BASE_URL). # # POSIX-compliant pipeline: all href parsing, filtering, and version @@ -1770,11 +1795,12 @@ _pick_radeon_wheel() { # for GNU extensions (grep -o, sort -V) that would break under BSD # or BusyBox coreutils. _pkg="$1" + _ver_prefix="${2:-}" [ -n "$_RADEON_LISTING" ] || return 1 [ -n "$_RADEON_PYTAG" ] || return 1 _tag="$_RADEON_PYTAG" _href=$(printf '%s\n' "$_RADEON_LISTING" \ - | awk -v pkg="$_pkg" -v tag="$_tag" ' + | awk -v pkg="$_pkg" -v tag="$_tag" -v ver_prefix="$_ver_prefix" ' BEGIN { max_pad = ""; max_url = "" } { line = $0 @@ -1788,7 +1814,7 @@ _pick_radeon_wheel() { base = p[n] sub(/[?#].*/, "", base) - prefix = pkg "-" + prefix = pkg "-" ver_prefix # Match cpXY-cpXY or cpXY-abi3 with any linux x86_64 # platform tag (linux_x86_64, manylinux_2_28_x86_64, # manylinux2014_x86_64, etc.) @@ -1822,6 +1848,12 @@ _pick_radeon_wheel() { TORCH_INDEX_URL=$(get_torch_index_url) +# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. +# All other ROCm tags and CUDA stay within <2.11.0. +case "$TORCH_INDEX_URL" in + */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; +esac + # Auto-detect GPU for AMD ROCm based # get_torch_index_url must have chosen */rocm* # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". @@ -1834,6 +1866,78 @@ case "$TORCH_INDEX_URL" in fi ;; esac +# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ─────── +# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug +# that causes a segfault in torch._grouped_mm (moe_utils.py line 167). +# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when +# _amd_gpu_radeon=true the installer silently lands on the broken combo. +# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2. +case "$TORCH_INDEX_URL" in + */rocm7.1|*/rocm7.1.*) + # Collect every gfx token in rocminfo / amd-smi enumeration order + # (skip duplicates), then index by HIP_VISIBLE_DEVICES / + # ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box + # where the user selected the dGPU does NOT get rerouted to the + # Strix per-gfx index. + _gfx_all="" + if command -v rocminfo >/dev/null 2>&1; then + _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + fi + if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + # PowerShell paths also probe `amd-smi static --asic`; mirror it + # so a host with hipinfo-less amd-smi reports the gfx target. + if [ -z "$_gfx_all" ]; then + _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}') + fi + fi + _runtime_gfx="" + if [ -n "$_gfx_all" ]; then + _vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" + _idx=0 + if [ -n "$_vis" ] && [ "$_vis" != "-1" ]; then + _first=${_vis%%,*} + case "$_first" in + ''|*[!0-9]*) _idx=0 ;; + *) _idx=$_first ;; + esac + fi + _runtime_gfx=$(printf '%s\n' "$_gfx_all" | awk -v idx="$_idx" ' + NF && !seen[$0]++ { vals[n++] = $0 } + END { + if (idx < 0 || idx >= n) idx = 0 + if (n > 0) print vals[idx] + }') + fi + _strix_gfx="" + case "$_runtime_gfx" in + gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;; + esac + if [ -n "$_strix_gfx" ]; then + echo "" >&2 + echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2 + echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2 + echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2 + echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2 + echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo "" >&2 + # AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's + # actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred + # over the pytorch.org rocm7.2 fallback because it exercises the real GPU + # kernel path. Set UNSLOTH_AMD_ROCM_MIRROR to override for air-gapped installs. + _amd_strix_base="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}" + # Strip ALL trailing slashes to match Python's .rstrip("/") -- a + # double-/triple-slash mirror URL would otherwise produce 404s on + # strict pip proxies (artifactory, sonatype). + while [ "${_amd_strix_base%/}" != "$_amd_strix_base" ]; do + _amd_strix_base="${_amd_strix_base%/}" + done + TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/" + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + _amd_gpu_radeon=false + fi + ;; +esac _TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL") if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then _TAURI_TORCH_INDEX_FAMILY="radeon" @@ -1841,27 +1945,93 @@ fi _TAURI_GPU_BRANCH=$(_tauri_gpu_branch "$_TAURI_TORCH_INDEX_FAMILY" "$_amd_gpu_radeon") tauri_diag_marker "$_TAURI_GPU_BRANCH" "$_TAURI_TORCH_INDEX_FAMILY" -# ── Print CPU-only hint when no GPU detected ── +# ── GPU detection summary (mirrors install.ps1 step "gpu" block) ── +if _has_usable_nvidia_gpu; then + step "gpu" "NVIDIA GPU detected" +elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then + # Probe gfx arch for the display label, honouring HIP_VISIBLE_DEVICES + _gpu_disp_gfx_all="" + _gpu_disp_mkt="" + if command -v rocminfo >/dev/null 2>&1; then + _gpu_disp_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + _gpu_disp_mkt=$(rocminfo 2>/dev/null | awk -F': ' \ + '/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + fi + if [ -z "$_gpu_disp_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gpu_disp_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_gpu_disp_gfx_all" ] && \ + _gpu_disp_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_gpu_disp_mkt" ] && command -v amd-smi >/dev/null 2>&1; then + _gpu_disp_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ + '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) + fi + _gpu_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" + _gpu_vis_idx=0 + if [ -n "$_gpu_vis" ] && [ "$_gpu_vis" != "-1" ]; then + _gpu_first="${_gpu_vis%%,*}" + case "$_gpu_first" in ''|*[!0-9]*) ;; *) _gpu_vis_idx=$_gpu_first ;; esac + fi + _gpu_disp_gfx=$(printf '%s\n' "$_gpu_disp_gfx_all" | awk -v idx="$_gpu_vis_idx" \ + 'NF && !seen[$0]++ { a[n++]=$0 } END { if(idx>=n) idx=0; if(n>0) print a[idx] }') + # UNSLOTH_ROCM_GFX_ARCH env override (mirrors install.ps1) + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then + _gpu_disp_gfx="${UNSLOTH_ROCM_GFX_ARCH}" + substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_gpu_disp_gfx" + # Name-based arch inference when tools don't report gfx (mirrors install.ps1 nameArchTable) + elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then + case "$_gpu_disp_mkt" in + *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 + *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 + *"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 iGPU + *"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 iGPU + *"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop + *"RX 7600"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 + *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU + esac + if [ -n "$_gpu_disp_gfx" ]; then + substep "gfx arch inferred from GPU name: $_gpu_disp_gfx" + substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$_gpu_disp_gfx to skip inference next time" + fi + fi + # ROCm version via hipconfig, then amd-smi + _gpu_rocm_ver="" + if command -v hipconfig >/dev/null 2>&1; then + _gpu_rocm_ver=$(hipconfig --version 2>/dev/null | awk 'NR==1 && /^[0-9]/{print; exit}' || true) + fi + if [ -z "$_gpu_rocm_ver" ] && command -v amd-smi >/dev/null 2>&1; then + _gpu_rocm_ver=$(amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ + 'NF>1{gsub(/[[:space:]]/,"", $2); print $2; exit}' || true) + fi + if [ -n "$_gpu_disp_gfx" ]; then + step "gpu" "AMD ROCm ($_gpu_disp_gfx)" + else + step "gpu" "AMD ROCm" + fi + _rocm_root="${ROCM_PATH:-${HIP_PATH:-/opt/rocm}}" + substep "ROCm: $_rocm_root" + [ -n "$_gpu_rocm_ver" ] && substep "hipconfig: $_gpu_rocm_ver" + [ -n "$_gpu_disp_mkt" ] && [ -n "$_gpu_disp_gfx" ] && substep "GPU: $_gpu_disp_mkt" +else + step "gpu" "none (CPU-only)" "$C_WARN" +fi + +# ── PyTorch wheel index note ── case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then - echo "" - echo " NOTE: No GPU detected (nvidia-smi and ROCm not found)." - echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference," - echo " re-run with --no-torch for a faster, lighter install:" - echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" - echo " AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" - echo "" + substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" + substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" + substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):" + substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" fi ;; - */rocm*) - echo "" + */rocm*|*/gfx*) if [ "$_amd_gpu_radeon" = true ]; then - echo " AMD Radeon + ROCm detected -- installing PyTorch wheels from repo.radeon.com" + substep "wheels: repo.radeon.com (Radeon)" else - echo " AMD ROCm detected -- installing ROCm-enabled PyTorch ($TORCH_INDEX_URL)" + substep "wheels: $TORCH_INDEX_URL" fi - echo "" ;; esac @@ -1943,24 +2113,23 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ "$_radeon_listing_ok" = true ]; then # Require torch, torchvision, torchaudio wheels to all resolve - # from the Radeon listing. If any is missing for this Python - # tag, fall through to the standard ROCm index instead of - # silently mixing Radeon wheels with PyPI defaults. + # from the Radeon listing. The repo often publishes multiple + # generations simultaneously, so picking the highest-version + # for each package independently can assemble a mismatched trio + # (e.g. torch 2.10 + torchvision 0.24). To prevent this, + # we identify the highest common minor version and downpair + # wheels if necessary to ensure a compatible set. _torch_whl=$(_pick_radeon_wheel "torch" 2>/dev/null) || _torch_whl="" _tv_whl=$(_pick_radeon_wheel "torchvision" 2>/dev/null) || _tv_whl="" _ta_whl=$(_pick_radeon_wheel "torchaudio" 2>/dev/null) || _ta_whl="" _tri_whl=$(_pick_radeon_wheel "triton" 2>/dev/null) || _tri_whl="" - # Sanity-check torch / torchvision / torchaudio are a - # matching release. The Radeon repo publishes multiple - # generations simultaneously, so picking the highest-version - # wheel for each package independently can assemble a - # mismatched trio (e.g. torch 2.9.1 + torchvision 0.23.0 + - # torchaudio 2.9.0 from the current rocm-rel-7.2.1 index). + # Check that torch and torchaudio share the same X.Y public # version prefix, and that torchvision's minor correctly - # pairs with torch's minor (torchvision = torch.minor - 5 + # pairs with torch's minor (torchvision = torch.minor + 15 # since torch 2.4 -> torchvision 0.19 -> torch 2.9 -> # torchvision 0.24). + # # URL-decode each wheel name so %2B -> + before version # extraction. Real Radeon wheel hrefs are percent-encoded # (torch-2.10.0%2Brocm7.2.0...), so a plain [+-] terminator @@ -1968,38 +2137,75 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # _radeon_versions_match would stay false for every real # listing, silently forcing a fallback to the generic # ROCm index. - _torch_ver="" - _tv_ver="" - _ta_ver="" - if [ -n "$_torch_whl" ]; then - _torch_name=$(printf '%s' "${_torch_whl##*/}" | sed 's/%2[Bb]/+/g') - _torch_ver=$(printf '%s\n' "$_torch_name" | sed -n 's|^torch-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p') - fi - if [ -n "$_tv_whl" ]; then - _tv_name=$(printf '%s' "${_tv_whl##*/}" | sed 's/%2[Bb]/+/g') - _tv_ver=$(printf '%s\n' "$_tv_name" | sed -n 's|^torchvision-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p') - fi - if [ -n "$_ta_whl" ]; then - _ta_name=$(printf '%s' "${_ta_whl##*/}" | sed 's/%2[Bb]/+/g') - _ta_ver=$(printf '%s\n' "$_ta_name" | sed -n 's|^torchaudio-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p') - fi + _extract_version() { + _whl=$1 + _pkg=$2 + if [ -n "$_whl" ]; then + _name=$(printf '%s' "${_whl##*/}" | sed 's/%2[Bb]/+/g') + printf '%s\n' "$_name" | sed -n "s|^${_pkg}-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p" + fi + } + + _torch_ver=$(_extract_version "$_torch_whl" "torch") + _tv_ver=$(_extract_version "$_tv_whl" "torchvision") + _ta_ver=$(_extract_version "$_ta_whl" "torchaudio") + _radeon_versions_match=false if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then - _torch_major=${_torch_ver%%.*} _torch_minor=${_torch_ver#*.} - _ta_major=${_ta_ver%%.*} _ta_minor=${_ta_ver#*.} - _tv_major=${_tv_ver%%.*} _tv_minor=${_tv_ver#*.} - # torchvision expected minor (e.g. torch 2.9 -> 0.24) - _expected_tv_minor=$((_torch_minor + 15)) - if [ "$_torch_major" = "$_ta_major" ] && \ - [ "$_torch_minor" = "$_ta_minor" ] && \ - [ "$_tv_major" = "0" ] && \ - [ "$_tv_minor" = "$_expected_tv_minor" ]; then - _radeon_versions_match=true - fi + _tv_equiv_minor=$((_tv_minor - 15)) + + # Determine initial target minor (lowest common denominator) + _target_minor=$_torch_minor + [ "$_tv_equiv_minor" -lt "$_target_minor" ] && _target_minor=$_tv_equiv_minor + [ "$_ta_minor" -lt "$_target_minor" ] && _target_minor=$_ta_minor + + # Loop downwards to find the first complete matching trio. + # This avoids aborting if the repo has gaps. + _attempts=0 + while [ "$_attempts" -lt 5 ] && [ "$_target_minor" -ge 0 ]; do + _expected_tv_minor=$((_target_minor + 15)) + + _curr_torch=$(_pick_radeon_wheel "torch" "2.${_target_minor}." 2>/dev/null) || _curr_torch="" + _curr_tv=$(_pick_radeon_wheel "torchvision" "0.${_expected_tv_minor}." 2>/dev/null) || _curr_tv="" + _curr_ta=$(_pick_radeon_wheel "torchaudio" "2.${_target_minor}." 2>/dev/null) || _curr_ta="" + + if [ -n "$_curr_torch" ] && [ -n "$_curr_tv" ] && [ -n "$_curr_ta" ]; then + # Extract versions from the wheels found in this iteration + _c_torch_ver=$(_extract_version "$_curr_torch" "torch") + _c_tv_ver=$(_extract_version "$_curr_tv" "torchvision") + _c_ta_ver=$(_extract_version "$_curr_ta" "torchaudio") + + # Parse Major.Minor for validation + _c_torch_major=${_c_torch_ver%%.*} + _c_torch_minor=${_c_torch_ver#*.} + _c_ta_major=${_c_ta_ver%%.*} + _c_ta_minor=${_c_ta_ver#*.} + _c_tv_major=${_c_tv_ver%%.*} + _c_tv_minor=${_c_tv_ver#*.} + + # Strict X.Y validation: allow patch versions to differ (e.g. torch 2.9.1 + vision 0.24.0) + # as long as the Major and Minor pairing is correct. + if [ "$_c_torch_major" = "$_c_ta_major" ] && \ + [ "$_c_torch_minor" = "$_c_ta_minor" ] && \ + [ "$_c_tv_major" = "0" ] && \ + [ "$_c_tv_minor" = "$((_c_torch_minor + 15))" ]; then + + _torch_whl=$_curr_torch + _tv_whl=$_curr_tv + _ta_whl=$_curr_ta + _tri_whl="" + _radeon_versions_match=true + break + fi + fi + _target_minor=$((_target_minor - 1)) + _attempts=$((_attempts + 1)) + done fi + if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \ [ "$_radeon_versions_match" != true ]; then substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py new file mode 100644 index 0000000000..606488cc7f --- /dev/null +++ b/scripts/verify_import_hoist.py @@ -0,0 +1,854 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. +"""Deterministic, scope-aware verifier for import-hoisting / alias-rename refactors. + +The risk when moving `from a import b as _b` (or `import b as _b`) to module top +and normalizing `_b` -> `b` is twofold: + + 1. DANGLING ALIAS - a `_b` reference is left un-normalized; it now resolves to + nothing (NameError) or, worse, to some *other* module-level `_b`. + 2. RENAME CLASH - `_b` was an alias on purpose because `b` already meant + something else in that scope; normalizing `_b` -> `b` silently re-points the + reference at the wrong object (no NameError, no pyflakes warning). + +This tool parses BEFORE (a git ref, default origin/main) and AFTER (default HEAD) +for each file, builds a real LEGB scope model (functions, classes, lambdas, +comprehensions, global/nonlocal, args, walrus, star-imports), and resolves every +Name load to its binding. It then compares, PER SCOPE: + + * UNRESOLVED-NEW : loads that resolve to nothing in AFTER but did in BEFORE + (or are newly present) -> catches dangling aliases. + * TARGET-MISSING : an import *target* (e.g. module `glob`, or + `importlib.metadata.version`) that a function resolved to + in BEFORE but no longer resolves to in AFTER -> catches a + function that lost access to a module it still uses. + Robust to alias renames because it compares the *target*, + not the local name. + * TARGET-CHANGED : a load whose resolved import target differs BEFORE vs + AFTER -> catches a rename that re-points to a different + module (the clash case). + * AMBIGUOUS-BIND : a name bound by BOTH an import and a non-import in the same + scope in AFTER (and not in BEFORE) -> the "alias was on + purpose / now collides" smell. + * MODULE-DUP-IMPORT: a module-level name imported and also defined/assigned at + module level (introduced by the change). + * NEW-UNUSED-IMPORT: a module-level import added in AFTER that nothing resolves + to (informational; re-exports are a known false positive). + +Usage: + verify_import_hoist.py [--before REF] [--after REF] ... # compare + verify_import_hoist.py --self-test # prove it catches bugs +Exit code 1 if any non-informational finding. +""" + +from __future__ import annotations + +import argparse +import ast +import builtins +import re as _re_mod +import subprocess +import sys +from dataclasses import dataclass, field + +_BUILTINS = set(dir(builtins)) | { + "__file__", + "__name__", + "__doc__", + "__package__", + "__spec__", + "__loader__", + "__builtins__", + "__class__", + "__annotations__", + "__dict__", + "__qualname__", + "__module__", + "__path__", + "__debug__", + "__import__", + "NotImplemented", + "Ellipsis", + "copyright", + "credits", + "license", + "help", + "exit", + "quit", + "__build_class__", + "__cached__", + "reveal_type", + "reveal_locals", +} + + +# ---------------------------------------------------------------- scope model + + +@dataclass +class Binding: + kind: str # 'import' | 'importfrom' | 'def' | 'class' | 'other' + target: str | None = None # canonical import target id, else None + + +@dataclass +class Scope: + kind: str # 'module' | 'function' | 'class' | 'lambda' | 'comp' + qualname: str + parent: "Scope | None" + bindings: dict[str, list[Binding]] = field(default_factory = dict) + globals: set[str] = field(default_factory = set) + nonlocals: set[str] = field(default_factory = set) + star_import: bool = False + + def add(self, name: str, b: Binding) -> None: + self.bindings.setdefault(name, []).append(b) + + +def _import_target(node: ast.AST, alias: ast.alias) -> tuple[str, str]: + """Return (bound_name, canonical_target_id) for one import alias.""" + if isinstance(node, ast.Import): + bound = alias.asname or alias.name.split(".")[0] + return bound, f"import:{alias.name}" + # ImportFrom + bound = alias.asname or alias.name + mod = ("." * (node.level or 0)) + (node.module or "") + return bound, f"from:{mod}:{alias.name}" + + +class _Builder(ast.NodeVisitor): + """Builds the scope tree + bindings, and records every (scope, Name-load).""" + + def __init__(self): + self.module = Scope("module", "", None) + self.uses: list[tuple[Scope, str, int]] = [] # (scope, name, lineno) hard loads + self.soft_uses: list[ + tuple[Scope, str, int] + ] = [] # annotations: count as "used" + # but never as "unresolved" + # (forward refs / string annos) + + def _visit_annotation(self, node, scope: Scope) -> None: + """Annotation context: with `from __future__ import annotations` these are + never evaluated (strings), and even otherwise they routinely contain forward + references. Record contained names as SOFT uses so an import used only in an + annotation still counts as used, but a forward-ref name is never 'unresolved'.""" + if node is None: + return + for n in ast.walk(node): + if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load): + self.soft_uses.append((scope, n.id, n.lineno)) + + # -- binding helpers -- + def _bind_targets(self, scope: Scope, target: ast.AST) -> None: + for n in ast.walk(target): + if isinstance(n, ast.Name) and isinstance(n.ctx, (ast.Store, ast.Del)): + self._bind_name(scope, n.id, Binding("other")) + elif isinstance(n, ast.Starred): + pass + + def _bind_name(self, scope: Scope, name: str, b: Binding) -> None: + if name in scope.globals: + self.module.add(name, b) + elif name in scope.nonlocals: + p = scope.parent + while p is not None and p.kind not in ("function", "lambda"): + p = p.parent + (p or self.module).add(name, b) + else: + scope.add(name, b) + + # -- generic dispatch within a scope -- + def _visit_body(self, stmts, scope: Scope) -> None: + for s in stmts: + self._visit_stmt(s, scope) + + def _visit_stmt(self, node: ast.AST, scope: Scope) -> None: + if isinstance(node, (ast.Import, ast.ImportFrom)): + star = isinstance(node, ast.ImportFrom) and any( + a.name == "*" for a in node.names + ) + if star: + scope.star_import = True + for alias in node.names: + if alias.name == "*": + continue + bound, target = _import_target(node, alias) + kind = "import" if isinstance(node, ast.Import) else "importfrom" + self._bind_name(scope, bound, Binding(kind, target)) + return + if isinstance(node, ast.Global): + scope.globals.update(node.names) + return + if isinstance(node, ast.Nonlocal): + scope.nonlocals.update(node.names) + return + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + self._bind_name(scope, node.name, Binding("def")) + # decorators / defaults evaluate in the ENCLOSING scope + for d in node.decorator_list: + self._visit_expr(d, scope) + self._visit_arg_defaults(node.args, scope) + child = Scope("function", f"{scope.qualname}.{node.name}", scope) + self._bind_type_params(node, child) + self._bind_args(node.args, child) + # arg + return annotations: soft uses (may be strings / forward refs) + for a in self._all_args(node.args): + self._visit_annotation(a.annotation, child) + self._visit_annotation(getattr(node, "returns", None), child) + self._visit_body(node.body, child) + return + if isinstance(node, ast.ClassDef): + self._bind_name(scope, node.name, Binding("class")) + for d in node.decorator_list: + self._visit_expr(d, scope) + for b in node.bases: + self._visit_expr(b, scope) + for kw in node.keywords: + self._visit_expr(kw.value, scope) + child = Scope("class", f"{scope.qualname}.{node.name}", scope) + self._bind_type_params(node, child) + self._visit_body(node.body, child) + return + if isinstance(node, ast.Match): + self._visit_expr(node.subject, scope) + for case in node.cases: + self._bind_pattern(case.pattern, scope) + if case.guard is not None: + self._visit_expr(case.guard, scope) + self._visit_body(case.body, scope) + return + if isinstance(node, getattr(ast, "TryStar", ())): # py3.11 except* + self._visit_body(node.body, scope) + for h in node.handlers: + if h.type is not None: + self._visit_expr(h.type, scope) + if h.name: + self._bind_name(scope, h.name, Binding("other")) + self._visit_body(h.body, scope) + self._visit_body(node.orelse, scope) + self._visit_body(node.finalbody, scope) + return + if isinstance(node, getattr(ast, "TypeAlias", ())): # py3.12 `type X = ...` + if isinstance(node.name, ast.Name): + self._bind_name(scope, node.name.id, Binding("other")) + self._visit_annotation(node.value, scope) + return + if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + val = node.value + if val is not None: + self._visit_expr(val, scope) + if isinstance(node, ast.AnnAssign) and node.annotation is not None: + self._visit_annotation(node.annotation, scope) + for t in targets: + self._bind_targets(scope, t) + # AugAssign target is also a load + if isinstance(node, ast.AugAssign): + self._record_loads(t, scope) + return + if isinstance(node, (ast.For, ast.AsyncFor)): + self._visit_expr(node.iter, scope) + self._bind_targets(scope, node.target) + self._visit_body(node.body, scope) + self._visit_body(node.orelse, scope) + return + if isinstance(node, (ast.With, ast.AsyncWith)): + for item in node.items: + self._visit_expr(item.context_expr, scope) + if item.optional_vars is not None: + self._bind_targets(scope, item.optional_vars) + self._visit_body(node.body, scope) + return + if isinstance(node, ast.Try): + self._visit_body(node.body, scope) + for h in node.handlers: + if h.type is not None: + self._visit_expr(h.type, scope) + if h.name: + self._bind_name(scope, h.name, Binding("other")) + self._visit_body(h.body, scope) + self._visit_body(node.orelse, scope) + self._visit_body(node.finalbody, scope) + return + # generic statement: visit all child expressions/stmts in same scope + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.stmt): + self._visit_stmt(child, scope) + else: + self._visit_expr(child, scope) + + # -- expressions -- + def _visit_arg_defaults(self, args: ast.arguments, scope: Scope) -> None: + for d in list(args.defaults) + [d for d in args.kw_defaults if d is not None]: + self._visit_expr(d, scope) + + def _all_args(self, args: ast.arguments) -> list[ast.arg]: + out = list(args.posonlyargs) + list(args.args) + list(args.kwonlyargs) + if args.vararg: + out.append(args.vararg) + if args.kwarg: + out.append(args.kwarg) + return out + + def _bind_args(self, args: ast.arguments, scope: Scope) -> None: + for a in self._all_args(args): + scope.add(a.arg, Binding("other")) + + def _bind_type_params(self, node, scope: Scope) -> None: + for tp in getattr(node, "type_params", []) or []: + name = getattr(tp, "name", None) + if isinstance(name, str): + scope.add(name, Binding("other")) + self._visit_annotation(getattr(tp, "bound", None), scope) + self._visit_annotation(getattr(tp, "default_value", None), scope) + + def _bind_pattern(self, pat, scope: Scope) -> None: + if pat is None: + return + if isinstance(pat, ast.MatchValue): + self._visit_expr(pat.value, scope) + elif isinstance(pat, ast.MatchSingleton): + pass + elif isinstance(pat, ast.MatchSequence): + for p in pat.patterns: + self._bind_pattern(p, scope) + elif isinstance(pat, ast.MatchStar): + if pat.name: + self._bind_name(scope, pat.name, Binding("other")) + elif isinstance(pat, ast.MatchMapping): + for k in pat.keys: + self._visit_expr(k, scope) + for p in pat.patterns: + self._bind_pattern(p, scope) + if pat.rest: + self._bind_name(scope, pat.rest, Binding("other")) + elif isinstance(pat, ast.MatchClass): + self._visit_expr(pat.cls, scope) + for p in pat.patterns: + self._bind_pattern(p, scope) + for p in pat.kwd_patterns: + self._bind_pattern(p, scope) + elif isinstance(pat, ast.MatchAs): + self._bind_pattern(pat.pattern, scope) + if pat.name: + self._bind_name(scope, pat.name, Binding("other")) + elif isinstance(pat, ast.MatchOr): + for p in pat.patterns: + self._bind_pattern(p, scope) + + def _record_loads(self, node: ast.AST, scope: Scope) -> None: + for n in ast.walk(node): + if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load): + self.uses.append((scope, n.id, n.lineno)) + + def _visit_expr(self, node: ast.AST, scope: Scope) -> None: + if isinstance(node, ast.Name): + if isinstance(node.ctx, ast.Load): + self.uses.append((scope, node.id, node.lineno)) + elif isinstance(node.ctx, (ast.Store, ast.Del)): + self._bind_name(scope, node.id, Binding("other")) + return + if isinstance(node, ast.Lambda): + self._visit_arg_defaults(node.args, scope) + child = Scope("lambda", f"{scope.qualname}.", scope) + self._bind_args(node.args, child) + self._visit_expr(node.body, child) + return + if isinstance( + node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp) + ): + child = Scope("comp", f"{scope.qualname}.", scope) + for i, gen in enumerate(node.generators): + # first iterable is evaluated in the enclosing scope + self._visit_expr(gen.iter, scope if i == 0 else child) + self._bind_targets(child, gen.target) + for cond in gen.ifs: + self._visit_expr(cond, child) + if isinstance(node, ast.DictComp): + self._visit_expr(node.key, child) + self._visit_expr(node.value, child) + else: + self._visit_expr(node.elt, child) + return + if isinstance(node, ast.NamedExpr): # walrus binds in enclosing scope + self._visit_expr(node.value, scope) + if isinstance(node.target, ast.Name): + self._bind_name(scope, node.target.id, Binding("other")) + return + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.stmt): + self._visit_stmt(child, scope) + else: + self._visit_expr(child, scope) + + def run(self, tree: ast.Module) -> None: + self._visit_body(tree.body, self.module) + + +# ---------------------------------------------------------------- resolution + + +def _any_star(scope: Scope) -> bool: + c = scope + while c is not None: + if c.star_import: + return True + c = c.parent + return False + + +def _resolve(scope: Scope, name: str): + """LEGB resolution. Returns (status, bindings) where status in + {'local','import','other','builtin','star','unresolved'}.""" + # global / nonlocal redirection + start = scope + if name in scope.globals: + chain = [_module_of(scope)] + elif name in scope.nonlocals: + chain = _enclosing_functions(scope) + else: + chain = _legb_chain(scope) + for i, sc in enumerate(chain): + if sc is None: + continue + if name in sc.bindings: + binds = sc.bindings[name] + if any(b.kind in ("import", "importfrom") for b in binds): + return "import", binds + return "other", binds + if name in _BUILTINS: + return "builtin", [] + if _any_star(start): + return "star", [] + return "unresolved", [] + + +def _module_of(scope: Scope) -> Scope: + while scope.parent is not None: + scope = scope.parent + return scope + + +def _enclosing_functions(scope: Scope) -> list[Scope]: + out = [] + p = scope.parent + while p is not None: + if p.kind in ("function", "lambda"): + out.append(p) + p = p.parent + out.append(_module_of(scope)) + return out + + +def _legb_chain(scope: Scope) -> list[Scope]: + """Immediate scope, then enclosing scopes skipping class scopes, then module.""" + chain = [scope] + p = scope.parent + while p is not None: + if ( + p.kind != "class" or p.parent is None + ): # module-level class never happens; keep module + if p.kind != "class": + chain.append(p) + p = p.parent + return chain + + +# ---------------------------------------------------------------- analysis + + +def _analyze(src: str): + tree = ast.parse(src) + b = _Builder() + b.run(tree) + # Per-scope: unresolved load names, and import targets it resolves to. + unresolved: dict[str, set[str]] = {} + targets_by_scope: dict[str, set[str]] = {} + target_by_use: dict[tuple[str, str], set[str]] = {} + for scope, name, _ln in b.uses: + status, binds = _resolve(scope, name) + if status == "unresolved": + unresolved.setdefault(scope.qualname, set()).add(name) + elif status == "import": + tids = {bd.target for bd in binds if bd.target} + targets_by_scope.setdefault(scope.qualname, set()).update(tids) + target_by_use.setdefault((scope.qualname, name), set()).update(tids) + # soft uses (annotations): only contribute to "used", never to "unresolved" + for scope, name, _ln in b.soft_uses: + status, binds = _resolve(scope, name) + if status == "import": + tids = {bd.target for bd in binds if bd.target} + targets_by_scope.setdefault(scope.qualname, set()).update(tids) + # module-level binding info for clash checks + module = b.module + module_imports = { + n: bs + for n, bs in module.bindings.items() + if any(x.kind in ("import", "importfrom") for x in bs) + } + module_dup = { + n + for n, bs in module.bindings.items() + if any(x.kind in ("import", "importfrom") for x in bs) + and any(x.kind not in ("import", "importfrom") for x in bs) + } + # ambiguous: any scope where a name is bound by import AND non-import + ambiguous: dict[str, set[str]] = {} + + def walk_scopes(scope: Scope): + for n, bs in scope.bindings.items(): + if any(x.kind in ("import", "importfrom") for x in bs) and any( + x.kind not in ("import", "importfrom") for x in bs + ): + ambiguous.setdefault(scope.qualname, set()).add(n) + # scope tree isn't stored; rebuild via uses is hard. We approximate with module only. + + walk_scopes(module) + return { + "unresolved": unresolved, + "targets_by_scope": targets_by_scope, + "target_by_use": target_by_use, + "module_import_targets": { + n: {x.target for x in bs if x.target} for n, bs in module_imports.items() + }, + "module_dup": module_dup, + "ambiguous": ambiguous, + } + + +def _git_show(ref: str, path: str) -> str | None: + try: + return subprocess.run( + ["git", "show", f"{ref}:{path}"], capture_output = True, text = True, check = True + ).stdout + except subprocess.CalledProcessError: + return None + + +def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]: + """Return list of (severity, message). severity in BLOCKER/WARN/INFO. + + Blocker signals (precise, no relocation false-positives): + UNRESOLVED-NEW - a load became undefined (dangling alias / removed import). + NEW-UNUSED-HOIST - a module-level import added by THIS change is resolved by + NO load. A correct hoist always wires its new import to a + reference; if the alias was left un-normalized OR renamed + to the wrong name, the hoisted import ends up unused. This + single signal catches BOTH user-described failure modes and + does NOT fire for code merely relocated to another file + (that removes the import, it doesn't add an unused one). + TARGET-CHANGED - the same (scope, name) load resolves to a different import + target before vs after (a same-name re-point). + """ + a = _analyze(before_src) + b = _analyze(after_src) + findings: list[tuple[str, str]] = [] + + def used_targets(analysis) -> set[str]: + out: set[str] = set() + for tids in analysis["targets_by_scope"].values(): + out |= tids + return out + + before_used = used_targets(a) + after_used = used_targets(b) + before_module_targets: set[str] = set() + for tids in a["module_import_targets"].values(): + before_module_targets |= tids + after_module_targets: set[str] = set() + for tids in b["module_import_targets"].values(): + after_module_targets |= tids + added_module_targets = after_module_targets - before_module_targets + + # 1. UNRESOLVED-NEW + for scope, names in b["unresolved"].items(): + new = names - a["unresolved"].get(scope, set()) + for n in sorted(new): + findings.append( + ( + "BLOCKER", + f"{path}: UNRESOLVED-NEW '{n}' in scope {scope} " + f"(undefined after change -> dangling alias / removed import)", + ) + ) + + # 2. HOISTED-IMPORT-UNUSED (the core botched-hoist / wrong-rename signal) + # A module-level import in AFTER that NO load resolves to, and which was + # either newly added by this change OR was actually used before. Excludes: + # - relocation (the import is REMOVED, so it's not in after at all) + # - stable pre-existing re-exports (unused before AND after, not newly added) + for n, tids in b["module_import_targets"].items(): + if tids & after_used: + continue # resolved by something -> fine + newly_added = bool(tids - before_module_targets) + was_used_before = bool(tids & before_used) + if newly_added or was_used_before: + why = ( + "added but unused" + if newly_added + else "was used before, now unused (references re-pointed)" + ) + findings.append( + ( + "BLOCKER", + f"{path}: HOISTED-IMPORT-UNUSED '{n}' ({sorted(tids)}) " + f"{why} -> un-normalized alias or wrong rename target?", + ) + ) + + # 3. TARGET-CHANGED (same scope+name resolves to a different import target) + for key, tafter in b["target_by_use"].items(): + tbefore = a["target_by_use"].get(key) + if tbefore and tbefore != tafter: + findings.append( + ( + "BLOCKER", + f"{path}: TARGET-CHANGED name '{key[1]}' in {key[0]} " + f"{sorted(tbefore)} -> {sorted(tafter)} (rename re-points module)", + ) + ) + + # 4. MODULE-DUP-IMPORT introduced + for n in sorted(b["module_dup"] - a["module_dup"]): + findings.append( + ( + "WARN", + f"{path}: MODULE-DUP-IMPORT '{n}' bound by import AND non-import " + f"at module level (possible clash)", + ) + ) + + # 5. AMBIGUOUS-BIND introduced (module scope) + for scope, names in b["ambiguous"].items(): + new = names - a["ambiguous"].get(scope, set()) + for n in sorted(new): + findings.append( + ("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}") + ) + + # 6. TARGET-MISSING (informational): a scope stopped resolving to an import + # target. Real bugs are already covered above; remaining cases are code + # relocated to another file (e.g. a moved helper). Shown for transparency. + for scope, tbefore in a["targets_by_scope"].items(): + tafter = b["targets_by_scope"].get(scope, set()) + for t in sorted(tbefore - tafter): + relocated = ( + "" + if t in added_module_targets + else " [target not re-added here -> likely relocated/deleted]" + ) + findings.append( + ("INFO", f"{path}: TARGET-MISSING {t} in scope {scope}{relocated}") + ) + return findings + + +# ---------------------------------------------------------------- self-test + +_SELF_TESTS = { + "dangling_alias": ( + # before: inline aliased import, used as _b + "import os\n" + "def f():\n" + " import glob as _b\n" + " return _b.glob('*')\n", + # after: hoisted to canonical, but reference NOT normalized -> _b dangles + "import os\n" "import glob\n" "def f():\n" " return _b.glob('*')\n", + "BLOCKER", + ), + "rename_clash": ( + # before: _b is a deliberate alias; `b` already means something else + "import re as _b\n" "b = 123\n" "def f():\n" " return _b.compile('x'), b\n", + # after: someone normalized _b -> b ; now f().b is the int, re is lost + "import re\n" "b = 123\n" "def f():\n" " return b.compile('x'), b\n", + "BLOCKER", # TARGET-MISSING from:.. or import:re in f + ), + "clean_rename": ( + "def f():\n" " import glob as _g\n" " return _g.glob('*')\n", + "import glob\n" "def f():\n" " return glob.glob('*')\n", + None, # expect NO blocker + ), + "clean_dedup_redundant": ( + "import sys\n" "def f():\n" " import sys\n" " return sys.argv\n", + "import sys\n" "def f():\n" " return sys.argv\n", + None, + ), + "from_import_dangling": ( + # from-import alias left un-normalized + "def f():\n" + " from importlib.metadata import version as _v\n" + " return _v('x')\n", + "from importlib.metadata import version\n" "def f():\n" " return _v('x')\n", + "BLOCKER", + ), + "local_var_clash": ( + # _b renamed to b, but b is a LOCAL variable in f -> import silently unused + "def f(b):\n" " import re as _b\n" " return _b.compile(b)\n", + "import re\n" + "def f(b):\n" + " return b.compile(b)\n", # 'b' is the param, not the module + "BLOCKER", + ), + "substring_safe": ( + # correct _copy->copy rename while a config_copy var exists: NO false positive + "def f(config):\n" + " import copy as _copy\n" + " config_copy = _copy.deepcopy(config)\n" + " return config_copy\n", + "import copy\n" + "def f(config):\n" + " config_copy = copy.deepcopy(config)\n" + " return config_copy\n", + None, + ), + "attr_access_not_a_use": ( + # x._b is attribute access, not a use of name _b; removing import _b is fine + "import os\n" + "def f(x):\n" + " import sys as _b\n" + " return x._b + _b.argv[0]\n", + "import os\n" "import sys\n" "def f(x):\n" " return x._b + sys.argv[0]\n", + None, + ), +} + + +def _self_test() -> int: + ok = True + for name, (before, after, expect) in _SELF_TESTS.items(): + findings = compare(before, after, f"<{name}>") + blockers = [m for sev, m in findings if sev == "BLOCKER"] + got = "BLOCKER" if blockers else None + passed = got == expect + ok = ok and passed + print(f"[{'PASS' if passed else 'FAIL'}] {name}: expect={expect} got={got}") + for sev, m in findings: + print(f" ({sev}) {m}") + print("\nSELF-TEST:", "ALL PASS" if ok else "FAILURES") + return 0 if ok else 1 + + +def _pyflakes_undefined(path: str) -> set[str] | None: + """Return the set of names pyflakes reports as 'undefined name' for `path`, + or None if pyflakes failed to run/parse the file.""" + try: + proc = subprocess.run( + [sys.executable, "-m", "pyflakes", path], capture_output = True, text = True + ) + except Exception: + return None + if "syntax error" in (proc.stdout + proc.stderr).lower(): + return None + names = set() + for line in proc.stdout.splitlines(): + m = _re_mod.search(r"undefined name '([^']+)'", line) + if m: + names.add(m.group(1)) + return names + + +def audit_files(paths: list[str]) -> int: + """Single-version robustness audit. For every file: confirm the analyzer does + not crash, then cross-check its 'unresolved' names against pyflakes. Any name + the resolver flags that pyflakes does NOT call undefined is a tool FALSE + POSITIVE (a resolver gap to fix).""" + n_files = n_err = n_fp = n_syntax = 0 + fp_detail: dict[str, set[str]] = {} + err_detail: dict[str, str] = {} + for path in paths: + n_files += 1 + try: + src = open(path, encoding = "utf-8").read() + except Exception as e: # unreadable + n_err += 1 + err_detail[path] = f"read: {e}" + continue + try: + res = _analyze(src) + except SyntaxError: + n_syntax += 1 + continue + except Exception as e: # analyzer crash -> robustness bug + n_err += 1 + err_detail[path] = f"{type(e).__name__}: {e}" + continue + tool_unresolved = set() + for names in res["unresolved"].values(): + tool_unresolved |= names + if not tool_unresolved: + continue + pf = _pyflakes_undefined(path) + if pf is None: + continue # pyflakes couldn't adjudicate; skip cross-check + false_pos = tool_unresolved - pf + if false_pos: + n_fp += 1 + fp_detail[path] = false_pos + print(f"audited files : {n_files}") + print(f"syntax-skipped : {n_syntax}") + print(f"analyzer errors : {n_err}") + for p, e in sorted(err_detail.items()): + print(f" ERROR {p}: {e}") + print(f"false-positive files: {n_fp} (resolver flagged a name pyflakes accepts)") + for p, names in sorted(fp_detail.items()): + print(f" FP {p}: {sorted(names)}") + ok = n_err == 0 and n_fp == 0 + print( + "\nAUDIT:", + "ROBUST (no crashes, no false positives vs pyflakes)" + if ok + else "NEEDS WORK (see above)", + ) + return 0 if ok else 1 + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--before", default = "origin/main") + ap.add_argument("--after", default = "HEAD") + ap.add_argument("--self-test", action = "store_true") + ap.add_argument( + "--audit", + action = "store_true", + help = "single-version robustness audit on filesystem paths", + ) + ap.add_argument("files", nargs = "*") + args = ap.parse_args() + + if args.self_test: + return _self_test() + if args.audit: + return audit_files(args.files) + + any_blocker = False + for path in args.files: + before = _git_show(args.before, path) + after = _git_show(args.after, path) + if after is None: + print(f"SKIP {path}: not found at {args.after}") + continue + if before is None: + before = "" # new file + findings = compare(before, after, path) + blockers = [f for f in findings if f[0] == "BLOCKER"] + warns = [f for f in findings if f[0] == "WARN"] + infos = [f for f in findings if f[0] == "INFO"] + status = ( + "CLEAN" + if not blockers and not warns + else ("BLOCKERS" if blockers else "WARNINGS") + ) + print(f"\n=== {path}: {status} ===") + for sev, m in blockers + warns + infos: + print(f" [{sev}] {m}") + any_blocker = any_blocker or bool(blockers) + print( + "\nOVERALL:", "FAIL (blockers found)" if any_blocker else "PASS (no blockers)" + ) + return 1 if any_blocker else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py new file mode 100644 index 0000000000..5650a60ee2 --- /dev/null +++ b/studio/backend/core/_torchao_stub.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared torchao Windows-ROCm import stub. + +torchao (pulled in by transformers.quantizers) imports +torch.distributed._functional_collectives at module level, which imports +distributed_c10d.py unconditionally — that file crashes on Windows ROCm because +torch._C._distributed_c10d (the RCCL backend) is absent. +torch/distributed/__init__.py itself is guarded by `if is_available()` so +`import torch.distributed` alone is safe; the crash only comes via torchao's +import chain. Stubbing torchao short-circuits it entirely. +_StubSubpackageFinder handles any depth of torchao.xxx.yyy imports. + +This logic used to be duplicated inline inside run_export_process() and +run_training_process(); it now lives here so both worker subprocesses call the +single `install_torchao_windows_rocm_stub()` entrypoint before importing +transformers / unsloth_zoo. +""" + +from __future__ import annotations + +import sys +import types +import importlib.abc +import importlib.machinery + +_STUB_SENTINEL = object() + + +# Metaclass for stub types so that isinstance(x, StubClass) returns False +# instead of raising TypeError ("arg 2 must be a type"). +# peft/tuners/lora/torchao.py does: +# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor +# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) +# If those names resolve to stub modules rather than types, isinstance() raises. +class _StubTypeMeta(type): + def __instancecheck__(cls, instance): + return False + + def __subclasscheck__(cls, subclass): + return False + + def __getattr__(cls, attr): + if attr.startswith("__"): + raise AttributeError(attr) + child = _StubTypeMeta(attr, (), {}) + setattr(cls, attr, child) + return child + + def __call__(cls, *args, **kwargs): + return None + + +def _make_stub_type(name): + """Stub class: accepted by isinstance() (always False), supports attr access.""" + return _StubTypeMeta(name, (), {}) + + +def _make_mod_stub(mod_name): + m = types.ModuleType(mod_name) + m.__path__ = [] + m.__package__ = mod_name + m._unsloth_stub = _STUB_SENTINEL + m.__spec__ = importlib.machinery.ModuleSpec(mod_name, loader = None, is_package = True) + + def _ga(attr, _m = m, _n = mod_name): + if attr.startswith("__"): + raise AttributeError(attr) + # Return a stub CLASS (not a module) so that isinstance(x, attr) + # works and returns False instead of raising TypeError. + child = _make_stub_type(f"{_n}.{attr}") + setattr(_m, attr, child) + return child + + m.__getattr__ = _ga + return m + + +class _StubSubpackageLoader(importlib.abc.Loader): + def __init__(self, mod_name): + self._mod_name = mod_name + + def create_module(self, spec): + return _make_mod_stub(self._mod_name) + + def exec_module(self, module): + pass + + +class _StubSubpackageFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target = None): + if "." not in fullname: + return None + parent = sys.modules.get(fullname.rsplit(".", 1)[0]) + if parent is None: + return None + if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL: + return None + return importlib.machinery.ModuleSpec( + fullname, _StubSubpackageLoader(fullname), is_package = True + ) + + +def install_torchao_windows_rocm_stub() -> None: + """Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash. + + No-op on every other platform (Windows CUDA included — there torchao is real + and shadowing it would break torchao-based quantization paths). Must run + before any import of transformers / unsloth_zoo. Safe to call once per worker + process. + """ + # Gate on the active torch runtime, not env-var presence -- HIP_PATH / + # ROCM_PATH stay set after a user installs the HIP SDK and reverts to a + # CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip + # but still encode "rocm" in torch.__version__, so accept either. + _is_win32_rocm = False + if sys.platform == "win32": + try: + import torch as _torch_probe + + _is_win32_rocm = bool( + getattr(getattr(_torch_probe, "version", None), "hip", None) + or "rocm" in getattr(_torch_probe, "__version__", "").lower() + ) + del _torch_probe + except Exception: + pass + if _is_win32_rocm: + # Register the finder only on Windows ROCm -- on other platforms there + # are no stub modules seeded, so appending is a pure accumulation. + sys.meta_path.append(_StubSubpackageFinder()) + # Seed torchao top-level + key submodules; the finder handles the rest. + for _tao_name in ( + "torchao", + "torchao.quantization", + "torchao.dtypes", + "torchao.float8", + "torchao.utils", + ): + if _tao_name not in sys.modules: + sys.modules[_tao_name] = _make_mod_stub(_tao_name) diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index f77b1966c4..defcff924b 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -439,6 +439,15 @@ def run_export_process( 'Install for better performance: pip install "triton-windows<3.7"' ) + # ── 1c. Stub torchao on Windows ROCm ── + # Shared with the training worker; see core/_torchao_stub.py for the full + # rationale (torchao -> torch.distributed._functional_collectives crashes on + # Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm. + # Must run before any import of transformers / unsloth_zoo. + from core._torchao_stub import install_torchao_windows_rocm_stub + + install_torchao_windows_rocm_stub() + # ── 2. Import ML libraries (fresh in this clean process) ── try: _send_response( diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 80ead53960..c8250bd4f5 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -17,6 +17,7 @@ import struct import structlog from loggers import get_logger import shutil +import signal import socket import subprocess import sys @@ -965,9 +966,6 @@ class LlamaCppBackend: 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 @@ -1238,6 +1236,33 @@ class LlamaCppBackend: 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. @@ -1255,8 +1280,6 @@ class LlamaCppBackend: 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( @@ -3158,6 +3181,14 @@ class LlamaCppBackend: 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( @@ -3167,6 +3198,24 @@ class LlamaCppBackend: ) 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. /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.) @@ -3875,10 +3924,6 @@ class LlamaCppBackend: Falls back to pgrep + /proc//exe on Linux when psutil is not installed. """ - import os - import signal - import sys - try: # -- Build the ownership allowlist -------------------------------- # Two kinds of matches: diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index b9643cac6a..0365b3ffd8 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -6,8 +6,10 @@ Unsloth Training Backend Integrates Unsloth training capabilities with the FastAPI backend """ +import gc import os import sys +import types # Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -42,7 +44,10 @@ from utils.hardware import ( get_visible_gpu_count, ) -torch._dynamo.config.recompile_limit = 64 +# recompile_limit was removed in some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2). +# Guard so training doesn't crash on RDNA2/RDNA3 with older ROCm torch wheels. +if hasattr(torch._dynamo.config, "recompile_limit"): + torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported from unsloth.chat_templates import get_chat_template @@ -417,8 +422,6 @@ class UnslothTrainer: in sys.modules. When the next training run calls dataset.map(num_proc=N), forked child processes inherit this stale state and deadlock. """ - import sys as _sys - # Remove cloned audio repo paths from sys.path base_dir = os.path.dirname(os.path.abspath(__file__)) audio_paths = [ @@ -433,15 +436,15 @@ class UnslothTrainer: removed_paths = [] for path in audio_paths: - if path in _sys.path: - _sys.path.remove(path) + if path in sys.path: + sys.path.remove(path) removed_paths.append(path) # Remove stale audio modules from sys.modules prefixes = ("snac", "whisper", "sparktts", "outetts") - removed_modules = [key for key in _sys.modules if key.startswith(prefixes)] + removed_modules = [key for key in sys.modules if key.startswith(prefixes)] for key in removed_modules: - del _sys.modules[key] + del sys.modules[key] if removed_paths or removed_modules: logger.info( @@ -538,10 +541,9 @@ class UnslothTrainer: # clear_unsloth_compiled_cache() deletes the disk cache, but the flag # prevents re-compilation — leaving missing cache files. Reloading # restores original class definitions so Unsloth can re-compile cleanly. - import sys as _sys import importlib - for _key, _mod in list(_sys.modules.items()): + for _key, _mod in list(sys.modules.items()): if "transformers.models." in _key and ".modeling_" in _key: if hasattr(_mod, "__UNSLOTH_PATCHED__"): try: @@ -657,6 +659,23 @@ class UnslothTrainer: f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" ) + # AMD ROCm hardware without native bfloat16 (e.g. RDNA2 / gfx103x) + # crashes with an LLVM error at the first bf16 kernel dispatch if + # dtype=None lets unsloth auto-pick bf16. Force float16 there so that + # path is never reached. NVIDIA keeps dtype=None so unsloth's own + # bf16/fp16/float32 auto-detection (including FORCE_FLOAT32 models) is + # honored -- older NVIDIA without bf16 (T4/V100) must NOT be coerced to + # float16 here, which the previous unconditional branch did wrongly. + # Derive ROCm inline (not hardware.IS_ROCM) because that flag is unset + # until detect_hardware() runs, which isn't guaranteed in this subprocess. + _is_rocm = ( + bool(getattr(torch.version, "hip", None)) + or "rocm" in torch.__version__.lower() + ) + _auto_dtype = ( + torch.float16 if (_is_rocm and not is_bfloat16_supported()) else None + ) + # Branch based on model type if self._audio_type == "csm": # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False @@ -666,7 +685,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, + dtype = _auto_dtype, auto_model = CsmForConditionalGeneration, load_in_4bit = False, device_map = device_map, @@ -683,7 +702,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastModel.from_pretrained( model_name = model_name, - dtype = None, + dtype = _auto_dtype, load_in_4bit = False, device_map = device_map, full_finetuning = full_finetuning, @@ -705,7 +724,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastLanguageModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, @@ -777,7 +796,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, @@ -791,7 +810,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastVisionModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, # Auto-detect + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, @@ -824,7 +843,7 @@ class UnslothTrainer: self.model, self.tokenizer = FastLanguageModel.from_pretrained( model_name = model_name, max_seq_length = max_seq_length, - dtype = None, # Auto-detect + dtype = _auto_dtype, load_in_4bit = load_in_4bit, device_map = device_map, full_finetuning = full_finetuning, @@ -1188,7 +1207,6 @@ class UnslothTrainer: We patch at both instance AND class level for maximum reliability, and strip non-TransformersKwargs params that Unsloth/PEFT inject. """ - import types import torch import torch.nn as nn from transformers.models.csm.modeling_csm import ( @@ -1730,7 +1748,6 @@ class UnslothTrainer: logger.info("Freeing SNAC codec model from GPU...\n") snac_model.to("cpu") del snac_model - import gc gc.collect() torch.cuda.empty_cache() @@ -1754,13 +1771,10 @@ class UnslothTrainer: Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens), format as special-token text strings for SFTTrainer with dataset_text_field="text". """ - import sys import torch import numpy as np import torchaudio.transforms as T - import subprocess - device = "cuda" if torch.cuda.is_available() else "cpu" # The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo, @@ -1960,7 +1974,6 @@ class UnslothTrainer: audio_tokenizer.model.cpu() audio_tokenizer.feature_extractor.cpu() del audio_tokenizer - import gc gc.collect() torch.cuda.empty_cache() @@ -1989,7 +2002,6 @@ class UnslothTrainer: OuteTTS AudioProcessor for speaker representations, PromptProcessor for training prompts. Outputs text strings for SFTTrainer with dataset_text_field="text". """ - import sys import io import tempfile import torch @@ -2173,7 +2185,6 @@ class UnslothTrainer: del whisper_model del audio_processor del prompt_processor - import gc gc.collect() torch.cuda.empty_cache() diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 632b38d75a..a825321597 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -21,6 +21,9 @@ import shutil import sys import time import traceback +import gc +import re +import types import subprocess as _sp from pathlib import Path from typing import Any, Callable @@ -70,6 +73,58 @@ _TILELANG_INSTALL_TIMEOUT_S = 600 _TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11") _FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS" +# Module-level handle so the torch.library.Library registration survives past +# run_training_process() and is not garbage collected mid-run. +_WINDOWS_ROCM_GROUPED_MM_LIB = None + +# Worker subprocesses inherit the parent env but not the parent's +# os.add_dll_directory registrations. Replicate main.py's Windows ROCm DLL +# setup at module load so the first `import torch` can find amdhip64.dll even +# when HIP_PATH\bin is not on the system PATH. Handles retained at module +# scope so they are not garbage collected. +_ROCM_DLL_HANDLES: list = [] +if sys.platform == "win32": + + def _add_rocm_dll_dirs_worker() -> None: + _candidates: list[str] = [] + for _var in ("HIP_PATH", "ROCM_PATH"): + _val = os.environ.get(_var) + if _val: + _candidates.append(os.path.join(_val, "bin")) + _default_root = os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm" + ) + + def _ver_key(name: str) -> tuple: + # Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string. + parts = [] + for chunk in name.split("."): + try: + parts.append((0, int(chunk))) + except ValueError: + parts.append((1, chunk)) + return tuple(parts) + + try: + if os.path.isdir(_default_root): + for _ver in sorted( + os.listdir(_default_root), key = _ver_key, reverse = True + ): + _bin = os.path.join(_default_root, _ver, "bin") + if os.path.isdir(_bin): + _candidates.append(_bin) + except OSError: + pass + for _d in _candidates: + if os.path.isdir(_d): + try: + _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d)) + except (OSError, AttributeError): + pass + + _add_rocm_dll_dirs_worker() + del _add_rocm_dll_dirs_worker + def _model_wants_causal_conv1d(model_name: str) -> bool: name = model_name.lower() @@ -320,11 +375,21 @@ def _install_package_wheel_first( f"{snippet}", ) else: - logger.error( - "Failed to install %s from PyPI:\n%s", - display_name, - result.stdout, - ) + if sys.platform == "win32": + # No prebuilt wheel and no source build toolchain on Windows -- + # this is expected for packages like causal-conv1d. Log at + # info so users aren't alarmed by what looks like an error. + logger.info( + "%s is not available on Windows (no prebuilt wheel); skipping", + display_name, + ) + logger.debug("Install output:\n%s", result.stdout) + else: + logger.error( + "Failed to install %s from PyPI:\n%s", + display_name, + result.stdout, + ) return False if is_hip: @@ -337,6 +402,9 @@ def _install_package_wheel_first( def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None: if not _model_wants_causal_conv1d(model_name): return + if sys.platform == "win32": + logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping") + return _install_package_wheel_first( event_queue = event_queue, @@ -404,6 +472,11 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool: """Install pinned FLA + fla-core with --no-deps. Returns True iff importable post-call.""" if os.getenv(_FLA_SKIP_ENV) == "1": return False + if sys.platform == "win32": + logger.info( + "Skipping flash-linear-attention install: no prebuilt wheel for Windows" + ) + return False if sys.version_info < _FLA_MIN_PYTHON: logger.info( "Skipping flash-linear-attention install: requires Python >= %d.%d, have %s", @@ -483,10 +556,17 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool: return False if result.returncode != 0: - logger.warning( - "flash-linear-attention install failed (continuing on torch fallback):\n%s", - result.stdout, - ) + if sys.platform == "win32": + logger.info( + "flash-linear-attention not available on Windows (no prebuilt wheel); " + "continuing on torch fallback" + ) + logger.debug("Install output:\n%s", result.stdout) + else: + logger.warning( + "flash-linear-attention install failed (continuing on torch fallback):\n%s", + result.stdout, + ) _send_status( event_queue, "flash-linear-attention install failed; continuing without it", @@ -607,15 +687,61 @@ def _tilelang_importable() -> bool: def _torch_has_hip() -> bool: - """True iff torch is a ROCm build; `torch.version.hip` is the only reliable signal on x86_64 ROCm.""" + """True iff torch is a ROCm build. + + `torch.version.hip` covers official PyTorch ROCm wheels; AMD SDK / Radeon + wheels can leave it unset but still encode "rocm" in `torch.__version__`. + """ try: import torch as _torch - return getattr(_torch.version, "hip", None) is not None + return bool( + getattr(_torch.version, "hip", None) + or "rocm" in getattr(_torch, "__version__", "").lower() + ) except Exception: return False +def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: + """Classify a ROCm device as unified-memory (APU) or discrete. + + Returns ``(gcn_arch, is_unified)`` where: + - ``gcn_arch`` is the canonical arch string (e.g. ``"gfx1151"``) when a + known attribute is present, or ``""`` when all arch attrs are absent. + - ``is_unified`` is ``True`` for AMD APUs with a shared GPU/system-RAM pool + (gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower + ``set_per_process_memory_fraction`` cap to leave headroom for the OS. + + Classification priority: + 1. ``gcnArchName`` / variant spellings (stable, naming-independent). + 2. Device-name substring match as a last-resort fallback when all arch + attrs are absent (AMD SDK / Radeon wheels may not populate them): + - gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M`` + - gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395), + ``Radeon 8050S`` (cut-down SKU) + """ + gcn_arch = "" + for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"): + _v = (getattr(props, _attr, "") or "").split(":")[0].strip() + if _v: + gcn_arch = _v + break + + if gcn_arch: + return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"} + + # Arch attrs absent — fall back to device-name matching. + dev_lower = (getattr(props, "name", "") or "").lower() + is_unified = ( + "890m" in dev_lower + or "880m" in dev_lower + or "8060s" in dev_lower + or "8050s" in dev_lower + ) + return gcn_arch, is_unified + + def _tilelang_platform_supported() -> bool: """True iff a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch. @@ -881,6 +1007,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None: _ensure_tilelang_backend_unconditional(eq) def _causal_conv1d_install(eq: Any) -> bool: + if sys.platform == "win32": + logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping") + return False ok = _install_package_wheel_first( event_queue = eq, import_name = "causal_conv1d", @@ -1133,7 +1262,6 @@ def _run_mlx_training(event_queue, stop_queue, config): Mirrors the event_queue protocol so the parent process pump works unchanged. """ import time - import gc import math import threading import queue as _queue @@ -1893,6 +2021,340 @@ def run_training_process( 'Install for better performance: pip install "triton-windows<3.7"' ) + # ── 1d. Stub torchao on Windows ROCm ── + # Shared with the export worker; see core/_torchao_stub.py for the full + # rationale (torchao -> torch.distributed._functional_collectives crashes on + # Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm. + # Must run before any import of transformers / unsloth_zoo. + from core._torchao_stub import install_torchao_windows_rocm_stub + + install_torchao_windows_rocm_stub() + + # ── 1e. Ensure torch.distributed helper attrs are present ── + # Single-GPU training never initialises the process group, so these helpers + # are never called — but transformers/trl import them unconditionally. + _td_stubs = { + "is_initialized": lambda: False, + "is_available": lambda: False, + "is_torchelastic_launched": lambda: False, + "get_rank": lambda: 0, + "get_world_size": lambda: 1, + "barrier": lambda: None, + } + + try: + import torch.distributed as _td + + for _name, _stub in _td_stubs.items(): + if not hasattr(_td, _name): + setattr(_td, _name, _stub) + except Exception: + _td_mock = types.ModuleType("torch.distributed") + for _name, _stub in _td_stubs.items(): + setattr(_td_mock, _name, _stub) + sys.modules["torch.distributed"] = _td_mock + try: + import torch as _torch + + _torch.distributed = _td_mock + except Exception: + pass + + # ── 1f. Windows ROCm runtime patches ── + # torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm ≤ 7.12 Windows), + # causing 0xC0000005 (access violation) during training. + # + # Root cause: the JitDecomp autograd decomposition system (NOT torch.compile) + # dispatches _grouped_mm → _fused_adagrad_ → _grouped_mm HIP → null crash. + # TORCHDYNAMO_DISABLE=1 stops the compiler frontend but does NOT stop + # JitDecomp, so we must also override the CUDA dispatch key for _grouped_mm + # with a safe Python fallback. + # + # Fixed in AMD's wheel: torch==2.11.0+rocm7.13.0 — the 3-D batch and grouped + # (with offs) variants of _grouped_mm now have working HIP kernels on gfx1200. + # We gate the dispatch override on HIP < 7.13 so users on the fixed wheel get + # the real GPU kernel rather than our Python fallback. + # + # Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0. + # + # Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None, + # Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor + # offs: optional group-split offsets (MoE-style variable-size batches) + # + # torch is already in sys.modules from section 1e's `import torch.distributed`. + # Module-level _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past + # function return / mid-run GC. + global _WINDOWS_ROCM_GROUPED_MM_LIB + if sys.platform == "win32": + _torch_for_rocm = sys.modules.get("torch") + # Broad check: torch.version.hip OR "rocm" in torch.__version__. + # AMD SDK / Radeon Windows wheels do not always populate + # torch.version.hip; without the broad check the BNB version pin, + # dynamo-disable, and _grouped_mm fallback below silently skip + # (matches the torchao stub gate above and main.py). + _build_version_for_rocm = ( + getattr(_torch_for_rocm, "__version__", "").lower() + if _torch_for_rocm is not None + else "" + ) + _is_win_rocm_torch = bool( + _torch_for_rocm is not None + and ( + getattr(getattr(_torch_for_rocm, "version", None), "hip", None) + or "rocm" in _build_version_for_rocm + ) + ) + if _is_win_rocm_torch: + # Disable dynamo (belt-and-suspenders; JitDecomp patch below is the + # real fix, but keeping dynamo off avoids any other compile paths). + if "TORCHDYNAMO_DISABLE" not in os.environ: + os.environ["TORCHDYNAMO_DISABLE"] = "1" + logger.info("Windows ROCm: torch.compile (dynamo) disabled") + + # BNB auto-detects the HIP version from torch.version.hip and uses + # it to choose which DLL to load (e.g. "7.13" → rocm713.dll). + # AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its + # version suffix does not always match the torch HIP version (e.g. + # torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still + # ships rocm72.dll). We detect the actual DLL name from the installed + # package and override BNB's auto-detection. "72" is a safe fallback + # if detection fails. Callers may override by pre-setting the var. + if "BNB_ROCM_VERSION" not in os.environ: + _bnb_rocm_ver = None + try: + import glob as _glob + import importlib.util as _ilu + import re as _re + + _bnb_spec = _ilu.find_spec("bitsandbytes") + if _bnb_spec and _bnb_spec.submodule_search_locations: + _all_vers: list[str] = [] + for _pkg_dir in _bnb_spec.submodule_search_locations: + for _dll in _glob.glob( + os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll") + ): + _m = _re.search( + r"libbitsandbytes_rocm(\d+)\.dll", + os.path.basename(_dll), + ) + if _m: + _all_vers.append(_m.group(1)) + # Pick the highest numeric suffix so that e.g. "713" + # wins over "72" when both variants are present. + # Filesystem glob order is not guaranteed, so always + # sort rather than stopping at the first match. + if _all_vers: + _bnb_rocm_ver = max(_all_vers, key = lambda v: int(v)) + except Exception: + pass + _bnb_rocm_ver = _bnb_rocm_ver or "72" + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver + logger.info( + "Windows ROCm: set BNB_ROCM_VERSION=%s " + "(detected from installed BNB wheel; " + "overrides torch.version.hip auto-detection)", + _bnb_rocm_ver, + ) + + # Parse HIP version for the kernel-fix gate below. + # torch.version.hip can be "7.13.99004", "7.2.0", etc. + # AMD SDK / Radeon wheels may leave torch.version.hip unset and + # encode the ROCm version in torch.__version__ instead + # (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back + # to that string when version.hip is missing. + def _hip_ver_at_least(major: int, minor: int) -> bool: + _hip_str = getattr( + getattr(_torch_for_rocm, "version", None), "hip", None + ) + if not _hip_str: + # Try the standard "+rocmX.Y.Z" embedded version first + # (e.g. "2.11.0+rocm7.13.0"). + _ver_match = re.search(r"rocm(\d+)\.(\d+)", _build_version_for_rocm) + if _ver_match: + return ( + int(_ver_match.group(1)), + int(_ver_match.group(2)), + ) >= (major, minor) + # AMD SDK / Radeon Windows wheels encode the build as + # "+rocmsdk" (e.g. "2.9.0+rocmsdk20251116") with no + # explicit rocmX.Y component. The rocmsdk format was + # introduced after the gfx120X null-kernel fix landed in + # ROCm 7.13, so any wheel with this suffix is new enough to + # have working HIP kernels. Treat as >= 7.13 rather than + # falling back to False and installing the Python workaround + # on a wheel that doesn't need it. + if "rocmsdk" in _build_version_for_rocm: + logger.debug( + "Windows ROCm: AMD SDK wheel detected (%r); " + "assuming HIP >= %d.%d (rocmsdk wheels post-date " + "the gfx120X null-kernel fix)", + _build_version_for_rocm, + major, + minor, + ) + return True + return False + try: + _parts = [int(x) for x in str(_hip_str).split(".")[:2]] + if len(_parts) < 2: + logger.warning( + "Windows ROCm: torch.version.hip %r has fewer than " + "two components; cannot compare against %d.%d", + _hip_str, + major, + minor, + ) + return False + return (_parts[0], _parts[1]) >= (major, minor) + except ValueError: + logger.warning( + "Windows ROCm: could not parse torch.version.hip %r as " + "a version number; assuming HIP < %d.%d", + _hip_str, + major, + minor, + ) + return False + + # _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12, + # causing 0xC0000005. AMD fixed it in ROCm 7.13 (torch 2.11+). + # Only install the Python fallback on the affected versions so users + # on 7.13+ get the real GPU kernel for MoE workloads. + if not _hip_ver_at_least(7, 13): + try: + import warnings as _warnings + + _gm_lib = _torch_for_rocm.library.Library("aten", "IMPL") + + def _grouped_mm_safe_impl( + self, mat2, offs = None, bias = None, out_dtype = None + ): + """Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12).""" + _t = _torch_for_rocm + if offs is None: + # No offsets: behave like the real op, which + # accepts either (M, K) x (K, N) -> mm, or 3-D + # batched inputs -> bmm. Picking torch.mm + # unconditionally previously raised "self must be + # a matrix" on 3-D MoE workloads. + if self.dim() == 3 and mat2.dim() == 3: + result = _t.bmm(self.contiguous(), mat2.contiguous()) + elif self.dim() == 3 and mat2.dim() == 2: + # Broadcast 2-D mat2 across the batch dim. + result = _t.matmul(self.contiguous(), mat2.contiguous()) + elif self.dim() == 2 and mat2.dim() == 3: + # Broadcast 2-D self across batch via matmul semantics. + result = _t.matmul(self.contiguous(), mat2.contiguous()) + else: + result = _t.mm(self.contiguous(), mat2.contiguous()) + else: + # Grouped case: offs[i] is the exclusive end-row of + # group i in `self`; mat2 may be 3-D or 2-D. + offs_list = offs.tolist() + pieces = [] + prev = 0 + for idx, end in enumerate(offs_list): + end = int(end) + a_part = self[prev:end].contiguous() + if mat2.dim() == 3: + b_part = mat2[idx].contiguous() + else: + b_part = mat2.contiguous() + pieces.append(_t.mm(a_part, b_part)) + prev = end + # Include any trailing rows not covered by offs + if prev < self.shape[0]: + a_tail = self[prev:].contiguous() + b_tail = ( + mat2[-1].contiguous() + if mat2.dim() == 3 + else mat2.contiguous() + ) + pieces.append(_t.mm(a_tail, b_tail)) + result = ( + _t.cat(pieces, dim = 0) + if pieces + else _t.zeros( + 0, + mat2.shape[-1], + device = self.device, + dtype = self.dtype, + ) + ) + if bias is not None: + result = result + bias + if out_dtype is not None: + result = result.to(out_dtype) + elif result.dtype != self.dtype: + result = result.to(self.dtype) + return result + + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") + _gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA") + + _WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC + logger.info( + "Windows ROCm: patched _grouped_mm CUDA dispatch " + "(null HIP kernel on gfx1200, ROCm ≤ 7.12 — " + "bypassed with Python mm fallback)" + ) + except Exception as _patch_exc: + logger.warning( + "Windows ROCm: could not patch _grouped_mm — " + "training may crash with 0xC0000005: %s", + _patch_exc, + ) + else: + logger.info( + "Windows ROCm: HIP >= 7.13 — _grouped_mm kernel is functional, " + "skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)" + ) + + # ── 1g. ROCm OOM guard ── + # 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 Python exception. set_per_process_memory_fraction caps the + # HIP allocator so PyTorch raises OutOfMemoryError before hitting the + # hardware limit, giving the UI a clean error instead of a system freeze. + # Only applied on ROCm -- NVIDIA CUDA has a graceful OOM path and does + # not need this cap. + # Unified-memory APUs (gfx1150 Strix Point / gfx1151 Strix Halo) share GPU + # and system RAM in one pool: 0.90 of 128 GB starves the OS. Use 0.80 there. + # Primary classifier: gcnArchName from device properties — stable within a + # product family and naming-independent. AMD SDK / Radeon wheels may omit + # gcnArchName or expose it under a variant spelling, so we try several attr + # names then fall back to known device-name markers as a last resort. + # Non-fatal: silently skipped if torch is not importable. + if _hw.IS_ROCM: + try: + import torch as _torch_mem + + if _torch_mem.cuda.is_available(): + # Classify unified vs discrete via _rocm_classify_unified_memory. + # See that function's docstring for classification priority. + _props = _torch_mem.cuda.get_device_properties(0) + _dev_name = _props.name + _gcn_arch, _is_unified = _rocm_classify_unified_memory(_props) + if _is_unified and not _gcn_arch: + logger.debug( + "ROCm OOM guard: gcnArchName absent -- inferred " + "unified memory from device name %r; applying 0.80 cap", + _dev_name, + ) + _mem_fraction = 0.80 if _is_unified else 0.90 + _torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction) + logger.info( + "ROCm OOM guard: set_per_process_memory_fraction(%.2f) — " + "%s memory host (%s, %s)", + _mem_fraction, + "unified" if _is_unified else "discrete", + _dev_name, + _gcn_arch or "unknown arch", + ) + except Exception as _oom_guard_err: + logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err) + # ── 2. Now import ML libraries (fresh in this clean process) ── try: _send_status(event_queue, "Importing Unsloth...") @@ -2347,14 +2809,38 @@ def run_training_process( ) except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } + _exc_str = str(exc).lower() + _is_oom = ( + "out of memory" in _exc_str + or "hip out of memory" in _exc_str + or "cuda out of memory" in _exc_str + or type(exc).__name__ == "OutOfMemoryError" ) + if _is_oom: + _oom_msg = ( + "GPU ran out of VRAM during training.\n" + "To fix: reduce max_seq_length (e.g. 2048–4096), enable " + "gradient_checkpointing=True, lower per_device_train_batch_size, " + "or use a smaller model / higher quantization." + ) + logger.error("Training stopped: GPU OOM — %s", exc) + event_queue.put( + { + "type": "error", + "error": _oom_msg, + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + else: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) def _send_status(event_queue: Any, message: str) -> None: diff --git a/studio/backend/main.py b/studio/backend/main.py index fbf1c31de3..6b8ac438c0 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -12,6 +12,110 @@ from pathlib import Path as _Path # Suppress annoying C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" +# ── Windows AMD ROCm DLL injection ────────────────────────────────────────── +# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with +# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. +if sys.platform == "win32": + # Retained at module scope -- os.add_dll_directory returns a handle that + # removes the search-path entry when garbage collected. + _ROCM_DLL_HANDLES: list = [] + + def _add_rocm_dll_dirs() -> None: + candidates = [] + # 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer + for _var in ("HIP_PATH", "ROCM_PATH"): + _val = os.environ.get(_var) + if _val: + candidates.append(os.path.join(_val, "bin")) + # 2. Standard AMD installer location: C:\Program Files\AMD\ROCm\\bin + # Scan all installed versions, newest first. + _default_root = os.path.join( + os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm" + ) + + def _ver_key(name: str) -> tuple: + # Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string. + parts = [] + for chunk in name.split("."): + try: + parts.append((0, int(chunk))) + except ValueError: + parts.append((1, chunk)) + return tuple(parts) + + try: + if os.path.isdir(_default_root): + for _ver in sorted( + os.listdir(_default_root), key = _ver_key, reverse = True + ): + _bin = os.path.join(_default_root, _ver, "bin") + if os.path.isdir(_bin): + candidates.append(_bin) + except OSError: + pass + for _d in candidates: + if os.path.isdir(_d): + try: + _ROCM_DLL_HANDLES.append(os.add_dll_directory(_d)) + except (OSError, AttributeError): + pass + + _add_rocm_dll_dirs() + del _add_rocm_dll_dirs + + # ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─ + # bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm.dll + # where comes from torch.version.hip (e.g. "7.13..." → "713"). + # The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without + # this the server process crashes with "Configured ROCm binary not found". + # Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION + # before any import that pulls in bitsandbytes (mirrors worker.py logic). + # Gate on the rocm bnb DLL (the exact file this configures) or HIP_PATH/ + # ROCM_PATH, not on torch.version.hip: that needed importing torch on every + # Windows host (NVIDIA/CPU included), adding seconds to startup. Radeon + # wheels without HIP_PATH still ship the rocm bnb DLL, so they are covered. + if "BNB_ROCM_VERSION" not in os.environ: + import glob as _glob + import logging as _logging + + _hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH")) + _bnb_rocm_ver = None + _found_rocm_bnb = False + try: + import importlib.util as _ilu + + _bnb_spec = _ilu.find_spec("bitsandbytes") + # submodule_search_locations (not spec.origin) handles editable installs. + if _bnb_spec and _bnb_spec.submodule_search_locations: + import re as _re_bnb + + _all_vers_main: list[str] = [] + for _pkg_dir in _bnb_spec.submodule_search_locations: + for _dll in _glob.glob( + os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll") + ): + _found_rocm_bnb = True + _km = _re_bnb.search( + r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll) + ) + if _km: + _all_vers_main.append(_km.group(1)) + if _all_vers_main: + _bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v)) + except Exception as _e: + _logging.getLogger(__name__).warning( + "Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'", + _e, + ) + # rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72"). + if _found_rocm_bnb or _hip_env: + _bnb_rocm_ver_final = _bnb_rocm_ver or "72" + os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final + _logging.getLogger(__name__).info( + "Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)", + _bnb_rocm_ver_final, + ) + # Ensure backend dir is on sys.path so _platform_compat is importable when # main.py is launched directly (e.g. `uvicorn main:app`). _backend_dir = str(_Path(__file__).parent) @@ -762,8 +866,6 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes: @font-face downloads to fail silently. Stripping the attribute makes them regular same-origin fetches that work on any protocol. """ - import re as _re - html = html_bytes.decode("utf-8") html = _re.sub(r'\s+crossorigin(?:="[^"]*")?', "", html) return html.encode("utf-8") diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py new file mode 100644 index 0000000000..e0b819d54b --- /dev/null +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs +(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS.""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.llama_cpp import LlamaCppBackend + + +def _fake_torch(hip, archs, *, cuda_ok = True): + t = types.ModuleType("torch") + t.version = types.SimpleNamespace(hip = hip) + t.cuda = types.SimpleNamespace( + is_available = lambda: cuda_ok, + device_count = lambda: len(archs), + get_device_properties = lambda i: types.SimpleNamespace(gcnArchName = archs[i]), + ) + return t + + +@pytest.mark.parametrize( + "hip,archs,expected", + [ + ("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped) + ("6.2.0", ["gfx1150"], True), # Strix Point APU + ("6.2.0", ["gfx1100"], False), # discrete RDNA3 + ("6.2.0", ["gfx1201"], False), # discrete RDNA4 + ("6.2.0", ["gfx942"], False), # MI300X (data center) + (None, ["sm_90"], False), # NVIDIA (no torch.version.hip) + ("6.2.0", ["gfx1100", "gfx1151"], True), # mixed dGPU + APU + ], +) +def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected): + monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip, archs)) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected + + +def test_cpu_no_cuda_returns_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False)) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is False + + +def test_missing_torch_returns_false(monkeypatch): + monkeypatch.setitem(sys.modules, "torch", None) + assert LlamaCppBackend._amd_apu_wants_unified_memory() is False diff --git a/studio/backend/tests/test_gguf_routing.py b/studio/backend/tests/test_gguf_routing.py new file mode 100644 index 0000000000..1b299cba19 --- /dev/null +++ b/studio/backend/tests/test_gguf_routing.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Tests for GGUF routing in detect_gguf_model. + +Regression test for the bug where a .gguf file temporarily appears +inaccessible on Windows during llama-server process teardown, causing +is_file() to return False and the model to be routed to the transformers +backend instead of llama-server. +""" + +import sys +import os +import types +from pathlib import Path +from unittest.mock import patch + +# Stub structlog before importing backend modules (mirrors other tests in this suite) +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _): + return lambda *a, **k: None + + sys.modules["structlog"] = types.SimpleNamespace( + get_logger = lambda *a, **k: _DummyLogger(), + BoundLogger = _DummyLogger, + ) + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from utils.models.model_config import detect_gguf_model + + +def test_detects_gguf_file_normally(tmp_path): + """Normal case: .gguf file exists and is accessible.""" + gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf" + gguf.write_bytes(b"") + result = detect_gguf_model(str(gguf)) + assert result is not None + assert result.endswith("gpt-oss-20b-MXFP4.gguf") + + +def test_detects_gguf_when_stat_raises_oserror(tmp_path): + """ + Regression: on Windows, both is_file() and exists() call stat() internally. + During the brief lock window after llama-server is killed, stat() raises + OSError, causing both to return False. detect_gguf_model must still route + to llama-server based on the file extension alone. + """ + gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf" + gguf.write_bytes(b"") + + original_stat = Path.stat + + def flaky_stat(self, **kwargs): + if self.suffix.lower() == ".gguf": + raise OSError("file temporarily inaccessible (Windows lock window)") + return original_stat(self, **kwargs) + + with patch.object(Path, "stat", flaky_stat): + result = detect_gguf_model(str(gguf)) + + assert result is not None, ( + "detect_gguf_model returned None when stat() raised OSError. " + "This causes the model to fall through to the transformers backend." + ) + + +def test_does_not_detect_mmproj_as_main_model(tmp_path): + """mmproj files must never be returned as the primary model.""" + mmproj = tmp_path / "mmproj-model-f16.gguf" + mmproj.write_bytes(b"") + result = detect_gguf_model(str(mmproj)) + assert result is None + + +def test_detects_gguf_in_directory(tmp_path): + """Directory containing a .gguf file is resolved to that file.""" + gguf = tmp_path / "model-Q4_K_M.gguf" + gguf.write_bytes(b"") + result = detect_gguf_model(str(tmp_path)) + assert result is not None + assert result.endswith("model-Q4_K_M.gguf") + + +def test_directory_named_like_gguf_scans_inside(tmp_path): + """A directory named *.gguf resolves the real .gguf inside, not itself.""" + gguf_dir = tmp_path / "mymodel.gguf" + gguf_dir.mkdir() + inner = gguf_dir / "model-Q4_K_M.gguf" + inner.write_bytes(b"") + result = detect_gguf_model(str(gguf_dir)) + assert result is not None + assert result.endswith("model-Q4_K_M.gguf") + + +def test_returns_none_for_non_gguf_path(tmp_path): + """Non-.gguf paths with no .gguf files inside return None.""" + result = detect_gguf_model(str(tmp_path)) + assert result is None diff --git a/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py new file mode 100644 index 0000000000..5d2d672890 --- /dev/null +++ b/studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py @@ -0,0 +1,430 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets. + +Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade +GitHub API are stubbed out so the suite runs without internet access and is +not subject to rate limits. +""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +_studio = Path(__file__).resolve().parent.parent.parent +if str(_studio) not in sys.path: + sys.path.insert(0, str(_studio)) + +_mod = importlib.import_module("install_llama_prebuilt") +HostInfo = _mod.HostInfo +resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None) +_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None) + +if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None: + pytest.skip("PR symbols not present - check branch", allow_module_level = True) + + +@pytest.fixture(autouse = True) +def _clear_lemonade_release_cache(): + """Prevent cross-test pollution of the lemonade release lru_cache when + future tests vary the fetch_json mock return value.""" + _cache = getattr(_mod, "_fetch_lemonade_release_cached", None) + if _cache is not None and hasattr(_cache, "cache_clear"): + _cache.cache_clear() + yield + if _cache is not None and hasattr(_cache, "cache_clear"): + _cache.cache_clear() + + +_STUB_TAG = "b1262" +_STUB_OS_PREFIXES = ("ubuntu", "windows") +_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X") + + +def _stub_lemonade_release() -> dict: + """Minimal lemonade release payload covering all supported GPU/OS combinations.""" + assets = [ + { + "name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip", + "browser_download_url": ( + f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/" + f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip" + ), + } + for prefix in _STUB_OS_PREFIXES + for family in _STUB_FAMILIES + ] + return {"tag_name": _STUB_TAG, "assets": assets} + + +def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo: + return HostInfo( + system = "Windows" if windows else "Linux", + machine = "amd64" if windows else "x86_64", + is_windows = windows, + is_linux = not windows, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = True, + rocm_gfx_target = gfx_target, + ) + + +def _lookup_family(gfx: str) -> str | None: + for prefix, family in _LEMONADE_GFX_FAMILIES: + if gfx.startswith(prefix): + return family + return None + + +# --------------------------------------------------------------------------- +# GPU family mapping +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "gfx,expected_family", + [ + ("gfx1151", "gfx1151"), + ("gfx1150", "gfx1150"), + ("gfx1201", "gfx120X"), + ("gfx1200", "gfx120X"), + ("gfx1100", "gfx110X"), + ("gfx1030", "gfx103X"), + ], +) +def test_gpu_family_mapping(gfx, expected_family): + assert _lookup_family(gfx) == expected_family + + +def test_unknown_gpu_not_in_families(): + assert _lookup_family("gfx999") is None + + +# --------------------------------------------------------------------------- +# Asset resolution - hits real lemonade GitHub API +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "gfx,os_prefix,windows", + [ + ("gfx1151", "ubuntu", False), + ("gfx1150", "ubuntu", False), + ("gfx1201", "ubuntu", False), + ("gfx1100", "ubuntu", False), + ("gfx1030", "ubuntu", False), + ("gfx1151", "windows", True), + ("gfx1100", "windows", True), + ], +) +def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows): + host = _make_rocm_host(gfx, windows = windows) + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + result = resolve_lemonade_rocm_choice( + host, os_prefix, "default", llama_tag = "latest" + ) + assert ( + result is not None + ), f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})" + assert _lookup_family(gfx) in result.name + assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm") + + +def test_unknown_gpu_falls_through_to_upstream(): + host = _make_rocm_host("gfx999") + result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest") + assert result is None + + +# --------------------------------------------------------------------------- +# Simple-policy dispatcher must plan a lemonade ROCm attempt for AMD-only hosts. +# This is the path setup.sh actually invokes (via --simple-policy), so the +# lemonade integration is useless if it isn't wired in here. +# --------------------------------------------------------------------------- + +direct_linux_release_plan = getattr(_mod, "direct_linux_release_plan", None) +direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None) + + +def _stub_unsloth_release(release_tag: str = "b9022") -> dict: + # Minimal payload that parse_direct_linux_release_bundle accepts. It + # requires at least one `app-{label}-linux-x64*.tar.gz` asset for the + # bundle to be recognised; we ship a bare CPU one so the planner has a + # baseline non-ROCm attempt to fall through to. + asset_name = f"app-{release_tag}-linux-x64.tar.gz" + return { + "tag_name": release_tag, + "name": release_tag, + "assets": [ + { + "name": asset_name, + "browser_download_url": f"https://example.invalid/{asset_name}", + }, + ], + } + + +@pytest.mark.skipif( + direct_linux_release_plan is None, + reason = "simple-policy dispatcher not present on this branch", +) +def test_simple_policy_plans_lemonade_for_rocm_host(): + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + plan = direct_linux_release_plan( + _stub_unsloth_release(), + host, + "unslothai/llama.cpp", + "latest", + ) + assert plan is not None, "ROCm host should not be skipped by simple-policy planner" + kinds = [a.install_kind for a in plan.attempts] + assert ( + "linux-rocm" in kinds + ), f"simple-policy planner did not include a lemonade ROCm attempt; got {kinds}" + rocm_attempt = next(a for a in plan.attempts if a.install_kind == "linux-rocm") + assert rocm_attempt.source_label == "lemonade" + assert "gfx1151" in rocm_attempt.name + + +@pytest.mark.skipif( + direct_upstream_release_plan is None, + reason = "simple-policy dispatcher not present on this branch", +) +def test_simple_policy_plans_lemonade_for_windows_hip_host(): + host = _make_rocm_host("gfx1151", windows = True) + release = { + "tag_name": "b9022", + "name": "b9022", + "assets": [], + } + with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()): + plan = direct_upstream_release_plan( + release, host, "ggml-org/llama.cpp", "latest" + ) + assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt" + kinds = [a.install_kind for a in plan.attempts] + assert ( + "windows-hip" in kinds + ), f"simple-policy planner did not include a lemonade HIP attempt; got {kinds}" + + +@pytest.mark.skipif( + direct_upstream_release_plan is None, + reason = "simple-policy dispatcher not present on this branch", +) +def test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavailable(): + """If lemonade returns None (e.g. gfx999 or transient API failure), the planner + must still include the upstream HIP asset rather than silently downgrading to CPU.""" + host = _make_rocm_host("gfx999", windows = True) + hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip" + release = { + "tag_name": "b9022", + "name": "b9022", + "assets": [ + { + "name": hip_asset, + "browser_download_url": f"https://example.invalid/{hip_asset}", + }, + ], + } + plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest") + assert plan is not None + kinds = [a.install_kind for a in plan.attempts] + assert ( + "windows-hip" in kinds + ), f"upstream HIP asset not included as fallback; got {kinds}" + hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip") + assert hip_attempt.source_label == "upstream" + + +# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ── + + +def test_lemonade_release_api_url_pinned_tag(): + """A pinned llama_tag must produce the /releases/tags/ URL.""" + assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262") + assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest") + assert _mod._lemonade_release_api_for("").endswith("/releases/latest") + + +def test_lemonade_release_api_url_encodes_tag(): + """Unexpected slashes / hashes in the tag must be URL-encoded so the URL + cannot be reshaped (defence in depth -- tags should already be sanitised + upstream).""" + url = _mod._lemonade_release_api_for("b1260/../latest") + assert "/releases/tags/b1260%2F..%2Flatest" in url + assert "//latest" not in url.split("/releases/tags/", 1)[1] + + +def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch): + """UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver.""" + monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1") + host = _make_rocm_host("gfx1151") + res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest") + assert res is None + + +def test_lemonade_resolver_rejects_non_github_url(monkeypatch): + """If the GitHub API response somehow contained an off-host download URL, + the resolver must refuse to use it (lemonade assets are not in the + approved-hash manifest).""" + bad_release = { + "tag_name": _STUB_TAG, + "assets": [ + { + "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", + "browser_download_url": "https://attacker.invalid/llama.zip", + }, + ], + } + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = bad_release): + res = resolve_lemonade_rocm_choice( + host, "ubuntu", "linux-rocm", llama_tag = "latest" + ) + assert res is None + + +def test_lemonade_resolver_rejects_http_scheme(): + assert not _mod._is_trusted_github_release_url( + "http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_accepts_github_cdn(): + # Real GitHub release CDN URLs carry the /github-production-release-asset- prefix. + assert _mod._is_trusted_github_release_url( + "https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_rejects_arbitrary_cdn_path(): + # A CDN URL without the release-asset path prefix must be rejected. + assert not _mod._is_trusted_github_release_url( + "https://objects.githubusercontent.com/abc/def", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_accepts_release_path(): + url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip" + assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm") + + +def test_lemonade_resolver_rejects_wrong_repo(): + """A github.com release URL for a different repo must be rejected.""" + assert not _mod._is_trusted_github_release_url( + "https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip", + "lemonade-sdk/llamacpp-rocm", + ) + + +def test_lemonade_resolver_rejects_empty_browser_download_url(): + """An asset entry with an empty browser_download_url must fall through.""" + release = { + "tag_name": _STUB_TAG, + "assets": [ + { + "name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip", + "browser_download_url": "", + }, + ], + } + host = _make_rocm_host("gfx1151") + with patch.object(_mod, "fetch_json", return_value = release): + res = resolve_lemonade_rocm_choice( + host, "ubuntu", "linux-rocm", llama_tag = "latest" + ) + assert res is None + + +def test_lemonade_runtime_patterns_include_hip_runtime(): + """linux-rocm overlay must use a broad lib glob to catch all bundled .so files. + + Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp, + ...) whose names change across ROCm releases. A broad ``lib*.so*`` glob + avoids having to enumerate every transitive dependency by name. + """ + from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice + + choice = AssetChoice( + repo = "lemonade-sdk/llamacpp-rocm", + tag = "b1262", + name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip", + url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip", + source_label = "lemonade", + install_kind = "linux-rocm", + ) + pats = runtime_patterns_for_choice(choice) + # The broad glob must be present so every .so in the lemonade bundle + # (including transitive deps added in future ROCm releases) gets overlaid. + assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}" + + +_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None) + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch): + """AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES; + on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100.""" + # Two GPUs; rocminfo reports each token twice (as in the real tool output). + probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1") + assert _pick_rocm_gfx_target(probe_out) == "gfx1100" + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch): + """CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None.""" + probe_out = "gfx1151\ngfx1100" + monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1") + assert _pick_rocm_gfx_target(probe_out) is None + + +@pytest.mark.skipif( + _pick_rocm_gfx_target is None, + reason = "_pick_rocm_gfx_target not present on this branch", +) +def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch): + """Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must + return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the + two gfx1100 entries into one and making index 2 out of range.""" + # Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU). + # Each GPU gets its own Agent section with a few token mentions. + probe_out = ( + "***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n" + "***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n" + ) + monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False) + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2") + assert _pick_rocm_gfx_target(probe_out) == "gfx1151" diff --git a/studio/backend/tests/test_log_filter_no_truncation.py b/studio/backend/tests/test_log_filter_no_truncation.py index d78643f5b9..d9a6e2bc4a 100644 --- a/studio/backend/tests/test_log_filter_no_truncation.py +++ b/studio/backend/tests/test_log_filter_no_truncation.py @@ -2,27 +2,11 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Regression tests for studio.backend.loggers.handlers.filter_sensitive_data. +Regression tests for loggers.handlers.filter_sensitive_data. -Context: filter_sensitive_data was originally written with a base64-detection -heuristic that truncated any string >100 chars containing ',' or '/' down to -20 chars + '...'. The block was dormant until PR #5246 wired the processor -into the structlog chain to redact native-path leases. Once active, the -heuristic ate normal log lines emitted by llama_cpp_backend (GGUF size -summary, mmproj selection, the full llama-server command line) and any -exception traceback that happened to contain a file path. - -These tests pin two properties: - -1. Long, comma- or slash-bearing log messages flow through filter_sensitive_data - unchanged. The exact strings exercised match the call sites at - studio/backend/core/inference/llama_cpp.py:2117, :2283, and :2312 that - were truncated in the original bug report. - -2. PR #5246's native-path lease redaction still fires for both the inline - ``native_path_lease=...`` regex form and the ``nativePathLease`` dict-key - form. This guards against future regressions that strip redaction along - with the truncation block. +Pins two properties: (1) long strings with commas/slashes pass through +unchanged (the base64-truncation heuristic from PR #5246 was too aggressive), +and (2) native-path lease redaction still fires for both inline and dict-key forms. """ from loggers.handlers import filter_sensitive_data diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py new file mode 100644 index 0000000000..2ce9b55789 --- /dev/null +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for _rocm_classify_unified_memory (ROCm OOM-guard classifier). + +Covers the three classification paths: + Path 1 – canonical gcnArchName attribute present. + Path 2 – gcnArchName absent, alternate-spelling attribute present. + Path 3 – ALL arch attrs absent; falls back to device-name substring match. + +Regression for: Strix Halo (gfx1151) misclassified as discrete on AMD SDK / +Radeon wheels that populate props.name = "Radeon 8060S Graphics" but do NOT +set any gcnArchName attribute. Without the 8060s/8050s name patterns the +fallback returned is_unified=False, applying the 0.90 fraction instead of +0.80 and leaving only ~12.8 GiB OS headroom on a 128 GiB unified-memory pool. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from core.training.worker import _rocm_classify_unified_memory + + +# ── helpers ────────────────────────────────────────────────────────────────── + + +def _props(**kwargs) -> SimpleNamespace: + """Build a fake device-properties object with the given attributes.""" + return SimpleNamespace(**kwargs) + + +# ── Path 1: canonical gcnArchName ──────────────────────────────────────────── + + +class TestCanonicalGcnArchName: + """gcnArchName is present and populated.""" + + @pytest.mark.parametrize( + "arch, expected_unified", + [ + ("gfx1150", True), # Strix Point + ("gfx1151", True), # Strix Halo + ("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete + ("gfx906", False), # MI50 — discrete server GPU + ("gfx1201", False), # RX 9070 XT — discrete + ], + ) + def test_canonical_attr(self, arch: str, expected_unified: bool) -> None: + props = _props(gcnArchName = arch, name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == arch + assert is_unified is expected_unified + + def test_arch_with_colon_suffix_stripped(self) -> None: + """gcnArchName can carry xnack/sramecc suffix; only the base is kept.""" + props = _props(gcnArchName = "gfx1151:xnack-", name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + def test_canonical_attr_wins_over_name(self) -> None: + """Arch attr takes priority; device name should be ignored.""" + # Discrete arch, but name looks like a unified SKU — arch must win. + props = _props(gcnArchName = "gfx1100", name = "Radeon 890M") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1100" + assert is_unified is False + + +# ── Path 2: alternate-spelling fallback ────────────────────────────────────── + + +class TestAlternateSpellingFallback: + """gcnArchName is missing but an alternate attr spelling is present.""" + + @pytest.mark.parametrize( + "attr_name", + ["gcn_arch_name", "arch_name", "gfx_arch_name"], + ) + def test_alternate_attr_unified(self, attr_name: str) -> None: + props = _props(**{attr_name: "gfx1151"}, name = "Radeon 8060S Graphics") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + @pytest.mark.parametrize( + "attr_name", + ["gcn_arch_name", "arch_name", "gfx_arch_name"], + ) + def test_alternate_attr_discrete(self, attr_name: str) -> None: + props = _props(**{attr_name: "gfx1201"}, name = "Radeon RX 9070 XT") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1201" + assert is_unified is False + + def test_first_non_empty_attr_wins(self) -> None: + """When multiple alternate attrs are present the first non-empty one wins.""" + props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant") + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "gfx1151" + assert is_unified is True + + +# ── Path 3: device-name fallback ───────────────────────────────────────────── + + +class TestDeviceNameFallback: + """ALL arch attrs absent — classifier must rely solely on device name.""" + + # --- unified-memory devices that MUST be detected --- + + @pytest.mark.parametrize( + "device_name", + [ + # gfx1150 Strix Point + "Radeon 890M", + "AMD Radeon 890M Graphics", + "RADEON 890M", # case-insensitive + "Radeon 880M", + "AMD Radeon 880M Graphics", + # gfx1151 Strix Halo — the regression case from the review + "Radeon 8060S Graphics", # Ryzen AI MAX+ 395 (as returned by torch) + "AMD Radeon 8060S", + "Radeon 8050S Graphics", # cut-down Strix Halo SKU + "AMD Radeon 8050S", + # case variants + "RADEON 8060S GRAPHICS", + "radeon 8050s", + ], + ) + def test_unified_memory_detected(self, device_name: str) -> None: + props = _props(name = device_name) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "", f"expected empty gcn_arch, got {gcn!r}" + assert ( + is_unified is True + ), f"device {device_name!r} should be classified as unified-memory" + + # --- discrete devices that must NOT be mis-classified --- + + @pytest.mark.parametrize( + "device_name", + [ + "Radeon RX 9070 XT", + "AMD Radeon RX 7900 XTX", + "Radeon RX 6900 XT", + "Radeon Pro W7900", + "AMD Instinct MI300X", + # Names that contain superficially similar substrings but are discrete + "Radeon RX 580", + "Radeon VII", + ], + ) + def test_discrete_not_misclassified(self, device_name: str) -> None: + props = _props(name = device_name) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert ( + is_unified is False + ), f"discrete device {device_name!r} should NOT be classified as unified-memory" + + def test_empty_name_returns_false(self) -> None: + """Completely absent name must not crash and must default to discrete.""" + props = _props() # no 'name' attr at all + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is False + + def test_none_name_returns_false(self) -> None: + props = _props(name = None) + gcn, is_unified = _rocm_classify_unified_memory(props) + assert gcn == "" + assert is_unified is False diff --git a/studio/backend/utils/hardware/amd.py b/studio/backend/utils/hardware/amd.py index fdb1ab4520..48d5890399 100644 --- a/studio/backend/utils/hardware/amd.py +++ b/studio/backend/utils/hardware/amd.py @@ -11,18 +11,35 @@ nvidia.py counterparts. import json import math import os +import platform import re import subprocess +import sys from typing import Any, Optional from loggers import get_logger from utils.native_path_leases import child_env_without_native_path_secret +from utils.subprocess_compat import windows_hidden_subprocess_kwargs logger = get_logger(__name__) +# amd-smi on Windows must initialise the full ROCm runtime on first call, which +# can take 15-25 s on cold hardware. Linux is consistently < 2 s. +_AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10 -def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: +# Circuit breaker: stop calling amd-smi after this many consecutive failures. +# On Windows, each failed call spawns a process that may show a UAC/DiskPart +# elevation prompt. Once we know amd-smi doesn't work we stop polling it. +_AMD_SMI_FAILURE_LIMIT = 3 +_amd_smi_consecutive_failures = 0 +_amd_smi_disabled = False + + +def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]: """Run amd-smi with the given arguments and return parsed JSON, or None.""" + global _amd_smi_consecutive_failures, _amd_smi_disabled + if _amd_smi_disabled: + return None try: result = subprocess.run( ["amd-smi", *args, "--json"], @@ -30,13 +47,40 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]: text = True, timeout = timeout, env = child_env_without_native_path_secret(), + **windows_hidden_subprocess_kwargs(), ) except (OSError, subprocess.TimeoutExpired) as e: - logger.warning("amd-smi query failed: %s", e) + if isinstance(e, FileNotFoundError): + # amd-smi ships with Adrenalin, not the HIP SDK -- absence is + # expected on HIP SDK-only Windows setups. Log at debug only. + logger.debug("amd-smi not found (not in PATH): %s", e) + else: + logger.warning("amd-smi query failed: %s", e) + _amd_smi_consecutive_failures += 1 + if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT: + logger.info( + "amd-smi not available (not installed; expected on HIP SDK-only systems); " + "GPU VRAM polling disabled" + ) + _amd_smi_disabled = True return None - if result.returncode != 0 or not result.stdout.strip(): + if result.returncode != 0: logger.warning("amd-smi returned code %d", result.returncode) + _amd_smi_consecutive_failures += 1 + if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT: + logger.info( + "amd-smi not available (not installed; expected on HIP SDK-only systems); " + "GPU VRAM polling disabled" + ) + _amd_smi_disabled = True return None + if not result.stdout.strip(): + # amd-smi exited successfully but produced no output (e.g. no GPUs + # visible on this query, or a version that emits nothing for --json). + # This is not a tool failure, so don't count against the circuit breaker. + logger.debug("amd-smi exited 0 but returned no output") + return None + _amd_smi_consecutive_failures = 0 # reset on success try: return json.loads(result.stdout) except json.JSONDecodeError: @@ -352,7 +396,7 @@ def get_visible_gpu_utilization( ) parsed_id = _parse_numeric(raw_id) if parsed_id is None: - logger.debug( + logger.warning( "amd-smi GPU id %r could not be parsed; falling back to " "enumeration index %d", raw_id, @@ -360,7 +404,15 @@ def get_visible_gpu_utilization( ) idx = fallback_idx else: - idx = int(parsed_id) + rounded = round(parsed_id) + if rounded != parsed_id: + logger.warning( + "amd-smi GPU id %r parsed as non-integer %r; truncating to %d", + raw_id, + parsed_id, + rounded, + ) + idx = int(rounded) if idx not in visible_set: continue metrics = _extract_gpu_metrics(gpu_data) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index ede37e2953..180fde8f13 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -16,8 +16,16 @@ Usage: ... """ +import copy +import gc +import glob import os import platform +import re +import subprocess +import sys +import types +from importlib.metadata import PackageNotFoundError, version as pkg_version import structlog from loggers import get_logger from enum import Enum @@ -120,11 +128,13 @@ def detect_hardware() -> DeviceType: # Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes. # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP. - if getattr(torch.version, "hip", None) is not None: + # AMD's repo.radeon.com SDK wheels (e.g. 2.9.0+rocmsdk20251116) do + # not set torch.version.hip, so fall back to checking __version__. + _hip_ver = getattr(torch.version, "hip", None) + if _hip_ver is not None or "rocm" in torch.__version__.lower(): IS_ROCM = True - print( - f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}" - ) + _hip_label = _hip_ver or torch.__version__ + print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}") else: print(f"Hardware detected: CUDA -- {device_name}") return DEVICE @@ -176,8 +186,6 @@ def clear_gpu_cache(): Clear GPU memory cache for the current device. Safe to call on any platform — no-ops gracefully. """ - import gc - gc.collect() device = get_device() @@ -359,8 +367,6 @@ def get_package_versions() -> Dict[str, Optional[str]]: Returns dict with keys: unsloth, torch, transformers, cuda. Missing packages yield None. """ - from importlib.metadata import version as pkg_version, PackageNotFoundError - packages = ("unsloth", "torch", "transformers") versions: Dict[str, Optional[str]] = {} @@ -466,7 +472,7 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: try: func = getattr(_backend, func_name) result = func(*args, **kwargs) - if result.get("available"): + if isinstance(result, dict) and result.get("available"): return result except Exception as e: logger.warning("%s %s query failed: %s", backend_name, func_name, e) @@ -479,9 +485,6 @@ def _read_apple_gpu_stats() -> Dict[str, Any]: Returns dict with utilization_pct, vram_used_bytes (system-wide GPU memory). Returns empty dict on failure. """ - import subprocess - import re - try: result = subprocess.run( ["ioreg", "-r", "-c", "AGXAccelerator"], @@ -506,6 +509,133 @@ def _read_apple_gpu_stats() -> Dict[str, Any]: } +def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]: + """Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent.""" + if platform.system() != "Linux": + return None + try: + files = glob.glob("/sys/class/drm/card*/device/gpu_busy_percent") + if not files: + return None + values = [int(open(f).read().strip()) for f in files] + return round(sum(values) / len(values), 1) + except Exception: + return None + + +def _rocm_linux_sysfs_temp_c() -> Optional[float]: + """Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C).""" + if platform.system() != "Linux": + return None + try: + files = glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input") + if not files: + return None + temps = [int(open(f).read().strip()) / 1000.0 for f in files] + return round(max(temps), 1) + except Exception: + return None + + +def _rocm_linux_sysfs_power_w() -> Optional[float]: + """Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts).""" + if platform.system() != "Linux": + return None + try: + for pattern in ( + "/sys/class/drm/card*/device/hwmon/hwmon*/power1_average", + "/sys/class/drm/card*/device/hwmon/hwmon*/power1_input", + ): + files = glob.glob(pattern) + if files: + watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files) + return round(watts, 1) + return None + except Exception: + return None + + +def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]: + """Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes).""" + if platform.system() != "Windows": + return None + try: + ps = ( + "$s=(Get-Counter '\\GPU Engine(*engtype_3D*)\\Utilization Percentage'" + " -ErrorAction SilentlyContinue).CounterSamples;" + "if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}" + ) + r = subprocess.run( + ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], + capture_output = True, + text = True, + timeout = 5, + ) + if r.returncode != 0 or not r.stdout.strip(): + return None + val = float(r.stdout.strip()) + return round(val, 1) if val >= 0 else None + except Exception: + return None + + +def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]: + """Query system-wide AMD GPU VRAM via Linux DRM sysfs. + + Reads /sys/class/drm/card*/device/mem_info_vram_* which the kernel + updates in real-time across all processes. No tools required. + Returns (used_gb, total_gb) or (None, None) on failure. + """ + if platform.system() != "Linux": + return None, None + try: + used_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_used") + total_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_total") + if not used_files or not total_files: + return None, None + used_bytes = sum(int(open(f).read().strip()) for f in used_files) + total_bytes = sum(int(open(f).read().strip()) for f in total_files) + if total_bytes == 0: + return None, None + return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + except Exception: + return None, None + + +def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]: + """Query system-wide dedicated GPU VRAM via Windows Performance Counters. + + Uses the same data source as Task Manager so it reflects cross-process + usage accurately. Works for any GPU vendor without amd-smi or nvidia-smi. + Returns (used_gb, total_gb) or (None, None) on failure. + """ + if platform.system() != "Windows": + return None, None + try: + ps = ( + "$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'" + " -ErrorAction SilentlyContinue).CounterSamples;" + "if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}" + ) + r = subprocess.run( + ["powershell", "-NoProfile", "-NonInteractive", "-Command", ps], + capture_output = True, + text = True, + timeout = 5, + ) + if r.returncode != 0 or not r.stdout.strip(): + return None, None + used_bytes = float(r.stdout.strip()) + if used_bytes < 0: + return None, None + import torch as _torch + + total_bytes = _torch.cuda.get_device_properties(0).total_memory + return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2) + except Exception: + return None, None + + def get_gpu_utilization() -> Dict[str, Any]: """Return a live snapshot of device utilization information.""" device = get_device() @@ -514,7 +644,78 @@ def get_gpu_utilization() -> Dict[str, Any]: result = _smi_query("get_primary_gpu_utilization") if result is not None: result["backend"] = _backend_label(device) + if IS_ROCM: + # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.) + _reconcile_primary_rocm_unified_memory( + result, _get_parent_visible_gpu_spec() + ) return result + # SMI tool unavailable or returned no usable data. On Windows, query + # the Performance Counter API (same source as Task Manager) for + # system-wide dedicated VRAM — covers cross-process usage that + # torch.cuda.mem_get_info cannot see from the Studio server process. + if IS_ROCM and platform.system() == "Windows": + _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() + if _win_used is not None and _win_total is not None: + _win_util = _rocm_windows_perf_counter_gpu_util_pct() + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": _win_util, + "temperature_c": None, + "vram_used_gb": _win_used, + "vram_total_gb": _win_total, + "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) + if _win_total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed. + if IS_ROCM and platform.system() == "Linux": + _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb() + if _linux_used is not None and _linux_total is not None: + _linux_util = _rocm_linux_sysfs_gpu_busy_pct() + _linux_temp = _rocm_linux_sysfs_temp_c() + _linux_power = _rocm_linux_sysfs_power_w() + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": _linux_util, + "temperature_c": _linux_temp, + "vram_used_gb": _linux_used, + "vram_total_gb": _linux_total, + "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) + if _linux_total > 0 + else None, + "power_draw_w": _linux_power, + "power_limit_w": None, + "power_utilization_pct": None, + } + # Last resort: torch mem_get_info (process-local). + _visible_spec = _get_parent_visible_gpu_spec() + _numeric_ids = _visible_spec.get("numeric_ids") or [0] + _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0] + _torch_devices = _torch_get_per_device_info(_primary_idx) + if _torch_devices: + _td = _torch_devices[0] + _total = _td["total_gb"] + _used = _td["used_gb"] + return { + "available": True, + "backend": _backend_label(device), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": _used, + "vram_total_gb": _total, + "vram_utilization_pct": round((_used / _total) * 100, 1) + if _total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } # MLX path: single _read_apple_gpu_stats() call carries both VRAM-used # bytes and GPU utilization %. psutil for unified-memory total is cheap. @@ -578,6 +779,77 @@ def get_gpu_utilization() -> Dict[str, Any]: return {"available": False, "backend": _backend_label(device)} +def _apply_unified_memory_correction( + device_metrics: Dict[str, Any], torch_info: Dict[str, Any] +) -> None: + """Per-device reconciliation: when torch reports a larger memory total + than amd-smi, overwrite the smi VRAM fields in place. + + Used by both the multi-device and primary-device reconciliation helpers + so the two endpoints stay in sync on AMD iGPUs with unified memory. + """ + torch_total_gb = torch_info["total_gb"] + smi_total_gb = device_metrics.get("vram_total_gb") or 0.0 + if torch_total_gb > smi_total_gb: + torch_used_gb = torch_info["used_gb"] + device_metrics["vram_total_gb"] = torch_total_gb + device_metrics["vram_used_gb"] = torch_used_gb + device_metrics["vram_utilization_pct"] = ( + round((torch_used_gb / torch_total_gb) * 100, 1) + if torch_total_gb > 0 + else None + ) + logger.debug( + "ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with " + "torch mem_get_info total (%.2f GB) for device %s", + smi_total_gb, + torch_total_gb, + torch_info.get("index"), + ) + + +def _reconcile_rocm_unified_memory( + utilization: Dict[str, Any], device_indices: list[int] +) -> None: + """Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo). + + amd-smi reports only the dedicated slice (~512 MB); torch sees the full + GTT pool (~128 GB). When torch total > smi total, overwrite per-device + VRAM fields so GPU selection uses the real available memory. + """ + torch_devices = _torch_get_per_device_info(device_indices) + if not torch_devices: + return + torch_by_index = {td["index"]: td for td in torch_devices} + for dev in utilization.get("devices", []): + td = torch_by_index.get(dev.get("index")) + if td is None: + continue + _apply_unified_memory_correction(dev, td) + + +def _reconcile_primary_rocm_unified_memory( + utilization: Dict[str, Any], parent_visible_spec: Dict[str, Any] +) -> None: + """Same fix as _reconcile_rocm_unified_memory for the flat primary-GPU dict.""" + numeric_ids = parent_visible_spec.get("numeric_ids") + if numeric_ids is None: + # No visibility env var set: torch ordinal 0 is the primary device. + primary_idx = [0] + elif len(numeric_ids) == 0: + # Empty mask (HIP_VISIBLE_DEVICES="" or "-1"): no GPU is visible to + # this process. Querying torch device 0 would raise a RuntimeError or + # return stale/wrong data, so bail out rather than writing bad values + # into the utilization dict. + return + else: + primary_idx = [int(numeric_ids[0])] + torch_devices = _torch_get_per_device_info(primary_idx) + if not torch_devices: + return + _apply_unified_memory_correction(utilization, torch_devices[0]) + + def get_visible_gpu_utilization() -> Dict[str, Any]: device = get_device() @@ -590,6 +862,10 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: ) if result is not None: result["backend"] = _backend_label(device) + numeric_ids = parent_visible_spec.get("numeric_ids") + if IS_ROCM and numeric_ids is not None: + # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.) + _reconcile_rocm_unified_memory(result, numeric_ids) return result # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) @@ -689,7 +965,15 @@ def _get_parent_visible_gpu_spec() -> Dict[str, Any]: # Use explicit None checks (not `or`) so empty string "" is honoured # as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES. cuda_visible = None - if IS_ROCM: + # Prefer ROCm masks only on a ROCm host, or when no CUDA mask is set, so a + # stale HIP_VISIBLE_DEVICES on an NVIDIA host can't override CUDA_VISIBLE_DEVICES. + _is_rocm_spec = IS_ROCM or ( + "CUDA_VISIBLE_DEVICES" not in os.environ + and ( + "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ + ) + ) + if _is_rocm_spec: hip_vis = os.environ.get("HIP_VISIBLE_DEVICES") rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES") if hip_vis is not None: @@ -865,7 +1149,57 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non def _determine_attention_impl_for_gpu_estimate(config) -> str: - import copy as _copy + # torch.distributed is incomplete on Windows ROCm — torch._C is a C + # extension (not a package), so Python cannot import the submodule + # torch._C._distributed_c10d that torch.distributed depends on. + # Inject an empty stub into sys.modules BEFORE importing torch.distributed + # so the import succeeds, then patch the missing process-group helpers. + if sys.platform == "win32" and IS_ROCM: + # Dummy class for any name torch.distributed tries to import from these stubs + class _Dummy: + pass + + for _c10d_name in ( + "torch._C._distributed_c10d", + "torch._C._distributed_autograd", + "torch._C._distributed_rpc", + ): + if _c10d_name not in sys.modules: + _stub = types.ModuleType(_c10d_name) + # torch.distributed imports these names from _distributed_c10d; + # provide no-op dummies so the import doesn't raise AttributeError. + for _sym in ( + "FakeProcessGroup", + "ProcessGroup", + "Work", + "Store", + "PrefixStore", + "FileStore", + "TCPStore", + "HashStore", + "Reducer", + "Logger", + "DistributedDebugLevel", + "GradBucket", + "BuiltinCommHookType", + ): + setattr(_stub, _sym, _Dummy) + sys.modules[_c10d_name] = _stub + + try: + import torch.distributed as _td + + for _attr, _stub in ( + ("is_initialized", lambda: False), + ("is_available", lambda: False), + ("get_rank", lambda: 0), + ("get_world_size", lambda: 1), + ("is_torchelastic_launched", lambda: False), + ): + if not hasattr(_td, _attr): + setattr(_td, _attr, _stub) + except ImportError: + pass from unsloth.models._utils import resolve_attention_implementation from transformers import AutoModel, AutoModelForCausalLM @@ -875,7 +1209,7 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str: # `sub_configs` and propagates to nested text_config / sub-configs, so a # shallow copy still mutates those shared inner objects on the cached # config returned by _load_config_for_gpu_estimate. Deepcopy isolates them. - config_copy = _copy.deepcopy(config) + config_copy = copy.deepcopy(config) model_class = None for auto_model in (AutoModelForCausalLM, AutoModel): @@ -1062,7 +1396,10 @@ def estimate_required_model_memory_gb( _determine_attention_impl_for_gpu_estimate(config) ) except Exception as e: - logger.warning( + # Log at debug: on Windows ROCm the torch.distributed stub does + # not implement Store, so this fires on every estimate call. + # It is expected and non-actionable -- eager is the safe fallback. + logger.debug( "Could not resolve attention implementation for '%s': %s", estimate_model, e, @@ -1552,14 +1889,35 @@ def apply_gpu_ids(gpu_ids) -> None: # parent process already set a ROCm visibility variable -- that # way a downstream ROCm process inherits the narrowed mask even # before Studio's hardware detection has classified the host. + # Final fallback: probe torch.version.hip so AMD workers without + # HIP_VISIBLE_DEVICES still get the correct ROCm visibility mask. _inherits_rocm_visibility = ( "HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ ) - if IS_ROCM or _inherits_rocm_visibility: + _is_rocm = IS_ROCM or _inherits_rocm_visibility + if not _is_rocm: + # torch.version.hip is a non-empty string on ROCm, None on CUDA. + # AMD SDK / Radeon ROCm wheels can leave torch.version.hip unset but + # still encode "rocm" in torch.__version__, matching detect_hardware(). + # Broad except: a probe failure must never crash a training worker. + try: + import torch as _torch + + _is_rocm = ( + getattr(_torch.version, "hip", None) is not None + or "rocm" in getattr(_torch, "__version__", "").lower() + ) + except Exception as e: + logger.debug( + "apply_gpu_ids: torch ROCm probe skipped (%s: %s)", + type(e).__name__, + e, + ) + if _is_rocm: os.environ["HIP_VISIBLE_DEVICES"] = value os.environ["ROCR_VISIBLE_DEVICES"] = value _visible_gpu_count = None - if IS_ROCM or _inherits_rocm_visibility: + if _is_rocm: logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value) else: logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) @@ -1652,8 +2010,6 @@ def safe_num_proc(desired: Optional[int] = None) -> int: Returns: A safe integer ≥ 1. """ - import sys - # Windows and macOS use 'spawn' for multiprocessing -- the overhead of # re-importing torch/transformers/unsloth per worker is typically slower # than single-process. @@ -1704,8 +2060,6 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``). Only ``num_proc=None`` guarantees in-process execution. """ - import sys - if sys.platform in ("win32", "darwin"): return None return safe_num_proc(desired) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 993995ee57..dc34444ccb 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1156,13 +1156,18 @@ def detect_gguf_model(path: str) -> Optional[str]: p = Path(path) # Case 1: direct .gguf file - if p.suffix.lower() == ".gguf" and p.is_file(): + if p.suffix.lower() == ".gguf": if _is_mmproj(p.name): return None - # Use absolute (not resolve) to preserve symlink names -- e.g. - # Ollama .studio_links/model.gguf -> blobs/sha256-... should - # keep the readable symlink name, not the opaque blob hash. - return str(p.absolute()) + # Extension is authoritative: don't gate on is_file()/exists(), which + # can fail in the Windows lock window after llama-server is killed. + try: + is_dir = p.is_dir() + except OSError: + is_dir = False # stat() unavailable in the lock window + if not is_dir: + return str(p.absolute()) # absolute() keeps symlink names readable + # Directory named "*.gguf": fall through to the dir scan below. # Case 2: directory containing .gguf files (skip mmproj) if p.is_dir(): diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 5c42e890d1..e0ce02261b 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -15,6 +15,7 @@ import urllib.request from typing import Callable from utils.native_path_leases import child_env_without_native_path_secret +from utils.subprocess_compat import windows_hidden_subprocess_kwargs _logger = logging.getLogger(__name__) @@ -106,6 +107,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non text = True, timeout = timeout, env = child_env_without_native_path_secret(), + **windows_hidden_subprocess_kwargs(), ) except subprocess.TimeoutExpired: return None diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 849e017ea8..4e4a6130cd 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -49,7 +49,7 @@ import { Edit03Icon, Globe02Icon, HelpCircleIcon, - Logout01Icon, + Logout05Icon, Search01Icon, PowerIcon, PencilEdit02Icon, @@ -796,7 +796,7 @@ export function AppSidebar() { void navigate({ to: "/login" }); }} > - + {t("shell.navigation.logOut")} setShutdownOpen(true)}> diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index b9d63ecee9..861a501570 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -9,6 +9,7 @@ from __future__ import annotations import argparse import errno import fnmatch +import functools import hashlib import json import os @@ -18,6 +19,7 @@ import re import shutil import site import socket +import struct import subprocess import sys import tarfile @@ -29,7 +31,7 @@ import urllib.parse import urllib.request import zipfile from contextlib import contextmanager -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace as dataclasses_replace try: from filelock import FileLock, Timeout as FileLockTimeout @@ -100,6 +102,39 @@ DEFAULT_PUBLISHED_SHA256_ASSET = os.environ.get( ) UPSTREAM_REPO = "ggml-org/llama.cpp" UPSTREAM_RELEASES_API = f"https://api.github.com/repos/{UPSTREAM_REPO}/releases/latest" + +LEMONADE_ROCM_REPO = "lemonade-sdk/llamacpp-rocm" +LEMONADE_ROCM_RELEASES_API = ( + f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/latest" +) + + +def _lemonade_release_api_for(llama_tag: str) -> str: + """Return the GitHub API URL for the lemonade release that matches a + requested llama.cpp tag. + + When llama_tag is unset or "latest", point at /releases/latest. When the + caller has pinned a specific tag (e.g. "b1260"), point at the same tag in + lemonade. Lemonade tracks `ggml-org/llama.cpp` build tags (e.g. "b1260") + but is NOT guaranteed to publish every upstream build -- lemonade may be + several builds behind ggml-org. Pinning to a specific tag that lemonade + skipped will produce a 404 and the caller falls through to the upstream + tarball; that is intentional so pinned installs stay reproducible. + Do NOT pass a `unslothai/llama.cpp` fork tag -- the fork uses its own + namespace and will always 404 against lemonade. + + The tag is URL-encoded with `safe=""` so an unexpected slash / hash / query + character cannot reshape the URL. + """ + normalized = (llama_tag or "").strip() + if not normalized or normalized.lower() == "latest": + return LEMONADE_ROCM_RELEASES_API + return ( + f"https://api.github.com/repos/{LEMONADE_ROCM_REPO}/releases/tags/" + f"{urllib.parse.quote(normalized, safe = '')}" + ) + + TEST_MODEL_URL = ( "https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf" ) @@ -126,6 +161,14 @@ DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS = env_int( 2, minimum = 1, ) +# Deeper macOS-only walk-back: upstream can ship a run of prebuilts built for a +# newer macOS than the host, only caught at validate time, so an older host must +# skip the whole run. Free on new hosts (first plan validates, extras unused). +DEFAULT_MAX_MACOS_RELEASE_FALLBACKS = env_int( + "UNSLOTH_LLAMA_MAX_MACOS_RELEASE_FALLBACKS", + 16, + minimum = 1, +) FORCE_COMPILE_DEFAULT_REF = os.environ.get("UNSLOTH_LLAMA_FORCE_COMPILE_REF", "master") DIRECT_LINUX_BUNDLE_PROFILES: dict[str, dict[str, Any]] = { @@ -179,6 +222,71 @@ DIRECT_LINUX_BUNDLE_PROFILES: dict[str, dict[str, Any]] = { }, } +# Lowest CUDA major we ship prebuilts for, and the highest major we probe for +# installed runtime libraries. Detection and runtime-line derivation are +# generated per major so a new toolkit (cuda14, ...) needs no code change while +# llama.cpp keeps the cudart64_.dll / libcudart.so. naming. +_MIN_CUDA_MAJOR = 12 +_MAX_PROBE_CUDA_MAJOR = 19 + +# Last ggml-org release whose Windows win-cuda-13 build is still sub-13.3 +# (cuda-13.1, b9360, 2026-05-27). Upstream bumped win-cuda-13 to 13.3 at b9365 +# and now ships only cuda-12.4 + cuda-13.3. cuda-12.4 predates Blackwell (ggml +# compiles sm_120 only at toolkit >= 12.8), so a Blackwell host on a 13.1/13.2 +# driver is gated off 13.3 and would drop to a CPU-only 12.4 build. b9360 is +# immutable, so we pin its cuda-13.1 build (plus paired cudart) as a GPU +# fallback for exactly those hosts. See unslothai/unsloth#5887. +_PINNED_BLACKWELL_FALLBACK_TAG = "b9360" +_PINNED_BLACKWELL_FALLBACK_RUNTIME = "13.1" +_PINNED_BLACKWELL_DRIVER_FLOOR = (13, 1) +_BLACKWELL_MIN_SM = 120 +# ggml compiles Blackwell sm_120 only at toolkit >= 12.8, so an in-release +# windows-cuda build at or above this already covers Blackwell and makes the +# older pinned 13.1 fallback unnecessary (cuda-12.4 is below it). +_BLACKWELL_MIN_TOOLKIT = (12, 8) +_PINNED_BLACKWELL_LLAMA_SHA256 = ( + "31ddb8b42d7ab4a47cab8c48c397519f580ca502df7e73f3ab396eacc16c8e8d" +) +_PINNED_BLACKWELL_CUDART_SHA256 = ( + "f96935e7e385e3b2d0189239077c10fe8fd7e95690fea4afec455b1b6c7e3f18" +) + + +def _cuda_runtime_lines_for_major(major: int) -> list[str]: + """Runtime lines a driver of this CUDA major can use, newest major first + down to the minimum we ship. A driver runs its own major and any older one + (backward compatibility).""" + return [f"cuda{m}" for m in range(major, _MIN_CUDA_MAJOR - 1, -1)] + + +def _resolve_linux_bundle_profile(bundle_profile: str) -> "dict[str, Any] | None": + """Profile (runtime line + sm coverage) for a linux-x64-cuda- + bundle. Known majors use their published coverage; an unknown future major + reuses the newest known major's coverage for the same class as a forward + default, with the post-build GPU smoke test as the backstop.""" + known = DIRECT_LINUX_BUNDLE_PROFILES.get(bundle_profile) + if known is not None: + return known + m = re.fullmatch( + r"cuda(?P\d+)-(?Polder|newer|portable)", bundle_profile + ) + if not m: + return None + base_key = max( + ( + k + for k, v in DIRECT_LINUX_BUNDLE_PROFILES.items() + if v["coverage_class"] == m.group("klass") + ), + key = lambda k: int(re.match(r"cuda(\d+)-", k).group(1)), + default = None, + ) + if base_key is None: + return None + profile = dict(DIRECT_LINUX_BUNDLE_PROFILES[base_key]) + profile["runtime_line"] = f"cuda{m.group('major')}" + return profile + @dataclass class HostInfo: @@ -196,6 +304,10 @@ class HostInfo: has_physical_nvidia: bool has_usable_nvidia: bool has_rocm: bool = False + rocm_gfx_target: str | None = None + # (major, minor) from platform.mac_ver(); None off macOS or if unparseable. + # Skips a macos prebuilt whose minimum-OS exceeds this host. + macos_version: tuple[int, int] | None = None @dataclass @@ -722,6 +834,26 @@ def windows_cuda_asset_aliases( return aliases +def _published_windows_cuda_runtime( + upstream_assets: dict[str, str], major: int, driver: tuple[int, int] | None +) -> str | None: + """Highest cuda-. published upstream that `driver` can run by + default CUDA compatibility, i.e. (major, minor) <= driver. None if nothing + qualifies. Gating on the driver (not just the major) keeps a 13.3 build off + a driver that only advertises 13.1, where it would otherwise rely on the + unguaranteed minor-version-compatibility path.""" + if driver is None: + return None + best: int | None = None + for name in upstream_assets: + m = re.search(r"-bin-win-cuda-(\d+)\.(\d+)-x64\.zip$", name) + if m and int(m.group(1)) == major: + minor = int(m.group(2)) + if (major, minor) <= driver and (best is None or minor > best): + best = minor + return f"{major}.{best}" if best is not None else None + + def format_byte_count(num_bytes: float) -> str: units = ["B", "KiB", "MiB", "GiB", "TiB"] value = float(num_bytes) @@ -1178,7 +1310,7 @@ def parse_direct_linux_release_bundle( inferred_labels: list[str] = [] linux_asset_re = re.compile( - r"^app-(?P