diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 33ac3b9bd8..0ef2ad1e9d 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -434,7 +434,7 @@ jobs: # ───────────────────────────────────────────────────────────── # Semgrep: design-flaw detection (catches what regex-pattern - # scanning of malicious authors cannot — first-party logic bugs + # scanning of malicious authors cannot, e.g. first-party logic bugs # like langchain-core CVE-2025-68664 dumps/dumpd injection, # n8n CVE-2025-68668 _pyodide.eval_code sandbox escape, marimo # CVE-2026-39987 unauth WebSocket). @@ -849,10 +849,13 @@ jobs: grep -q "Standalone pre-install package scanner" scripts/scan_packages.py - name: Scan declared + transitive Python deps - # scan_packages.py exits 1 on CRITICAL/HIGH findings, 0 on - # clean. We swallow the exit because the baseline isn't - # triaged yet; surface the findings in the workflow summary. - # Drop continue-on-error after the first clean run on main. + # scan_packages.py exits 1 on NON-baselined CRITICAL/HIGH + # findings, 0 otherwise. It scans code-only (docstrings and + # comments are blanked first) and suppresses reviewed + # known-good findings via scripts/scan_packages_baseline.json, + # so legitimate-library noise no longer red-fails the gate. + # The step stays advisory until SCAN_ENFORCE=1 (see env below); + # then PIPESTATUS propagates the scanner's exit code. # # `--with-deps` walks PyPI metadata to enumerate every # transitive dep the declared set would install, then scans @@ -869,6 +872,14 @@ jobs: # downloads in exchange for wall-clock parallelism. env: SHARD_FILES: ${{ matrix.shard.files }} + # Enforcement switch. "1" = blocking: a non-baselined CRITICAL/HIGH + # fails the build. scan_packages.py scans code-only (docstrings/comments + # stripped), fetches sdist-only packages directly from PyPI (no build) + # so every shard resolves, and honors the reviewed allowlist at + # scripts/scan_packages_baseline.json, so only NON-baselined + # CRITICAL/HIGH cause its exit 1. The committed baseline makes all three + # shards exit 0 today; set this back to "0" to return to advisory. + SCAN_ENFORCE: "1" run: | set +e mkdir -p logs @@ -884,12 +895,14 @@ jobs: fi done echo "::endgroup::" + rc=0 if [ ${#REQ_ARGS[@]} -eq 0 ]; then echo "[security-audit] shard ${{ matrix.shard.id }}: no PyPI specs, nothing to scan" \ | tee "$LOG" else python scripts/scan_packages.py --with-deps "${REQ_ARGS[@]}" \ 2>&1 | tee "$LOG" + rc=${PIPESTATUS[0]} fi { echo "## scan_packages :: shard ${{ matrix.shard.id }}" @@ -897,11 +910,19 @@ jobs: echo "### Files in this shard" for f in $SHARD_FILES; do echo "- audit-reqs/$f.txt"; done echo + echo "scan_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)" + echo echo '### Findings (tail)' echo '```' tail -200 "$LOG" echo '```' } >> "$GITHUB_STEP_SUMMARY" + # Advisory by default; blocking once SCAN_ENFORCE=1 and the baseline + # is committed. PIPESTATUS is captured above so `tee` does not mask the + # scanner's exit code. + if [ "$SCAN_ENFORCE" = "1" ]; then + exit "$rc" + fi - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 if: always() @@ -975,24 +996,37 @@ jobs: python3 -c "import ast; ast.parse(open('scripts/scan_npm_packages.py').read())" - name: Scan npm tarballs (declared + transitive, no install) - # The script exits 1 on HIGH/CRITICAL findings; we capture the - # full log and surface it in the step summary either way. It - # never runs `npm install`, never executes anything from a - # downloaded tarball, and only fetches from registry.npmjs.org. - # Initially non-blocking so the baseline can settle; drop - # continue-on-error once the baseline is clean for a week. + # scan_npm_packages.py exits 1 on NON-baselined HIGH/CRITICAL + # findings, 0 otherwise. It scans code-only (JS/TS comments are + # blanked first) and honors a reviewed allowlist at + # scripts/scan_npm_packages_baseline.json. It never runs + # `npm install`, never executes anything from a downloaded + # tarball, and only fetches from registry.npmjs.org. The npm + # corpus is clean (the baseline is empty), so the gate is + # enforcing (SCAN_ENFORCE=1) and any new finding fails the build. + env: + SCAN_ENFORCE: "1" run: | - set -o pipefail + set +e LOG=logs-scan-npm.txt python3 scripts/scan_npm_packages.py 2>&1 | tee "$LOG" + rc=${PIPESTATUS[0]} { echo "## scan_npm_packages" echo + echo "scan_npm_packages.py exit code: $rc (enforce=$SCAN_ENFORCE)" + echo echo '### Findings (tail)' echo '```' tail -300 "$LOG" echo '```' } >> "$GITHUB_STEP_SUMMARY" + # Blocking: the npm corpus is clean, so any non-baselined + # HIGH/CRITICAL is new and should fail the build. PIPESTATUS is + # captured above so `tee` does not mask the scanner's exit code. + if [ "$SCAN_ENFORCE" = "1" ]; then + exit "$rc" + fi - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 if: always() diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 88c7344683..b394b308e4 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -224,7 +224,8 @@ jobs: tests/sh/test_mac_intel_compat.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_tauri_install_exit_order.sh \ - tests/sh/test_torch_constraint.sh; do + tests/sh/test_torch_constraint.sh \ + tests/sh/test_torch_flavor.sh; do echo "::group::$s" bash "$s" echo "::endgroup::" diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index cffb33f71d..c2c4fa03bf 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -317,7 +317,7 @@ jobs: timeout-minutes: 25 env: # Tool calling is the highest-volume GGUF in this workflow - # (Qwen3.5-2B at IQ3_XXS = ~890 MiB). Caching HF_HOME would + # (Qwen3.5-2B at Q4_K_XL = ~1.28 GiB). Caching HF_HOME would # store xet chunks + blobs + snapshots = ~4 GiB compressed -- # 4-5x file-size inflation, dominated by xet chunks. Use main's # `--local-dir gguf-cache` pattern to cache the flat .gguf only. @@ -326,8 +326,11 @@ jobs: # path keeps the test off HF_HOME entirely so the cache size # tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images # jobs still cover the gguf_variant resolution path. + # Q4_K_XL, not IQ3_XXS: at IQ3_XXS this model emits malformed + # tool calls that llama-server's peg-native parser rejects with a + # 500. Mac/Windows already use Q4_K_XL for the same reason. GGUF_REPO: unsloth/Qwen3.5-2B-GGUF - GGUF_FILE: Qwen3.5-2B-UD-IQ3_XXS.gguf + GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf STUDIO_PORT: '18889' steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -772,6 +775,9 @@ jobs: kill "${STUDIO_PID}" 2>/dev/null || true sleep 2 ss -tln | grep ":${STUDIO_PORT}" || true + # Capture backend + llama-server logs so a 500 has a server-side traceback. + mkdir -p logs/server-logs + cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true - name: Upload logs # Always upload so green runs are still reviewable. @@ -784,6 +790,7 @@ jobs: path: | logs/studio.log logs/install.log + logs/server-logs/ retention-days: 7 # ───────────────────────────────────────────────────────────────────── diff --git a/.github/workflows/studio-ui-smoke.yml b/.github/workflows/studio-ui-smoke.yml index de106e201f..dcf9fd26af 100644 --- a/.github/workflows/studio-ui-smoke.yml +++ b/.github/workflows/studio-ui-smoke.yml @@ -276,6 +276,10 @@ jobs: run: | kill "${STUDIO_IME_PID}" 2>/dev/null || true sleep 2 + # Capture backend + llama-server logs (all three Studios share this + # dir) so a stray 500 has a server-side traceback. + mkdir -p logs/server-logs + cp -r ~/.unsloth/studio/logs/. logs/server-logs/ 2>/dev/null || true - name: Upload Playwright artifacts # Always upload so a green run's screenshots stay reviewable -- @@ -289,6 +293,7 @@ jobs: logs/studio_extra.log logs/studio_ime.log logs/install.log + logs/server-logs/ logs/playwright logs/playwright_extra logs/playwright_ime diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index a772a6d102..a6f1401067 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -65,17 +65,20 @@ jobs: with: persist-credentials: false - # Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit - # test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke. - - name: setup.ps1 unit test (Resolve-CudaToolkit) + # Fast GPU-free gate: parse install.ps1 + setup.ps1 and run the PowerShell + # unit tests (CUDA-toolkit + torch-flavor helpers) before the heavy GGUF smoke. + - name: PowerShell installer unit tests shell: pwsh run: | - $errs = $null - [void][System.Management.Automation.Language.Parser]::ParseFile( - (Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs) - if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } - Write-Host "setup.ps1 parsed with no errors" + foreach ($f in @('install.ps1', 'studio/setup.ps1')) { + $errs = $null + [void][System.Management.Automation.Language.Parser]::ParseFile( + (Resolve-Path $f).Path, [ref]$null, [ref]$errs) + if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 } + Write-Host "$f parsed with no errors" + } pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1 + pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1 - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 with: diff --git a/.github/workflows/studio-windows-ui-smoke.yml b/.github/workflows/studio-windows-ui-smoke.yml index 40d8e530cd..00458d213b 100644 --- a/.github/workflows/studio-windows-ui-smoke.yml +++ b/.github/workflows/studio-windows-ui-smoke.yml @@ -136,6 +136,17 @@ jobs: } } + - name: Seed a legacy launch-studio.vbs (upgrade-cleanup check) + # Simulate a pre-hardening install so the post-install assertion below + # proves the installer DELETES an existing launch-studio.vbs (the exact + # Kaspersky-flagged file), not merely stops generating it. + shell: pwsh + run: | + $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio' + New-Item -ItemType Directory -Force -Path $appDir | Out-Null + Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode + Write-Host "seeded legacy launch-studio.vbs at $appDir" + - name: Install Studio (--local, --no-torch) # install.ps1 is the supported Windows installer. install.sh # has no Windows branch (apt-get / brew calls). The PS1 @@ -192,6 +203,69 @@ jobs: echo "install.ps1 installed the Windows prebuilt llama.cpp:" cat "$INFO" + - name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut) + # The shortcut launch path is otherwise untested here (the steps below + # boot `unsloth studio` directly). Guard against re-introducing the VBS + # that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk + # pointing anywhere other than hidden PowerShell over launch-studio.ps1. + shell: pwsh + run: | + $appDir = Join-Path $env:LOCALAPPDATA 'Unsloth Studio' + if (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.vbs')) { + throw "regression: launch-studio.vbs exists (the Kaspersky VBS-FP shape)" + } + if (-not (Test-Path -LiteralPath (Join-Path $appDir 'launch-studio.ps1'))) { + throw "missing launch-studio.ps1 in $appDir" + } + $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk' + if (-not (Test-Path -LiteralPath $lnk)) { + $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk' + } + if (-not (Test-Path -LiteralPath $lnk)) { throw "no Unsloth Studio.lnk on Desktop or Start Menu" } + $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk) + Write-Host "shortcut target: $($sc.TargetPath)" + Write-Host "shortcut args: $($sc.Arguments)" + if ($sc.TargetPath -match 'wscript\.exe$') { throw "shortcut still targets wscript.exe (VBS host)" } + if ($sc.TargetPath -notmatch 'powershell\.exe$') { throw "unexpected shortcut target: $($sc.TargetPath)" } + if ($sc.Arguments -notmatch '-WindowStyle Hidden') { + throw "shortcut must launch windowless (-WindowStyle Hidden)" + } + Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)" + + - name: Launch Studio via the shortcut and assert health + # Run the exact command the .lnk stores (hidden PowerShell over + # launch-studio.ps1) and confirm it brings the backend up. This is the + # only step that proves the shortcut launch is not silently broken. + # Default port range is 8888-8908; the later UI tests use 18896/18897, so + # there is no conflict, and we tear this server down before they boot. + shell: pwsh + run: | + $lnk = Join-Path ([Environment]::GetFolderPath('Desktop')) 'Unsloth Studio.lnk' + if (-not (Test-Path -LiteralPath $lnk)) { + $lnk = Join-Path $env:APPDATA 'Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk' + } + $sc = (New-Object -ComObject WScript.Shell).CreateShortcut($lnk) + Write-Host "launching: $($sc.TargetPath) $($sc.Arguments)" + Start-Process -FilePath $sc.TargetPath -ArgumentList $sc.Arguments -WorkingDirectory $sc.WorkingDirectory + $foundPort = 0 + foreach ($i in 1..180) { + foreach ($port in 8888..8908) { + try { + $r = Invoke-RestMethod -Uri "http://127.0.0.1:$port/api/health" -TimeoutSec 1 + if ($r.status -eq 'healthy' -and $r.service -eq 'Unsloth UI Backend') { $foundPort = $port; break } + } catch {} + } + if ($foundPort) { break } + Start-Sleep -Seconds 1 + } + # Tear down the shortcut-launched server before the main UI tests boot. + try { + $owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess + if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null } + } catch {} + if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" } + Write-Host "Studio healthy on port $foundPort (launched via the shortcut)" + - name: Add Studio shim to GITHUB_PATH # install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe # and adds that dir to the User PATH via the Windows registry. diff --git a/.gitignore b/.gitignore index a839633790..9f7d4b8c60 100644 --- a/.gitignore +++ b/.gitignore @@ -235,3 +235,5 @@ package-lock.json !studio/backend/core/data_recipe/oxc-validator/package-lock.json !studio/package-lock.json llama.cpp/ +# Stray "~" dir some tools create from a literal ~ TMPDIR; never part of the repo. +/~/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cffbf73cd5..bf2a0c8e7c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.16 + rev: v0.15.17 hooks: - id: ruff args: diff --git a/install.ps1 b/install.ps1 index 396398d347..f797b80296 100644 --- a/install.ps1 +++ b/install.ps1 @@ -494,6 +494,37 @@ function Install-UnslothStudio { } } + # Retry Invoke-InstallCommand on transient uv download failures with backoff. + # Returns the last exit code on permanent failure so rollback still fires. + function Invoke-InstallCommandRetry { + param( + [Parameter(Mandatory = $true, Position = 0)][ScriptBlock]$Command, + [string]$Label = "install step" + ) + # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables). + # TryParse with bounds avoids an Int32 overflow throw. Bounds: 1..100 retries, 0..3600s. + $maxAttempts = 3 + $parsedAttempts = 0 + if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRIES, [ref]$parsedAttempts) -and $parsedAttempts -ge 1 -and $parsedAttempts -le 100) { + $maxAttempts = $parsedAttempts + } + $delay = 3 + $parsedDelay = 0 + if ([int]::TryParse($env:UNSLOTH_INSTALL_RETRY_DELAY, [ref]$parsedDelay) -and $parsedDelay -ge 0 -and $parsedDelay -le 3600) { + $delay = $parsedDelay + } + $attempt = 1 + while ($true) { + $code = Invoke-InstallCommand $Command + if ($code -eq 0) { return 0 } + if ($attempt -ge $maxAttempts) { return $code } + substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow" + Start-Sleep -Seconds $delay + $attempt++ + $delay = $delay * 2 + } + } + function New-StudioShortcuts { param( [Parameter(Mandatory = $true)][string]$UnslothExePath @@ -518,7 +549,6 @@ function Install-UnslothStudio { } $appDir = $StudioDataDir $launcherPs1 = Join-Path $appDir "launch-studio.ps1" - $launcherVbs = Join-Path $appDir "launch-studio.vbs" $desktopDir = [Environment]::GetFolderPath("Desktop") $desktopLink = if ($desktopDir -and $desktopDir.Trim()) { Join-Path $desktopDir "Unsloth Studio.lnk" @@ -811,19 +841,30 @@ exit 0 # even when install.ps1 is executed from PowerShell 7. $utf8Bom = New-Object System.Text.UTF8Encoding($true) [System.IO.File]::WriteAllText($launcherPs1, $launcherContent, $utf8Bom) - # shell.Run(cmd, 0, ...) already hides the window, so -WindowStyle Hidden - # is redundant; omitting it trims an AV-heuristic token (Kaspersky FP). - $vbsContent = @" -Set shell = CreateObject("WScript.Shell") -cmd = "powershell -NoProfile -ExecutionPolicy Bypass -File ""$launcherPs1""" -shell.Run cmd, 0, False -"@ - # WSH handles UTF-16LE reliably for .vbs files with non-ASCII paths. - Set-Content -LiteralPath $launcherVbs -Value $vbsContent -Encoding Unicode -Force + # No .vbs launcher is written. A WScript.Shell .vbs that spawns a hidden + # ExecutionPolicy-Bypass PowerShell is exactly the shape VBS-dropper + # heuristics score (e.g. Kaspersky HEUR:Trojan.VBS.Agent.gen). The .lnk + # shortcuts instead point straight at powershell.exe running + # launch-studio.ps1 with a hidden window (selected below). + + # Delete any launch-studio.vbs left by a pre-hardening install. New + # installs no longer generate it, but an upgrade that merely stopped + # generating it would leave the exact file AV flags on disk, so remove + # it explicitly. Covers default and env-mode installs (same $appDir). + $legacyLauncherVbs = Join-Path $appDir "launch-studio.vbs" + if (Test-Path -LiteralPath $legacyLauncherVbs) { + Remove-Item -LiteralPath $legacyLauncherVbs -Force -ErrorAction SilentlyContinue + } # Prefer bundled icon from local clone/dev installs. # If not available, best-effort download from raw GitHub. # We only attach the icon if the resulting file has a valid ICO header. + # Snapshot the existing icon first so we can tell whether it actually + # changed and gate the heavier icon-cache refresh on a real change. + $preIconHash = $null + if (Test-Path -LiteralPath $iconPath) { + try { $preIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash } catch {} + } $hasValidIcon = $false if ($bundledIcon -and (Test-Path -LiteralPath $bundledIcon)) { try { @@ -859,6 +900,24 @@ shell.Run cmd, 0, False } } + # Did the icon content actually change vs the previous install? + # Only a real change (or a first/removed icon) should trigger the heavy + # refresh; a no-op reinstall with no icon at all must not. + $iconChanged = $false + if ($hasValidIcon) { + if (-not $preIconHash) { + $iconChanged = $true + } else { + try { + $postIconHash = (Get-FileHash -LiteralPath $iconPath -Algorithm SHA256).Hash + $iconChanged = ($postIconHash -ne $preIconHash) + } catch { $iconChanged = $true } + } + } elseif ($preIconHash) { + # A previously present icon was removed or invalidated. + $iconChanged = $true + } + # Env-mode: skip persistent Desktop / Start Menu .lnk shortcuts # that may point at a deleted workspace; launcher + icon stay. if ($StudioRedirectMode -eq 'env') { @@ -866,8 +925,22 @@ shell.Run cmd, 0, False return } - $wscriptExe = Join-Path $env:SystemRoot "System32\wscript.exe" - $shortcutArgs = "//B //Nologo `"$launcherVbs`"" + # Whether this is effectively a first install (no pre-existing .lnk). + # Used to gate the heavier icon-cache refresh below so a no-op reinstall + # does not repeatedly clear caches / restart StartMenuExperienceHost -- + # a behavioral cluster AV heuristics can score as dropper-like. + $firstInstall = -not ( + ($desktopLink -and (Test-Path -LiteralPath $desktopLink)) -or + ($startMenuLink -and (Test-Path -LiteralPath $startMenuLink)) + ) + + # Launch transport for the shortcuts: powershell.exe runs + # launch-studio.ps1 with a hidden window. We deliberately avoid a + # .vbs/WScript.Shell wrapper -- that script-engine shape is what AV + # VBS-dropper heuristics score (Kaspersky HEUR:Trojan.VBS.Agent.gen). + $powershellForLnk = Join-Path $env:SystemRoot "System32\WindowsPowerShell\v1.0\powershell.exe" + $shortcutTarget = $powershellForLnk + $shortcutArgs = "-NoProfile -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$launcherPs1`"" try { $wshell = New-Object -ComObject WScript.Shell @@ -877,9 +950,11 @@ shell.Run cmd, 0, False if (-not $linkPath -or [string]::IsNullOrWhiteSpace($linkPath)) { continue } try { $shortcut = $wshell.CreateShortcut($linkPath) - $shortcut.TargetPath = $wscriptExe + $shortcut.TargetPath = $shortcutTarget $shortcut.Arguments = $shortcutArgs $shortcut.WorkingDirectory = $appDir + # Start minimized so the brief PowerShell console flash is muted. + $shortcut.WindowStyle = 7 $shortcut.Description = "Launch Unsloth Studio" if ($hasValidIcon) { $shortcut.IconLocation = "$iconPath,0" @@ -893,15 +968,13 @@ shell.Run cmd, 0, False } if ($createdShortcutCount -gt 0) { substep "Created Unsloth Studio shortcut" - # Force Explorer to re-read each new shortcut's icon so it renders - # immediately instead of a stale/generic entry (a same-name .lnk - # recreated across reinstalls keeps Explorer's cached per-item icon). - # The reliable, non-disruptive fix (no explorer restart) is a per-item - # SHChangeNotify SHCNE_UPDATEITEM + SHCNF_PATHW per .lnk; the global - # SHCNE_ASSOCCHANGED broadcast alone does NOT recover a stale item. - # Also clear the on-disk icon cache (covers heavier staleness). - try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {} - try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {} + # Always do the cheap, non-disruptive per-item refresh so a + # rewritten same-name .lnk renders with its new target/icon + # immediately (a same-name .lnk recreated across reinstalls keeps + # Explorer's cached per-item icon). The reliable fix (no explorer + # restart) is a per-item SHChangeNotify SHCNE_UPDATEITEM + + # SHCNF_PATHW per .lnk; the global SHCNE_ASSOCCHANGED broadcast + # alone does NOT recover a stale item. try { Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int eventId, uint flags, string item1, System.IntPtr item2);' -ErrorAction SilentlyContinue # SHCNE_UPDATEITEM (0x00002000) + SHCNF_PATHW (0x0005) per shortcut @@ -911,21 +984,31 @@ shell.Run cmd, 0, False # SHCNE_ASSOCCHANGED (0x08000000) global refresh (belt-and-suspenders) [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, $null, [System.IntPtr]::Zero) } catch {} - # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN - # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT - # invalidate, so a rewritten same-name shortcut shows the old tile - # until the host restarts. Drop only the render caches (NEVER - # start2.bin -- the pinned layout) and let the host rebuild. - # Best-effort; Win10 has no such host (Test-Path skips it). - try { - $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState" - if (Test-Path -LiteralPath $smehTemp) { - Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue | - Remove-Item -Force -ErrorAction SilentlyContinue - Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue - Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue - } - } catch {} + # Heavier on-disk icon-cache clear + StartMenuExperienceHost tile + # rebuild only when the icon actually changed or this is a first + # install. Running "clear icon cache + kill StartMenuExperienceHost" + # on every no-op reinstall is a dropper-like behavioral cluster and + # is unnecessary when the icon is unchanged (the per-item notify + # above already refreshes the rewritten shortcut). + if ($firstInstall -or $iconChanged) { + try { & "$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache 2>$null } catch {} + try { & "$env:SystemRoot\System32\ie4uinit.exe" -show 2>$null } catch {} + # Win11's Start Menu (StartMenuExperienceHost) keeps its OWN + # pre-rendered tile-icon cache that ie4uinit/explorer restart do NOT + # invalidate, so a rewritten same-name shortcut shows the old tile + # until the host restarts. Drop only the render caches (NEVER + # start2.bin -- the pinned layout) and let the host rebuild. + # Best-effort; Win10 has no such host (Test-Path skips it). + try { + $smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState" + if (Test-Path -LiteralPath $smehTemp) { + Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue | + Remove-Item -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue + Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue + } + } catch {} + } } else { substep "no Unsloth Studio shortcuts were created" "Yellow" } @@ -1193,7 +1276,7 @@ shell.Run cmd, 0, False # ── Install uv ── Write-TauriLog "STEP" "Installing uv package manager" - $UvMinVersion = "0.7.22" + $UvMinVersion = "0.8.16" function Test-UvVersionOk { $cmd = Get-Command uv -ErrorAction SilentlyContinue if (-not $cmd) { return $false } @@ -1264,6 +1347,15 @@ shell.Run cmd, 0, False $env:UV_COMPILE_BYTECODE_TIMEOUT = "180" } + # uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read + # timeout for large wheel downloads. User-provided values are preserved. + if (-not $env:UV_HTTP_RETRIES) { + $env:UV_HTTP_RETRIES = "5" + } + if (-not $env:UV_HTTP_TIMEOUT) { + $env:UV_HTTP_TIMEOUT = "180" + } + # ── Create venv (migrate old layout if possible, otherwise fresh) ── # Pass the resolved executable path to uv so it does not re-resolve # a version string back to a conda interpreter. @@ -1565,7 +1657,9 @@ shell.Run cmd, 0, False $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") { + if ($hipOut -match "(?i)gcnArchName") { + # hipinfo can crash after printing gcnArchName (#6043). + # Once the arch is printed, keep the ROCm wheel path. $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 } @@ -1575,8 +1669,13 @@ shell.Run cmd, 0, False } else { $ROCmGpuLabel = "AMD ROCm" } + if ($LASTEXITCODE -ne 0) { + Write-Host " [INFO] hipinfo exited with code $LASTEXITCODE but reported gcnArchName -- treating as ROCm-capable (see #6043)" -ForegroundColor Cyan + } } elseif ($LASTEXITCODE -ne 0) { - # hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected") + # hipinfo ran but returned a HIP runtime error without any gcnArchName + # output (e.g. "no ROCm-capable device detected"), or crashed before + # printing device info. $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 @@ -1816,6 +1915,64 @@ shell.Run cmd, 0, False substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow" return "$baseUrl/cu126" } + + # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── + # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, + # matching setup.ps1's stale-venv parse. + function ConvertTo-TorchFlavorTag { + param([string]$TorchVersion) + if (-not $TorchVersion) { return $null } + if ($TorchVersion -match '\+(cu\d+)') { return $Matches[1] } + if ($TorchVersion -match '\+rocm') { return 'rocm' } + if ($TorchVersion -match '\+cpu') { return 'cpu' } + return 'cpu' + } + + # Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a + # gfx* leaf -> rocm). $null on an unknown leaf (odd mirror) so repair no-ops. + function Get-ExpectedTorchFlavorTag { + param([string]$TorchIndexUrl, [string]$ROCmIndexUrl) + if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' } + if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null } + $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if ($leaf -match '^cu\d+$') { return $leaf } + if ($leaf -eq 'cpu') { return 'cpu' } + if ($leaf -match '^rocm') { return 'rocm' } + if ($leaf -match '^gfx') { return 'rocm' } + return $null + } + + # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses + # ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference. + function Get-InstalledTorchTag { + param([string]$PythonExe) + if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $PythonExe + $psi.Arguments = '-c "import torch; print(torch.__version__)"' + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.UseShellExecute = $false + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + # Drain BOTH streams async, then WaitForExit. A synchronous ReadToEnd() + # before the wait would block forever if a wedged "import torch" never + # closes stdout; leaving the redirected stderr undrained would deadlock a + # child that floods it past the pipe buffer. Async reads let a noisy-but- + # exiting probe finish, while a truly hung one still hits the 30s timeout + # and is killed -- bounded either way. + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + $finished = $proc.WaitForExit(30000) + if (-not $finished) { try { $proc.Kill() } catch {}; return $null } + $torchVer = $outTask.GetAwaiter().GetResult().Trim() + [void]$errTask.GetAwaiter().GetResult() + if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null } + return ConvertTo-TorchFlavorTag $torchVer + } catch { return $null } + } + $TorchIndexUrl = Get-TorchIndexUrl # ===== Windows-on-ARM + NVIDIA GPU -> automatic WSL2 fallback (N1X "RTX Spark" / DGX Spark-class) ===== @@ -2286,21 +2443,21 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below # is --no-deps). All transitive deps are torch-free. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2314,7 +2471,7 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } substep "overlaying unsloth-zoo from git main..." - $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } + $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } if ($zooOverlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) @@ -2327,7 +2484,7 @@ shell.Run cmd, 0, False 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 } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { 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) @@ -2335,7 +2492,7 @@ shell.Run cmd, 0, False } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2347,21 +2504,21 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install no-torch runtime deps" { uv pip install --python $VenvPython --no-deps -r $NoTorchReq } } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2376,7 +2533,7 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } substep "overlaying unsloth-zoo from git main..." - $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } + $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } if ($zooOverlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) @@ -2387,7 +2544,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) @@ -2399,13 +2556,13 @@ shell.Run cmd, 0, False return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } substep "overlaying unsloth-zoo from git main..." - $zooOverlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } + $zooOverlayExit = Invoke-InstallCommandRetry -Label "overlay unsloth-zoo (git main)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth-zoo "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" } if ($zooOverlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit) } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) @@ -2413,6 +2570,50 @@ shell.Run cmd, 0, False } } + # ── Enforce the installed torch flavor matches the detected GPU build ── + # PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv + # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on + # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is + # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* + # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install + # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. + if (-not $SkipTorch) { + $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl + if ($expectedTorchTag -and $expectedTorchTag -ne 'cpu') { + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { + if ($expectedTorchTag -eq 'rocm' -and $ROCmIndexUrl) { + # AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path + # would have force-reinstalled. Repair from the same repo.amd.com index. + $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } + substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec torchvision torchaudio } + if ($torchFixExit -ne 0) { + Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) + } + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + } elseif ($expectedTorchTag -ne 'rocm') { + # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. + substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + if ($torchFixExit -ne 0) { + Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red + return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) + } + $installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython + } + } + # Safety net (incl. AMD): GPU build expected but still CPU -> warn loudly. + if ($installedTorchTag -eq 'cpu') { + Write-Host "" + Write-Host " [WARN] PyTorch is CPU-only but a $expectedTorchTag GPU build was expected for this machine." -ForegroundColor Yellow + Write-Host " [WARN] Training and GPU inference will run on CPU until this is fixed." -ForegroundColor Yellow + Write-Host " [WARN] Re-run this installer, or reinstall the GPU build manually for your GPU." -ForegroundColor Yellow + } + } + } + # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped # for --local: the editable install above already makes _PACKAGE_ROOT in # unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__). diff --git a/install.sh b/install.sh index 59f4c80e9c..a83c0d5181 100755 --- a/install.sh +++ b/install.sh @@ -163,6 +163,40 @@ run_install_cmd() { return $_rc } +# Retry run_install_cmd on transient uv download failures with backoff. Returns +# the last exit code on permanent failure so the set -e rollback trap still fires. +: "${UNSLOTH_INSTALL_RETRIES:=3}" +: "${UNSLOTH_INSTALL_RETRY_DELAY:=3}" +run_install_cmd_retry() { + _ricr_label="$1" + # Sanitize overrides to a default of 3 (a typo must not disable retries; =1 disables). + # Length guard precedes the numeric test so a huge value can't overflow `[ -ge ]`. + # 0?* rejects leading-zero delays ("08"/"09" break the later $((delay*2)) as octal); + # bare "0" stays valid. Bounds: 1..100 retries, 0..3600s base delay. + case "$UNSLOTH_INSTALL_RETRIES" in + ''|*[!0-9]*|0) _ricr_max=3 ;; + *) if [ "${#UNSLOTH_INSTALL_RETRIES}" -le 3 ] && [ "$UNSLOTH_INSTALL_RETRIES" -ge 1 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRIES" -le 100 ] 2>/dev/null; then _ricr_max=$UNSLOTH_INSTALL_RETRIES; else _ricr_max=3; fi ;; + esac + case "$UNSLOTH_INSTALL_RETRY_DELAY" in + ''|*[!0-9]*|0?*) _ricr_delay=3 ;; + *) if [ "${#UNSLOTH_INSTALL_RETRY_DELAY}" -le 4 ] && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -ge 0 ] 2>/dev/null && [ "$UNSLOTH_INSTALL_RETRY_DELAY" -le 3600 ] 2>/dev/null; then _ricr_delay=$UNSLOTH_INSTALL_RETRY_DELAY; else _ricr_delay=3; fi ;; + esac + _ricr_attempt=1 + while :; do + # AND-OR (not `if`) preserves the real failure code: $? after a non-taken + # `if` is 0 in sh/dash/bash, which would break the rollback path. + run_install_cmd "$@" && return 0 + _ricr_rc=$? + if [ "$_ricr_attempt" -ge "$_ricr_max" ]; then + return "$_ricr_rc" + fi + substep "retrying \"$_ricr_label\" after transient failure (attempt $((_ricr_attempt + 1))/$_ricr_max, waiting ${_ricr_delay}s)..." "$C_WARN" + sleep "$_ricr_delay" || true + _ricr_attempt=$((_ricr_attempt + 1)) + _ricr_delay=$((_ricr_delay * 2)) + done +} + # Install bitsandbytes on AMD ROCm hosts. Uses the continuous-release_main # wheel for the ROCm 4-bit GEMV fix (bnb PR #1887, post-0.49.2); bnb <= 0.49.2 # NaNs at decode shape on every AMD GPU. Falls back to PyPI >=0.49.1 if the @@ -1229,6 +1263,10 @@ if (-not \$targetExe) { exit 1 } # AppData\Local (renders blank); a profile-path icon renders everywhere. \$iconDir = Join-Path \$env:USERPROFILE '.unsloth' \$iconPath = Join-Path \$iconDir 'unsloth.ico' +\$preIconHash = \$null +if (Test-Path -LiteralPath \$iconPath) { + try { \$preIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash } catch {} +} if (-not (Test-Path -LiteralPath \$iconPath)) { try { New-Item -ItemType Directory -Force -Path \$iconDir | Out-Null @@ -1244,9 +1282,11 @@ if (Test-Path -LiteralPath \$iconPath) { (Join-Path \$env:APPDATA 'Microsoft\Windows\Start Menu\Programs') ) \$created = @() +\$firstShortcut = \$false foreach (\$dir in \$locations) { if (-not \$dir -or -not (Test-Path \$dir)) { continue } \$linkPath = Join-Path \$dir '$_css_lnk_name_ps' + if (-not (Test-Path -LiteralPath \$linkPath)) { \$firstShortcut = \$true } \$shortcut = \$WshShell.CreateShortcut(\$linkPath) \$shortcut.TargetPath = \$targetExe \$shortcut.Arguments = '$_css_sc_args_ps' @@ -1255,27 +1295,43 @@ foreach (\$dir in \$locations) { \$shortcut.Save() \$created += \$linkPath } -# Force Explorer to re-read EACH new shortcut's icon so it renders immediately -# instead of a stale/blank (generic) icon. The reliable, NON-disruptive fix -# (no explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM, -# SHCNF_PATHW, ) -- the global SHCNE_ASSOCCHANGED alone does not recover a -# stale item. Also clear the on-disk icon cache for heavier staleness. -try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {} -try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {} +\$iconChanged = \$false +if (\$hasIcon) { + if (-not \$preIconHash) { + \$iconChanged = \$true + } else { + try { + \$postIconHash = (Get-FileHash -LiteralPath \$iconPath -Algorithm SHA256).Hash + \$iconChanged = (\$postIconHash -ne \$preIconHash) + } catch { \$iconChanged = \$true } + } +} elseif (\$preIconHash) { + \$iconChanged = \$true +} +# Per-item refresh always (cheap, non-disruptive) so the rewritten .lnk renders +# immediately instead of a stale/blank (generic) icon. The reliable fix (no +# explorer restart) is a PER-ITEM SHChangeNotify(SHCNE_UPDATEITEM, SHCNF_PATHW, +# ) -- the global SHCNE_ASSOCCHANGED alone does not recover a stale item. try { Add-Type -Namespace UnslothShell -Name IconRefresh -MemberDefinition '[System.Runtime.InteropServices.DllImport("shell32.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)] public static extern void SHChangeNotify(int e, uint f, string a, System.IntPtr b);' -ErrorAction SilentlyContinue foreach (\$p in \$created) { try { [UnslothShell.IconRefresh]::SHChangeNotify(0x00002000, 0x0005, \$p, [System.IntPtr]::Zero) } catch {} } [UnslothShell.IconRefresh]::SHChangeNotify(0x08000000, 0, \$null, [System.IntPtr]::Zero) } catch {} -# Win11 Start Menu keeps its own tile-icon cache (preserve start2.bin). -try { - \$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState' - if (Test-Path -LiteralPath \$smeh) { - Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue - Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue - Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue - } -} catch {} +# Heavier on-disk icon-cache clear + StartMenuExperienceHost tile rebuild +# (preserve start2.bin) only on first install or a real icon change, so a no-op +# WSL reinstall does not run a dropper-like clear-cache + kill cluster each time. +if (\$created.Count -gt 0 -and (\$firstShortcut -or \$iconChanged)) { + try { & "\$env:SystemRoot\System32\ie4uinit.exe" -ClearIconCache } catch {} + try { & "\$env:SystemRoot\System32\ie4uinit.exe" -show } catch {} + try { + \$smeh = Join-Path \$env:LOCALAPPDATA 'Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState' + if (Test-Path -LiteralPath \$smeh) { + Get-ChildItem -LiteralPath \$smeh -Filter 'TileCache_*' -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue + Remove-Item -LiteralPath (Join-Path \$smeh 'StartUnifiedTileModelCache.dat') -Force -ErrorAction SilentlyContinue + Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue + } + } catch {} +} WSLPS1_EOF # Convert WSL path to Windows path for powershell.exe @@ -1458,12 +1514,19 @@ fi # ── Install uv ── tauri_log "STEP" "Installing uv package manager" -UV_MIN_VERSION="0.7.22" +UV_MIN_VERSION="0.8.16" # When bytecode compilation is enabled, large installs can exceed uv's 60s default on slow machines. Default to 180s, preserving overrides ("0" disables). : "${UV_COMPILE_BYTECODE_TIMEOUT:=180}" export UV_COMPILE_BYTECODE_TIMEOUT +# uv >= 0.8.16 retries HTTP/2 streaming body errors; raise retries and read +# timeout for large wheel downloads. ":=" preserves any user override. +: "${UV_HTTP_RETRIES:=5}" +export UV_HTTP_RETRIES +: "${UV_HTTP_TIMEOUT:=180}" +export UV_HTTP_TIMEOUT + version_ge() { # returns 0 if $1 >= $2 _a=$1 @@ -1929,6 +1992,45 @@ get_torch_index_url() { else echo "$_base/cpu"; fi } +# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── +# torch.__version__ ($1) -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu. +_torch_flavor_tag() { + case "$1" in + *+cu[0-9]*) printf '%s\n' "$1" | sed -n 's/.*+\(cu[0-9][0-9]*\).*/\1/p' ;; + *+rocm*) echo "rocm" ;; + *+cpu*) echo "cpu" ;; + "") echo "" ;; + *) echo "cpu" ;; + esac +} + +# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> +# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. +_expected_torch_flavor_tag() { + _u="${1%/}" + _leaf="${_u##*/}" + case "$_leaf" in + cu[0-9]*) echo "$_leaf" ;; + cpu) echo "cpu" ;; + rocm*|gfx*) echo "rocm" ;; + *) echo "" ;; + esac +} + +# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv +# resolves (torch + every transitive dep) via --index-url -- the same URLs the +# fresh-install paths above already use -- so a stale wheel is auto-repairable. +# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. +_torch_index_repairable() { + _u="${1%/}" + _leaf="${_u##*/}" + case "$_leaf" in + cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; + *) echo "no" ;; + esac +} + get_radeon_wheel_url() { # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/, @@ -2422,28 +2524,28 @@ if [ "$_MIGRATED" = true ]; then # PyPI metadata still declares torch as a hard dep), then install # runtime deps (typer, safetensors, transformers, etc.) with --no-deps # to prevent transitive torch resolution. - run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ + run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.7" unsloth-zoo + "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. - run_install_cmd "install pydantic (with deps for compatible core)" \ + run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then - run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" + run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi else - run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.7" unsloth-zoo + "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" fi @@ -2452,13 +2554,13 @@ if [ "$_MIGRATED" = true ]; then # fresh reinstall. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" # Repair ROCm torch if overwritten during migrated install _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." - run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" \ --force-reinstall @@ -2584,7 +2686,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then 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" - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" else @@ -2596,30 +2698,30 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # filelock / sympy / networkx which are not in the # Radeon listing. if [ -n "$_tri_whl" ]; then - run_install_cmd "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install triton + PyTorch" uv pip install --python "$_VENV_PY" \ --find-links "$_RADEON_BASE_URL" \ "$_tri_whl" "$_torch_whl" "$_tv_whl" "$_ta_whl" else - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ --find-links "$_RADEON_BASE_URL" \ "$_torch_whl" "$_tv_whl" "$_ta_whl" fi fi else substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." - run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). @@ -2628,7 +2730,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # which is only useful once torch is present for training. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" ;; esac @@ -2639,31 +2741,31 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ + run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.7" unsloth-zoo + "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" # Same pydantic-with-deps trick as the migrated branch. - run_install_cmd "install pydantic (with deps for compatible core)" \ + run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then - run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" + run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.7" unsloth-zoo + run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ + --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" elif [ -n "${UNSLOTH_INSTALL_REF:-}" ] && [ "${UNSLOTH_INSTALL_REF}" != "main" ] && [ "$PACKAGE_NAME" = "unsloth" ]; then @@ -2675,7 +2777,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then --upgrade-package unsloth --upgrade-package unsloth-zoo \ "unsloth @ git+https://github.com/unslothai/unsloth@${UNSLOTH_INSTALL_REF}" unsloth-zoo else - run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ --upgrade-package unsloth -- "$PACKAGE_NAME" fi # aarch64 + NVIDIA (DGX Spark / GB10 / N1X): unsloth's x86_64-oriented cuXXX @@ -2696,11 +2798,11 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. if [ "$SKIP_TORCH" = false ]; then case "$TORCH_INDEX_URL" in - */rocm*) + */rocm*|*/gfx*) _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) if [ -z "$_has_hip" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." - run_install_cmd "repair ROCm torch" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" \ --force-reinstall @@ -2713,15 +2815,50 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.6.7" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." - run_install_cmd "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ + run_install_cmd_retry "overlay unsloth-zoo (git main)" uv pip install --python "$_VENV_PY" \ --no-deps --reinstall-package unsloth-zoo \ "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME" + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" --torch-backend=auto -- "$PACKAGE_NAME" + fi +fi + +# ── Enforce the installed torch flavor matches the detected GPU build ── +# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv +# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on +# CPU. Reinstall the right wheel triplet when a GPU build is expected; if it +# can't be reinstalled, warn loudly. --no-torch / CPU-only / macOS: no-op. +if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then + _expected_torch_tag=$(_expected_torch_flavor_tag "$TORCH_INDEX_URL") + # Only act when a GPU build is expected (cuXXX / rocm); cpu and unknown skip. + if [ -n "$_expected_torch_tag" ] && [ "$_expected_torch_tag" != "cpu" ]; then + _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) + _installed_torch_tag="" + [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") + # Repair when flavor is wrong AND the index is plain --index-url reinstallable + # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. + if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ + && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then + substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." + run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ + "$TORCH_CONSTRAINT" torchvision torchaudio \ + --index-url "$TORCH_INDEX_URL" \ + --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio + _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) + _installed_torch_tag="" + [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") + fi + # Safety net (incl. AMD/WSL): GPU build expected but still CPU -> warn loudly. + if [ "$_installed_torch_tag" = "cpu" ]; then + substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" + substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" + substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + fi fi fi diff --git a/pyproject.toml b/pyproject.toml index 17293ae4fa..d59293017e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,7 +72,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.6", "wheel>=0.42.0", "packaging", "numpy", @@ -93,7 +93,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.6", "torchvision", "unsloth[triton]", ] @@ -583,7 +583,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.5", + "unsloth_zoo>=2026.6.6", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index fe90afa7e6..c1d156d40a 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -412,7 +412,7 @@ BLOCKED_NPM_VERSIONS: dict[str, set[str]] = { "@uipath/functions-tool": {"1.0.1"}, "@uipath/access-policy-sdk": {"0.3.1"}, "@uipath/platform-tool": {"1.0.1"}, - # Mini Shai-Hulud May-12 wave: @mistralai/* (npm) — separate from PyPI mistralai + # Mini Shai-Hulud May-12 wave: @mistralai/* (npm), separate from PyPI mistralai # (https://www.aikido.dev/blog/mini-shai-hulud-is-back-tanstack-compromised). "@mistralai/mistralai": {"2.2.2", "2.2.3", "2.2.4"}, "@mistralai/mistralai-gcp": {"1.7.1", "1.7.2", "1.7.3"}, @@ -916,6 +916,204 @@ def _evidence( LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare") +# ───────────────────────────────────────────────────────────────────── +# Code-only scanning for JS/TS sources. Blank `//` and `/* */` comments +# before matching (the top FP source: scary strings in JSDoc/changelog +# comments), tracking string/template/regex context so a `//` inside +# "http://..." is not mistaken for a comment. Strings are NOT blanked +# (droppers hide payloads there). Fail open on lexer confusion: the raw +# text is still scanned. JS sibling of scan_packages.py::_strip_noncode. +# ───────────────────────────────────────────────────────────────────── +_JS_FAMILY_SUFFIXES = (".js", ".mjs", ".cjs", ".ts", ".tsx", ".jsx") + +# Keywords after which a `/` begins a regex literal (not division). +_REGEX_PRECEDING_KEYWORDS = frozenset( + { + "return", + "typeof", + "instanceof", + "in", + "of", + "new", + "delete", + "void", + "throw", + "yield", + "await", + "do", + "else", + "case", + } +) +_IDENT_CHARS = frozenset("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_$") + + +def _slash_is_regex(prev_tok: str) -> bool: + """Disambiguate a lone ``/``: regex literal vs division operator. + + Biased toward regex when ambiguous -- regex state never blanks, so a + wrong guess only costs FP reduction (or a fail-open), never a missed + detection. + """ + if prev_tok == "": + return True # start of file -> expression position + if prev_tok in _REGEX_PRECEDING_KEYWORDS: + return True + last = prev_tok[-1] + if last.isalnum() or last in "_$)]": + return False # previous token ends a value -> division + return True # operators, punctuation, `{`, `}` -> regex (safe bias) + + +def _strip_js_noncode(text: str) -> str: + """Blank JS/TS comments, preserving byte geometry. Fail-open on confusion.""" + if "//" not in text and "/*" not in text: + return text # nothing to strip + n = len(text) + out = list(text) + nl = ("\n", "\r") + + def _blank(a: int, b: int) -> None: + for k in range(a, b): + if out[k] not in nl: + out[k] = " " + + state = "code" + prev_tok = "" + tmpl_stack: list[str] = [] + i = 0 + try: + while i < n: + c = text[i] + nxt = text[i + 1] if i + 1 < n else "" + if state == "code": + if c == "/" and nxt == "/": + start = i + i += 2 + while i < n and text[i] not in nl: + i += 1 + _blank(start, i) + continue + if c == "/" and nxt == "*": + start = i + i += 2 + closed = False + while i < n: + if text[i] == "*" and i + 1 < n and text[i + 1] == "/": + i += 2 + closed = True + break + i += 1 + if not closed: + return text # unterminated block comment + _blank(start, i) + continue + if c == "'": + state = "sq" + i += 1 + continue + if c == '"': + state = "dq" + i += 1 + continue + if c == "`": + state = "tmpl" + i += 1 + continue + if c == "/": + if _slash_is_regex(prev_tok): + state = "regex" + i += 1 + continue + prev_tok = "/" + i += 1 + continue + if c.isspace(): + i += 1 + continue + if c in _IDENT_CHARS: + j = i + while j < n and text[j] in _IDENT_CHARS: + j += 1 + prev_tok = text[i:j] + i = j + continue + if c == "}" and tmpl_stack: + state = tmpl_stack.pop() + i += 1 + continue + prev_tok = c + i += 1 + continue + elif state in ("sq", "dq"): + q = "'" if state == "sq" else '"' + if c == "\\": + i += 2 + continue + if c == q: + state = "code" + prev_tok = "_v" + i += 1 + continue + if c in nl: + return text # unterminated string literal + i += 1 + continue + elif state == "tmpl": + if c == "\\": + i += 2 + continue + if c == "`": + state = "code" + prev_tok = "_v" + i += 1 + continue + if c == "$" and nxt == "{": + tmpl_stack.append("tmpl") + state = "code" + prev_tok = "{" + i += 2 + continue + i += 1 + continue + elif state == "regex": + if c == "\\": + i += 2 + continue + if c == "[": + state = "regex_cc" + i += 1 + continue + if c == "/": + state = "code" + prev_tok = "_v" + i += 1 + continue + if c in nl: + return text # unterminated regex literal + i += 1 + continue + elif state == "regex_cc": + if c == "\\": + i += 2 + continue + if c == "]": + state = "regex" + i += 1 + continue + if c in nl: + return text + i += 1 + continue + else: + return text + if state != "code" or tmpl_stack: + return text # unterminated construct -> fail open + except Exception: + return text + return "".join(out) + + def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: findings: list[Finding] = [] try: @@ -1042,6 +1240,14 @@ def _host_in_outbound_context(text: str, host: str) -> bool: def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: findings: list[Finding] = [] + # Code-only scanning for JS/TS sources: blank comments before matching so + # an IOC host / `eval(atob)` example / campaign marker quoted in a comment + # cannot manufacture a false positive. Assigned string literals (where real + # droppers hide base64 payloads) are preserved. Non-JS text (json/yaml/sh/ + # py/html) is scanned as-is -- this lexer only understands JS comments. + if rel.lower().endswith(_JS_FAMILY_SUFFIXES): + text = _strip_js_noncode(text) + # IOC substrings (literal, case-sensitive). for needle, (sev, why) in KNOWN_IOC_STRINGS.items(): if needle in text: @@ -1236,6 +1442,131 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N pass +# ───────────────────────────────────────────────────────────────────── +# Baseline allowlist: triaged known-good HIGH/CRITICAL findings so the gate +# can enforce without red-failing on rare legitimate-library behavior. +# Matched on ``(normalized package, package-relative path, pattern)`` -- not +# evidence text -- so a version bump does not reopen a finding, but a *new* +# kind of finding in a listed file is a different pattern and still fails. +# Mirrors scan_packages.py. Regenerate with ``--write-baseline``. +# ───────────────────────────────────────────────────────────────────── + +_DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json") + +# Bumped when the entry-key semantics change. v2 keys on the package-relative +# path; v1 stored only a basename, so a v1 entry could suppress a same-named file +# in a different directory. A pre-v2 baseline with entries is ignored (fail +# closed) rather than mis-applied. +_BASELINE_SCHEMA_VERSION = 2 + + +def _norm_pkg_name(display: str) -> str: + """``@scope/pkg@1.2.3`` / ``pkg@1.2.3`` -> name without the version. + + The version is the LAST ``@``-separated field; a leading ``@`` (scope) + is preserved. Lower-cased (npm names are case-insensitive). Sentinels + like ```` / ```` pass through unchanged. + """ + s = (display or "").strip() + at = s.rfind("@") + if at > 0: # >0 so a leading @scope is not treated as the version sep + s = s[:at] + return s.lower() + + +_NPM_TARBALL_ROOT = "package/" + + +def _relpath_in_package(filename: str) -> str: + """Path within the published package, stable across version bumps. npm + tarballs root every file at ``package/``; strip it so the key is the real + source path (``dist/index.js``) and a new file with the same basename in a + different directory is not silently suppressed.""" + f = (filename or "").replace("\\", "/") + return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f + + +def _finding_key(f: Finding) -> tuple[str, str, str]: + """Stable allowlist key: normalized package, package-relative path, pattern.""" + return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern) + + +def _load_baseline(path: str) -> set[tuple[str, str, str]]: + """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" + try: + with open(path, "r", encoding = "utf-8") as fh: + data = json.load(fh) + except FileNotFoundError: + return set() + except (OSError, json.JSONDecodeError) as exc: + print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) + return set() + entries = data.get("entries", []) + if entries and data.get("version") != _BASELINE_SCHEMA_VERSION: + print( + f" [WARN] baseline schema v{data.get('version')} predates package-relative " + f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.", + file = sys.stderr, + ) + return set() + keys: set[tuple[str, str, str]] = set() + for e in entries: + try: + keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"])) + except (KeyError, TypeError): + continue + return keys + + +def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int: + """Persist at-or-above-threshold findings as an allowlist for triage.""" + entries = [] + seen: set[tuple[str, str, str]] = set() + for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)): + if _SEVERITY_RANK[f.severity] > threshold_rank: + continue + key = _finding_key(f) + if key in seen: + continue + seen.add(key) + entries.append( + { + "package": _norm_pkg_name(f.package), + "file": _relpath_in_package(f.filename), + "pattern": f.pattern, + "severity": f.severity, + "evidence": (f.evidence or f.detail)[:240], + } + ) + doc = { + "_comment": ( + "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL " + "finding manually judged benign. Matched on (package, " + "package-relative path, pattern); evidence/severity are for review " + "only. Regenerate with --write-baseline AFTER reviewing every line." + ), + "version": _BASELINE_SCHEMA_VERSION, + "entries": entries, + } + with open(path, "w", encoding = "utf-8") as fh: + json.dump(doc, fh, indent = 2, sort_keys = False) + fh.write("\n") + print(f" Wrote {len(entries)} baseline entr(y/ies) to {path}") + return len(entries) + + +def _partition_baseline( + findings: list[Finding], baseline: set[tuple[str, str, str]] +) -> tuple[list[Finding], list[Finding]]: + """Split findings into (active, suppressed) by allowlist membership.""" + if not baseline: + return list(findings), [] + active, suppressed = [], [] + for f in findings: + (suppressed if _finding_key(f) in baseline else active).append(f) + return active, suppressed + + def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser( description = "Pre-install npm tarball content scanner.", @@ -1263,6 +1594,30 @@ def main(argv: list[str] | None = None) -> int: "Medium and below print but exit 0." ), ) + parser.add_argument( + "--baseline", + metavar = "FILE", + default = None, + help = ( + "Allowlist JSON of triaged known-good findings to suppress. " + "Defaults to scan_npm_packages_baseline.json next to this script " + "if present." + ), + ) + parser.add_argument( + "--no-baseline", + action = "store_true", + help = "Ignore the auto-discovered baseline allowlist.", + ) + parser.add_argument( + "--write-baseline", + metavar = "FILE", + default = None, + help = ( + "Write the current at/above-threshold findings to FILE as an " + "allowlist, then exit 0. Review every entry before committing it." + ), + ) args = parser.parse_args(argv) lockfile = Path(args.lockfile).resolve() @@ -1341,7 +1696,46 @@ def main(argv: list[str] | None = None) -> int: "critical": CRITICAL, }[args.fail_on] threshold_rank = _SEVERITY_RANK[threshold] - blocking = [f for f in all_findings if _SEVERITY_RANK[f.severity] <= threshold_rank] + + # --write-baseline: persist the full current at/above-threshold set as the + # new allowlist (ignoring any loaded baseline), then exit 0. A hard error + # means the scan was incomplete, so warn -- a baseline baked from a partial + # run would silently allow whatever failed to download. + if args.write_baseline: + if hard_errors: + print( + f" [WARN] {len(hard_errors)} hard error(s): baseline may be " + "incomplete (some packages did not scan).", + file = sys.stderr, + ) + _write_baseline(args.write_baseline, all_findings, threshold_rank) + return 0 + + # Baseline allowlist: suppress triaged, known-good findings so the CI gate + # can be enforcing without red-failing on legitimate-library noise. + if args.no_baseline: + baseline_path = None + elif args.baseline: + baseline_path = args.baseline + elif os.path.isfile(_DEFAULT_BASELINE_PATH): + baseline_path = _DEFAULT_BASELINE_PATH + else: + baseline_path = None + baseline = _load_baseline(baseline_path) if baseline_path else set() + active, suppressed = _partition_baseline(all_findings, baseline) + + if suppressed: + crit_s = sum(1 for f in suppressed if f.severity == CRITICAL) + high_s = sum(1 for f in suppressed if f.severity == HIGH) + print( + f"\n[scan-npm] {len(suppressed)} finding(s) suppressed by baseline " + f"{baseline_path} ({crit_s} CRITICAL, {high_s} HIGH).", + flush = True, + ) + + # Exit code: 1 on a hard error, or a NON-baselined finding at/above the + # threshold. This is the signal CI gates on once the baseline is clean. + blocking = [f for f in active if _SEVERITY_RANK[f.severity] <= threshold_rank] if hard_errors or blocking: if blocking: print( diff --git a/scripts/scan_npm_packages_baseline.json b/scripts/scan_npm_packages_baseline.json new file mode 100644 index 0000000000..61d8e74023 --- /dev/null +++ b/scripts/scan_npm_packages_baseline.json @@ -0,0 +1,5 @@ +{ + "_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.", + "version": 2, + "entries": [] +} diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py index 861b35617b..4be9fc5efb 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -33,10 +33,24 @@ Examples: python scan_packages.py --fix -r requirements.txt python scan_packages.py --fix --max-search 20 -r requirements.txt + # Triage to a baseline once, then gate on anything NEW + python scan_packages.py -r requirements.txt --write-baseline scripts/scan_packages_baseline.json + python scan_packages.py -r requirements.txt # auto-loads the baseline, exits 0 if only baselined findings remain + +False positives: + .py files are scanned code-only: comments and bare docstrings/doctests are + blanked before pattern matching (line numbers preserved), so prose, usage + examples and `>>>` doctests cannot trip a finding. Residual findings that + are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored + test fixture) are suppressed via a reviewed baseline allowlist, matched on + (package, basename(file), check). A NEW kind of finding in an already-listed + file is a different check and still fails. This mirrors the Hugging Face Hub + approach (ClamAV/picklescan: low-FP, signature/structural, surface status). + Exit codes: - 0 -- no CRITICAL or HIGH findings - 1 -- CRITICAL or HIGH findings detected - 2 -- no packages specified + 0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline) + 1 -- non-baselined CRITICAL or HIGH findings detected + 2 -- no packages specified, or scan incomplete (pip download failure) """ import argparse @@ -50,6 +64,8 @@ import subprocess import sys import tarfile import tempfile +import tokenize +import urllib.parse import urllib.request import zipfile from dataclasses import dataclass, field @@ -213,17 +229,26 @@ RE_ARCHIVE_STAGING = re.compile( ) # Anti-analysis / sandbox evasion / debugger detection +# NB: deliberately does NOT include a bare ``platform.system() ... Linux/Windows +# /Darwin`` branch. Under re.DOTALL that matched across the whole file -- any +# cross-platform library (typer, packaging, pandas, pymupdf, ...) trips it -- so +# it had ~zero precision and only generated false positives. OS detection alone +# is not an anti-analysis signal; the debugger/VM/long-sleep signals below are. RE_ANTI_ANALYSIS = re.compile( r"\bptrace\b" r"|\bsys\s*\.\s*gettrace\s*\(" r"|\bsys\s*\.\s*settrace\b" r"|\bTracerPid\b" - r"|\b/proc/self/status\b" + # /proc/self/status is read to scrape TracerPid for anti-debug. A leading + # \b here is unsatisfiable (\b never holds between a non-word boundary and + # "/"), so the old pattern was dead; a lookbehind that only forbids a + # preceding word char or path separator lets `open("/proc/self/status")` + # and `cat /proc/self/status` match while avoiding mid-path partials. + r"|(? list[Finding]: return findings +# A STRING after one of these tokens (and before a NEWLINE) is a bare +# docstring/doctest/prose statement -- the dominant FP source -- so we blank it. +# A string after `=` or `(` is real code and is never blanked. +_LINE_START_TOKENS = frozenset({tokenize.NEWLINE, tokenize.NL, tokenize.INDENT, tokenize.DEDENT}) + + +def _is_fstring(tok_string: str) -> bool: + """True if a STRING token is an f-string (3.10/3.11 emit one STRING token). + + A bare f-string statement evaluates its expressions at import, so unlike an + inert docstring it must never be blanked. + """ + q = min((tok_string.find(c) for c in "'\"" if c in tok_string), default = -1) + return q > 0 and "f" in tok_string[:q].lower() + + +def _strip_noncode(content: str, blank_comments: bool = True) -> str: + """Blank comments and bare docstrings so IOC patterns see code only. + + Removed regions become spaces (newlines kept) so line numbers stay exact for + _extract_evidence. Fails open on tokenizer errors (the raw text is still + fully scanned, so a real detection is never lost). ``blank_comments=False`` + keeps comments (only strings/docstrings blanked) to isolate the span that + exec() could actually run. + """ + try: + toks = list(tokenize.generate_tokens(io.StringIO(content).readline)) + except (tokenize.TokenError, IndentationError, SyntaxError, ValueError): + return content + + spans: list[tuple[int, int, int, int]] = [] # (srow, scol, erow, ecol) + prev_significant = tokenize.NEWLINE # start-of-file behaves like a new line + n = len(toks) + for i, tok in enumerate(toks): + ttype = tok.type + if ttype == tokenize.COMMENT: + if blank_comments: + spans.append((*tok.start, *tok.end)) + continue # transparent; never advances prev_significant + if ( + ttype == tokenize.STRING + and prev_significant in _LINE_START_TOKENS + and not _is_fstring(tok.string) # f-strings execute; never blank them + ): + # Bare string only if it is the whole statement: next significant + # token must close the logical line. + j = i + 1 + while j < n and toks[j].type in (tokenize.COMMENT, tokenize.NL): + j += 1 + if j < n and toks[j].type == tokenize.NEWLINE: + spans.append((*tok.start, *tok.end)) + prev_significant = ttype + continue + if ttype in ( + tokenize.NL, + tokenize.NEWLINE, + tokenize.INDENT, + tokenize.DEDENT, + tokenize.ENCODING, + ): + prev_significant = ttype + continue + prev_significant = ttype + + if not spans: + return content + + buf = content.splitlines(keepends = True) + for srow, scol, erow, ecol in spans: + for row in range(srow, erow + 1): + line = buf[row - 1] + if line.endswith("\n"): + body, nl = line[:-1], "\n" + elif line.endswith("\r"): + body, nl = line[:-1], "\r" + else: + body, nl = line, "" + start = scol if row == srow else 0 + end = ecol if row == erow else len(body) + end = min(end, len(body)) + if start < end: + body = body[:start] + (" " * (end - start)) + body[end:] + buf[row - 1] = body + nl + return "".join(buf) + + +# Payload carriers that are suspicious when hidden in a blanked region (a +# docstring/string) of a file that can dynamically execute strings. +_HIDDEN_PAYLOAD_PATTERNS = ( + (RE_LARGE_BLOB, "large base64 blob"), + (RE_EMBEDDED_KEYS, "embedded key material"), + (RE_MAY12_IOC, "Shai-Hulud IOC string"), + (RE_OBFUSCATION, "marshal/compile/obfuscation"), +) + + +def _hidden_payload_findings( + original: str, stripped: str, filename: str, package: str +) -> list[Finding]: + """Flag payloads that live only in the blanked (docstring/string) region of + a file that contains exec/eval. Such a string is invisible to code-only + scanning yet ``exec(__doc__)`` / ``exec()`` could still run it.""" + if not RE_EXEC_EVAL.search(stripped): + return [] + # Only docstrings/strings run via exec(__doc__)/exec(); comments cannot. + # Isolate that span: keep comments as real code, take what string-blanking + # removed (length-preserved, so offsets stay exact for _extract_evidence). + code = _strip_noncode(original, blank_comments = False) + removed = "".join(o if o != s else " " for o, s in zip(original, code)) + out = [] + + def _hidden(pat): + # Carrier present in a blanked region but NOT in real code. A carrier in + # real code is already caught by the normal check, so restricting to + # blanked-only avoids re-flagging legitimate in-code constants. + return bool(pat.search(removed)) and not pat.search(stripped) + + for pat, label in _HIDDEN_PAYLOAD_PATTERNS: + if _hidden(pat): + out.append( + Finding( + HIGH, + package, + filename, + "exec/eval with payload hidden in a docstring/string", + f"{label}: {_extract_evidence(removed, pat)}", + ) + ) + # Fetch-then-run dropper: a network call AND an os/subprocess exec that both + # live in the blanked region. Search the removed span directly (not "absent + # from real code") so a benign visible network/subprocess call cannot mask + # the docstring payload. + if RE_NETWORK.search(removed) and RE_SUBPROCESS.search(removed): + out.append( + Finding( + HIGH, + package, + filename, + "exec/eval with hidden network+exec payload", + f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}", + ) + ) + return out + + def check_py_file(content: str, filename: str, package: str) -> list[Finding]: """Run all .py-specific checks.""" - findings = [] + # Code-only scanning: strip comments/docstrings up front so prose, doctests + # and usage examples cannot manufacture false positives. Aligns with the + # Hugging Face Hub model (ClamAV/picklescan: low-FP, signature/structural). + original = content + content = _strip_noncode(content) + findings = _hidden_payload_findings(original, content, filename, package) basename = os.path.basename(filename) is_setup = basename in ("setup.py", "setup.cfg") is_init = basename == "__init__.py" @@ -937,7 +1112,13 @@ def _extract_evidence( pattern: re.Pattern, max_matches: int = 3, ) -> str: - """Pull matching lines as evidence snippets.""" + """Pull matching lines as evidence snippets. + + Falls back to a whole-content search when the pattern only matches across + line boundaries (several IOC regexes use ``re.DOTALL``). Without this an + anti-analysis / archive-staging finding could report empty evidence, making + the baseline entry impossible to review. + """ lines = content.splitlines() matches = [] for i, line in enumerate(lines, 1): @@ -948,7 +1129,17 @@ def _extract_evidence( matches.append(f"L{i}: {snippet}") if len(matches) >= max_matches: break - return " | ".join(matches) if matches else "" + if matches: + return " | ".join(matches) + # Multiline (DOTALL) match: report the line where the match begins. + m = pattern.search(content) + if m: + line_no = content.count("\n", 0, m.start()) + 1 + snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else "" + if len(snippet) > 160: + snippet = snippet[:160] + "..." + return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: " + return "" # Non-Python checkers @@ -1390,6 +1581,394 @@ _PIP_DOWNLOAD_PIN_FLAGS = [ _RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]") +# sdist fallback. `--only-binary :all:` never builds an sdist (no setup.py +# exec), but a wheel-less project then can't be fetched at all and one such +# package fails the whole --with-deps resolve (exit 2) -- a coverage hole. So on +# resolve failure we drop to per-spec and fetch any sdist-only package's raw +# tarball from the PyPI JSON API for scan_archive() to read statically: no pip, +# no build, same no-exec guarantee. Transport failures are still exit 2; only +# "no wheel" is downgraded to a direct fetch. + +# How many levels of indirect-dep recovery to chase (a wheel dep whose own child +# is sdist-only, and so on). Bounded with dedup so recovery always terminates. +_MAX_DEP_FOLLOWUP_DEPTH = 2 +_SDIST_DOWNLOAD_TIMEOUT = 180 +# Never fetch an archive larger than we would be willing to scan (iter_archive_files cap). +_MAX_SDIST_BYTES = HARD_MAX_TOTAL_BYTES +# Direct sdist bytes only ever come from PyPI's own CDN; refuse anything else. +_TRUSTED_PYPI_HOSTS = frozenset({"files.pythonhosted.org", "pypi.org", "pypi.python.org"}) + + +def _spec_pin_version(spec: str) -> str | None: + """Return the ``==X.Y.Z`` pin from a spec, or None if unpinned.""" + m = _RE_PYPI_SPEC_VERSION.search(spec) + return m.group(1) if m else None + + +def _pypi_json(name: str, version: str | None = None) -> dict | None: + """Fetch PyPI metadata JSON (read-only HTTPS GET, no exec); None on error. + With ``version`` it fetches that release's document, whose ``requires_dist`` + is accurate for the pin (the project-level doc describes only the latest).""" + url = "https://pypi.org/pypi/" + urllib.parse.quote(name, safe = "") + if version: + url += "/" + urllib.parse.quote(version, safe = "") + url += "/json" + try: + req = urllib.request.Request(url, headers = {"Accept": "application/json"}) + with urllib.request.urlopen(req, timeout = 30) as resp: + if getattr(resp, "status", 200) != 200: + return None + data = resp.read(16 * 1024 * 1024) # metadata is small; cap regardless + return json.loads(data.decode("utf-8", errors = "replace")) + except Exception: + return None + + +def _release_files(meta: dict, version: str | None) -> list[dict]: + """Files for a pinned version, else the latest release's. A pin that is + absent or empty returns [] (never the latest) so a yanked/bad pin fails + closed instead of a different artifact being scanned in its place.""" + if version is not None: + return meta.get("releases", {}).get(version) or [] + return meta.get("urls", []) or [] + + +def _release_has_wheel(meta: dict, version: str | None) -> bool: + """True if the (pinned or latest) release publishes any bdist_wheel.""" + return any(f.get("packagetype") == "bdist_wheel" for f in _release_files(meta, version)) + + +def _is_trusted_pypi_url(url: str) -> bool: + """Only download sdist bytes from PyPI's own hosts, over HTTPS.""" + try: + parsed = urllib.parse.urlparse(url) + except Exception: + return False + return parsed.scheme == "https" and parsed.hostname in _TRUSTED_PYPI_HOSTS + + +_MARKER_ENV_VARS = ( + "sys_platform", + "platform_system", + "platform_machine", + "platform_release", + "platform_version", + "platform_python_implementation", + "os_name", + "python_version", + "python_full_version", + "implementation_name", + "implementation_version", +) + + +def _marker_holds_by_default(marker: str) -> bool: + """Keep (scan) a dep unless its marker is purely ``extra``-gated. The scanner + runs on one OS/Python but a package may be installed on another, so a marker + that can be true on a different target (``sys_platform == 'win32'``, + ``python_version == '3.13'``) is always kept; only a marker depending solely + on ``extra`` and false with no extra requested is dropped. Conservative: on + any uncertainty, keep (over-scan, never silently skip).""" + m = marker.strip() + if not m or "extra" not in m: + return True # no extra gate: installed by default on some target -> scan + if any(v in m for v in _MARKER_ENV_VARS): + return True # also platform/python gated: true on some target -> scan + # Pure extra marker: decide by evaluating with no extra requested. + try: + from packaging.markers import Marker, default_environment + + env = default_environment() + env["extra"] = "" + return bool(Marker(m).evaluate(env)) + except Exception: + # packaging missing/unparseable: drop only a pure positive extra-equality. + return re.fullmatch(r"\s*extra\s*==\s*['\"][^'\"]+['\"]\s*", m) is None + + +def _requires_dist_names(meta: dict) -> list[str]: + """Transitive dep specs (name + version specifier) from metadata, to recover + a sdist-only package's tree. The specifier is kept so a pinned malicious + version is fetched, not latest. Drops deps whose marker cannot hold for a + default install.""" + info = meta.get("info", {}) or {} + reqs = info.get("requires_dist") or [] + specs: list[str] = [] + for r in reqs: + if not isinstance(r, str): + continue + head = r + if ";" in r: + head, marker = r.split(";", 1) + if not _marker_holds_by_default(marker): + continue + if not _RE_NAME.match(head.strip()): + continue + # "torch (>=1.10)" / "torch >=1.10" -> "torch>=1.10" (pip-friendly). + specs.append(re.sub(r"\s+", "", head).replace("(", "").replace(")", "")) + return specs + + +def _requires_dist_for( + name: str, + version: str | None, + project_meta: dict, + errors: list[str] | None = None, +) -> list[str]: + """Declared deps for the pinned version, read from that release's metadata + (its ``requires_dist`` can differ from latest). Unpinned uses the + project-level (latest) document. A pinned version whose own metadata cannot + be fetched returns [] (never latest's deps) and, when ``errors`` is given, + records an incomplete-scan error so a partial tree is not read as "no deps".""" + if not version: + return _requires_dist_names(project_meta) + vmeta = _pypi_json(name, version) + if vmeta is None: + msg = f"metadata fetch failed for pinned {name}=={version}; dependency scan incomplete" + if errors is None: + print(f" [WARN] {msg}", file = sys.stderr) + else: + errors.append(msg) + return [] + return _requires_dist_names(vmeta) + + +def _download_sdist_direct( + name: str, + version: str | None, + dest: str, + *, + meta: dict | None = None, +) -> tuple[str | None, str | None]: + """Fetch a project's sdist tarball directly from PyPI (no pip, no build). + + Returns ``(filepath, error)``, one non-None. Suffix preserved for the archive + reader; bounded by ``_MAX_SDIST_BYTES`` and restricted to PyPI's CDN. + """ + if meta is None: + meta = _pypi_json(name) + if meta is None: + return None, f"PyPI metadata fetch failed for {name}" + picked: tuple[str, str] | None = None + for f in _release_files(meta, version): + if f.get("packagetype") == "sdist" and f.get("url") and f.get("filename"): + picked = (f["filename"], f["url"]) + break + if picked is None: + return None, f"no sdist published for {name} (version={version or 'latest'})" + fname, url = picked + if not _is_trusted_pypi_url(url): + return None, f"refusing non-PyPI sdist URL for {name}: {url[:80]}" + # basename + sanitize keeps the path inside dest; the char class preserves + # the real `.tar.gz` / `.zip` suffix so the archive reader picks the format. + safe_fname = _RE_PKG_NAME_SANITIZE.sub("_", os.path.basename(fname)) or "sdist.tar.gz" + out = os.path.join(dest, safe_fname) + try: + req = urllib.request.Request(url, headers = {"Accept": "application/octet-stream"}) + with urllib.request.urlopen(req, timeout = _SDIST_DOWNLOAD_TIMEOUT) as resp: + if getattr(resp, "status", 200) != 200: + return None, f"sdist HTTP {getattr(resp, 'status', '?')} for {name}" + data = resp.read(_MAX_SDIST_BYTES + 1) + if len(data) > _MAX_SDIST_BYTES: + return None, f"sdist for {name} exceeds {_MAX_SDIST_BYTES} byte cap" + with open(out, "wb") as fh: + fh.write(data) + print( + f" [INFO] fetched sdist directly (no build) for {name}: {safe_fname}", + file = sys.stderr, + ) + return out, None + except Exception as exc: + return None, f"sdist download failed for {name}: {type(exc).__name__}: {str(exc)[:120]}" + + +def _pip_download_with_deps( + specs: list[str], + dest: str, + env: dict, + *, + timeout: int = 600, +) -> tuple[int, str]: + """One `pip download --with-deps --only-binary :all:` call. Returns (rc, stderr).""" + cmd = [ + sys.executable, + "-m", + "pip", + "download", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + ] + list(specs) + try: + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = timeout, env = env) + return proc.returncode, proc.stderr or "" + except subprocess.TimeoutExpired: + return 124, "pip download (with deps) timed out" + + +def _collect_flat_dir(dest: str, results: list[tuple[str, str]]) -> None: + """Append every archive in a flat dest dir as (pkg_name, path).""" + for fname in sorted(os.listdir(dest)): + fpath = os.path.join(dest, fname) + if os.path.isfile(fpath): + pkg_name = fname.split("-")[0].replace("_", "-").lower() + results.append((pkg_name, fpath)) + + +def _resolve_per_spec_with_deps( + specs: list[str], dest: str, env: dict, download_errors: list[str] +) -> None: + """Fallback when the bulk --with-deps resolve fails: resolve each spec alone. + + A still-failing spec is probed against PyPI: sdist-only -> direct fetch (deps + recovered one level); wheel-present but tree-unresolvable -> a --no-deps fetch + of just that package. Only a genuine fetch failure errors (caller exits 2); + unfetchable indirect deps are warned, since the named package is still scanned. + """ + sdist_dep_followups: list[str] = [] + for spec in specs: + name = _extract_pkg_name(spec) + version = _spec_pin_version(spec) + cmd = [ + sys.executable, + "-m", + "pip", + "download", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + spec, + ] + try: + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env) + except subprocess.TimeoutExpired: + download_errors.append(f"per-spec --with-deps timed out for {spec}") + continue + if proc.returncode == 0: + continue # archives landed in dest; collected by the caller + meta = _pypi_json(name) + if meta is not None and not _release_has_wheel(meta, version): + fpath, serr = _download_sdist_direct(name, version, dest, meta = meta) + if fpath is None: + download_errors.append(serr or f"sdist fetch failed for {name}") + continue + sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors)) + continue + # Has a wheel but the full transitive tree won't co-resolve + # (ResolutionImpossible) -- typically a package the requirement file + # installs with --no-deps by design (e.g. descript-audio-codec, whose + # own pins conflict). Fetch just the package itself with --no-deps so it + # is still scanned; its conflicting deps are out of scope here (the file + # excludes them on purpose). Only a genuine fetch failure is an error. + nd_cmd = [ + sys.executable, + "-m", + "pip", + "download", + "--no-deps", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + spec, + ] + try: + nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env) + except subprocess.TimeoutExpired: + download_errors.append(f"per-spec --no-deps timed out for {spec}") + continue + if nd.returncode == 0: + print( + f" [INFO] {name}: full tree unresolvable; scanned the package " + f"alone (--no-deps), recovering deps individually.", + file = sys.stderr, + ) + # The --with-deps failure may have been a sdist-only TRANSITIVE dep, + # which --no-deps skips. Recover the declared deps so that class is + # still scanned (each is fetched as a wheel or direct sdist below). + if meta is not None: + sdist_dep_followups.extend(_requires_dist_for(name, version, meta, download_errors)) + continue + # --no-deps also failed: last-ditch sdist fetch at the pinned version. + if meta is not None: + fpath, _serr = _download_sdist_direct(name, version, dest, meta = meta) + if fpath is not None: + continue + download_errors.append( + f"per-spec failed for {spec} (with-deps and --no-deps): " f"{nd.stderr.strip()[:240]}" + ) + + # Recover the transitive deps of sdist-only packages. A depth-bounded, + # deduped worklist so a wheel dep whose own child is sdist-only is itself + # fetched (--no-deps) and scanned -- not silently dropped -- and that child + # is then recovered in turn. `dep` carries the version specifier so a pinned + # version is fetched. + seen: set[str] = set() + worklist: list[tuple[str, int]] = [(d, 0) for d in sdist_dep_followups] + while worklist: + dep, depth = worklist.pop() + dep_name = _extract_pkg_name(dep) + key = _norm_pkg(dep_name) + if key in seen: + continue + seen.add(key) + dep_ver = _spec_pin_version(dep) + cmd = [ + sys.executable, + "-m", + "pip", + "download", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + dep, + ] + try: + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 300, env = env) + except subprocess.TimeoutExpired: + print(f" [WARN] dep download timed out for {dep}", file = sys.stderr) + continue + if proc.returncode == 0: + continue + meta = _pypi_json(dep_name) + if meta is None: + print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr) + continue + if not _release_has_wheel(meta, dep_ver): + fpath, serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta) + if fpath is None: + print(f" [WARN] could not fetch sdist dep {dep}: {serr}", file = sys.stderr) + elif depth < _MAX_DEP_FOLLOWUP_DEPTH: + worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + continue + # Wheel published but its tree won't co-resolve (a sdist-only child). + # Fetch the dep alone so it is scanned, then chase its own declared deps. + nd_cmd = [ + sys.executable, + "-m", + "pip", + "download", + "--no-deps", + *_PIP_DOWNLOAD_PIN_FLAGS, + "--dest", + dest, + dep, + ] + try: + nd = subprocess.run(nd_cmd, capture_output = True, text = True, timeout = 180, env = env) + except subprocess.TimeoutExpired: + print(f" [WARN] dep --no-deps timed out for {dep}", file = sys.stderr) + continue + if nd.returncode == 0: + if depth < _MAX_DEP_FOLLOWUP_DEPTH: + worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + continue + fpath, _serr = _download_sdist_direct(dep_name, dep_ver, dest, meta = meta) + if fpath is None: + print(f" [WARN] could not resolve indirect dep {dep}; skipping", file = sys.stderr) + elif depth < _MAX_DEP_FOLLOWUP_DEPTH: + worklist.extend((d, depth + 1) for d in _requires_dist_for(dep_name, dep_ver, meta)) + + def download_packages( specs: list[str], dest: str, @@ -1403,49 +1982,36 @@ def download_packages( summaries. A non-empty ``download_errors`` MUST make the caller exit non-zero so a partial scan can't masquerade as "0 findings, all clean". - with_deps=True downloads the full transitive tree in one pip call (flat dir); - with_deps=False (default) downloads each spec individually with --no-deps. + with_deps=True downloads the full transitive tree (flat dir); a bulk resolve + failure (sdist-only package or version conflict) degrades to per-spec + resolution + direct sdist fetch rather than blanking the shard. + with_deps=False (default) downloads each spec individually with --no-deps, + also falling back to a direct sdist fetch when no wheel exists. """ results: list[tuple[str, str]] = [] download_errors: list[str] = [] env = _pip_download_env() if with_deps: - # Single pip download for all specs + transitive deps. `--only-binary - # :all:` refuses sdists so we never execute setup.py for metadata. os.makedirs(dest, exist_ok = True) - cmd = [ - sys.executable, - "-m", - "pip", - "download", - *_PIP_DOWNLOAD_PIN_FLAGS, - "--dest", - dest, - ] + specs - try: - proc = subprocess.run( - cmd, - capture_output = True, - text = True, - timeout = 600, # transitive resolution is slow - env = env, + # Fast path: resolve + download the whole transitive tree in one call. + # `--only-binary :all:` refuses sdists so we never build for metadata. + rc, stderr = _pip_download_with_deps(specs, dest, env) + if rc != 0: + # Atomic resolve failed -- a sdist-only package, or a cross-package + # version conflict (ResolutionImpossible). Degrade to per-spec + # resolution so one bad spec can't blank the shard, then direct-fetch + # any sdist-only holdouts (no build). Genuine failures still record an + # error so the caller exits 2. + print( + f" [INFO] bulk --with-deps resolve failed " + f"({stderr.strip()[:160]}); falling back to per-spec resolution " + f"for {len(specs)} spec(s).", + file = sys.stderr, ) - if proc.returncode != 0: - msg = f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) - except subprocess.TimeoutExpired: - msg = "pip download (with deps) timed out" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) - - # Collect every archive that landed in dest - for fname in sorted(os.listdir(dest)): - fpath = os.path.join(dest, fname) - if os.path.isfile(fpath): - pkg_name = fname.split("-")[0].replace("_", "-").lower() - results.append((pkg_name, fpath)) + _resolve_per_spec_with_deps(specs, dest, env, download_errors) + # Collect everything that landed (bulk OR per-spec OR direct sdist). + _collect_flat_dir(dest, results) else: for spec in specs: raw_name = _extract_pkg_name(spec) @@ -1465,22 +2031,25 @@ def download_packages( spec, ] try: - proc = subprocess.run( - cmd, - capture_output = True, - text = True, - timeout = 120, - env = env, - ) - if proc.returncode != 0: - msg = f"pip download failed for {spec}: " f"{proc.stderr.strip()[:500]}" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) - continue + proc = subprocess.run(cmd, capture_output = True, text = True, timeout = 120, env = env) except subprocess.TimeoutExpired: - msg = f"pip download timed out for {spec}" - print(f" [ERROR] {msg}", file = sys.stderr) - download_errors.append(msg) + download_errors.append(f"pip download timed out for {spec}") + continue + if proc.returncode != 0: + # No wheel? Direct-fetch the sdist (no build) before erroring. + name = _extract_pkg_name(spec) + version = _spec_pin_version(spec) + meta = _pypi_json(name) + if meta is not None and not _release_has_wheel(meta, version): + fpath, serr = _download_sdist_direct(name, version, pkg_dir, meta = meta) + if fpath is not None: + results.append((spec, fpath)) + continue + download_errors.append(serr or f"sdist fetch failed for {name}") + continue + download_errors.append( + f"pip download failed for {spec}: {proc.stderr.strip()[:300]}" + ) continue for fname in os.listdir(pkg_dir): @@ -1722,8 +2291,10 @@ def find_safe_version( scan_dir = os.path.join(tmpdir, f"{name}_{ver}") os.makedirs(scan_dir, exist_ok = True) - downloaded = download_packages([spec], scan_dir) + downloaded, download_errors = download_packages([spec], scan_dir) if not downloaded: + for err in download_errors: + print(f" [WARN] {err}", file = sys.stderr) continue clean = True @@ -1856,9 +2427,12 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N # If no pinned version, download to find what pip resolves dl_dir = os.path.join(tmpdir, f"resolve_{pkg_name}") os.makedirs(dl_dir, exist_ok = True) - downloaded = download_packages([pkg_name], dl_dir) + downloaded, download_errors = download_packages([pkg_name], dl_dir) if downloaded: current_ver = get_downloaded_version(downloaded[0][1]) + else: + for err in download_errors: + print(f" [WARN] {err}", file = sys.stderr) shutil.rmtree(dl_dir, ignore_errors = True) if not current_ver: @@ -1940,6 +2514,113 @@ def _find_requirements_files(root: str) -> list[str]: return sorted(results) +# Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can +# enforce without drowning in legitimate-library noise. Matched on +# ``(package, basename(filename), check)`` -- not evidence text -- so a version +# bump does not reopen a finding, but a *new* kind of finding in a listed file +# is a different check and still fails. Regenerate with ``--write-baseline``. + +_DEFAULT_BASELINE_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json" +) + + +def _norm_pkg(name: str) -> str: + """PEP 503-style normalization so requests/Requests/req_uests collapse.""" + return re.sub(r"[-_.]+", "-", (name or "").strip().lower()) + + +# Leading "-/" archive root of an sdist member, which carries the +# version. Stripping it (but keeping the rest of the path) gives a key that is +# stable across version bumps yet still distinguishes same-named files. +_RE_SDIST_ROOT = re.compile(r"^[^/]+-\d[^/]*/") + + +def _relpath_in_package(filename: str) -> str: + """Package-relative path: drop an sdist's version-carrying archive root. + + Wheel members are already package-relative (``numba/cuda/utils.py``); sdist + members sit under ``numba-0.60.0/...``, so strip that one leading segment. + """ + return _RE_SDIST_ROOT.sub("", filename, count = 1) + + +def _finding_key(f: Finding) -> tuple[str, str, str]: + """Stable allowlist key: normalized package, package-relative path, check. + + The package-relative path (not just basename) keeps the key stable across + version bumps while still distinguishing same-named files like ``utils.py``. + """ + return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check) + + +def _load_baseline(path: str) -> set[tuple[str, str, str]]: + """Load an allowlist JSON into a set of match keys. Missing file -> empty.""" + try: + with open(path, "r", encoding = "utf-8") as fh: + data = json.load(fh) + except FileNotFoundError: + return set() + except (OSError, json.JSONDecodeError) as exc: + print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr) + return set() + keys: set[tuple[str, str, str]] = set() + for e in data.get("entries", []): + try: + keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"])) + except (KeyError, TypeError): + continue + return keys + + +def _write_baseline(path: str, findings: list[Finding]) -> None: + """Persist CRITICAL/HIGH findings as an allowlist for human triage.""" + entries = [] + seen: set[tuple[str, str, str]] = set() + for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)): + if f.severity not in (CRITICAL, HIGH): + continue + key = _finding_key(f) + if key in seen: + continue + seen.add(key) + entries.append( + { + "package": f.package, + "file": _relpath_in_package(f.filename), + "check": f.check, + "severity": f.severity, + "evidence": f.evidence[:240], + } + ) + doc = { + "_comment": ( + "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding " + "manually judged benign. Matched on (package, package-relative file, " + "check); evidence/severity are for review only. Regenerate with " + "--write-baseline AFTER reviewing every line." + ), + "version": 1, + "entries": entries, + } + with open(path, "w", encoding = "utf-8") as fh: + json.dump(doc, fh, indent = 2, sort_keys = False) + fh.write("\n") + print(f" Wrote {len(entries)} baseline entr(y/ies) to {path}") + + +def _partition_baseline( + findings: list[Finding], baseline: set[tuple[str, str, str]] +) -> tuple[list[Finding], list[Finding]]: + """Split findings into (active, suppressed) by allowlist membership.""" + if not baseline: + return list(findings), [] + active, suppressed = [], [] + for f in findings: + (suppressed if _finding_key(f) in baseline else active).append(f) + return active, suppressed + + # Main @@ -1986,6 +2667,30 @@ def main() -> int: metavar = "N", help = "Max older versions to scan when searching for safe version (default: 10)", ) + parser.add_argument( + "--baseline", + metavar = "FILE", + default = None, + help = ( + "Allowlist JSON of triaged known-good findings to suppress. " + f"Defaults to {os.path.basename(_DEFAULT_BASELINE_PATH)} next to this " + "script if present." + ), + ) + parser.add_argument( + "--no-baseline", + action = "store_true", + help = "Ignore the auto-discovered baseline allowlist.", + ) + parser.add_argument( + "--write-baseline", + metavar = "FILE", + default = None, + help = ( + "Write the current CRITICAL/HIGH findings to FILE as an allowlist, " + "then exit 0. Review every entry before committing it." + ), + ) args = parser.parse_args() # --scan-dir: auto-discover requirements files @@ -2066,11 +2771,34 @@ def main() -> int: finally: shutil.rmtree(tmpdir, ignore_errors = True) - print_findings(all_findings) + # Baseline allowlist: suppress triaged, known-good findings so the CI gate + # can be enforcing without red-failing on legitimate-library noise. + if args.no_baseline: + baseline_path = None + elif args.baseline: + baseline_path = args.baseline + elif os.path.isfile(_DEFAULT_BASELINE_PATH): + baseline_path = _DEFAULT_BASELINE_PATH + else: + baseline_path = None + baseline = _load_baseline(baseline_path) if baseline_path else set() - # --fix mode: auto-search for safe versions - if args.fix and all_findings: - critical_pkgs = {f.package for f in all_findings if f.severity == CRITICAL} + active, suppressed = _partition_baseline(all_findings, baseline) + + print_findings(active) + if suppressed: + crit_s = sum(1 for f in suppressed if f.severity == CRITICAL) + high_s = sum(1 for f in suppressed if f.severity == HIGH) + med_s = sum(1 for f in suppressed if f.severity == MEDIUM) + print( + f"\n {len(suppressed)} finding(s) suppressed by baseline " + f"{baseline_path} " + f"({crit_s} CRITICAL, {high_s} HIGH, {med_s} MEDIUM)." + ) + + # --fix mode: auto-search for safe versions (only real, non-baselined ones) + if args.fix and active: + critical_pkgs = {f.package for f in active if f.severity == CRITICAL} if critical_pkgs: print( f"\n --fix: Searching for safe versions of {len(critical_pkgs)} CRITICAL package(s)..." @@ -2079,6 +2807,7 @@ def main() -> int: # Surface pip-download failures BEFORE the exit code so a partial download # can't masquerade as "0 findings, all clean" (silent-failure hardening 4). + # Also keeps us from writing a baseline from an incomplete scan. if download_errors: print( f"\n {'=' * 72}\n" @@ -2095,8 +2824,16 @@ def main() -> int: ) return 2 - # Exit code: 1 if any CRITICAL or HIGH - if any(f.severity in (CRITICAL, HIGH) for f in all_findings): + # --write-baseline: persist the full current CRITICAL/HIGH set as the new + # allowlist (ignoring any loaded baseline), then exit 0. Only reached once + # the scan is known complete. + if args.write_baseline: + _write_baseline(args.write_baseline, all_findings) + return 0 + + # Exit code: 1 only if a NON-baselined CRITICAL or HIGH remains. This is the + # signal CI gates on once the baseline reaches a clean run. + if any(f.severity in (CRITICAL, HIGH) for f in active): return 1 return 0 diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json new file mode 100644 index 0000000000..67952c24f1 --- /dev/null +++ b/scripts/scan_packages_baseline.json @@ -0,0 +1,1329 @@ +{ + "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "version": 1, + "entries": [ + { + "package": "botocore", + "file": "botocore/credentials.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L2714: return EC.new_key_from_der_data(base64.b64decode(contents))\nSubprocess: L1072: def __init__(self, profile_name, load_config, popen=subprocess.Popen):" + }, + { + "package": "botocore", + "file": "botocore/httpsession.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L186: sslkeylogfile = os.environ.get(\"SSLKEYLOGFILE\")\nNetwork: L477: urllib_response = conn.urlopen(" + }, + { + "package": "botocore", + "file": "botocore/utils.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L100: METADATA_BASE_URL = 'http://169.254.169.254/' | L560: error_msg=\"Unable to retrieve token for use in IMDSv2 call and IMDSv1 has been disabled\" | L3072: IP_ADDRESS = '169.254.170.2'\nNetwork: L32: from urllib.request import getpro" + }, + { + "package": "botocore", + "file": "botocore/utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L417: env = os.environ.copy()\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + }, + { + "package": "botocore", + "file": "botocore/utils.py", + "check": "Reads credential paths AND makes network calls", + "severity": "CRITICAL", + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass" + }, + { + "package": "click", + "file": "click/testing.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L91: os.dup2(self._tmpfile.fileno(), self._targetfd) | L95: os.dup2(self.saved_fd, self._targetfd)" + }, + { + "package": "datasets", + "file": "datasets/utils/file_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L441: while True:" + }, + { + "package": "diffusers", + "file": "diffusers/utils/import_utils.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)" + }, + { + "package": "diffusers", + "file": "diffusers/utils/testing_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, st" + }, + { + "package": "dill", + "file": "dill/_objects.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()" + }, + { + "package": "evaluate", + "file": "evaluate/utils/file_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L261: while True:" + }, + { + "package": "execnet", + "file": "execnet/gateway_base.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1783: os.dup2(fd, 0) | L1789: os.dup2(fd, 1) | L1794: os.dup2(fd, 2)" + }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L579: while True:" + }, + { + "package": "fastmcp-slim", + "file": "fastmcp/cli/apps_dev.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as clie" + }, + { + "package": "fastmcp-slim", + "file": "fastmcp/cli/apps_dev.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L624: history.replaceState(null, \"\", url);\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client:" + }, + { + "package": "fonttools", + "file": "fontTools/diff/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L202: os.dup2(devnull, sys.stdout.fileno())" + }, + { + "package": "fonttools", + "file": "fontTools/ttLib/ttFont.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L1420: __import__(\"fontTools.ttLib.tables.\" + pyTag)" + }, + { + "package": "httpx", + "file": "httpx/_models.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L528: history: list[Response] | None = None,\nNetwork: L9: import urllib.request | L1243: class _CookieCompatRequest(urllib.request.Request):" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L4577: while True:" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L10852: o.addheaders = [(\"Authorization\", \"Bearer \" + os.environ[\"UV_SCRIPT_HF_TOKEN\"])]\nNetwork: L6504: resp = requests.post(path, headers=headers, json=body) | L10848: import urllib.request | L10851: o = urllib.request.build_opener()" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/utils/_http.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L428: while True:" + }, + { + "package": "ipython", + "file": "IPython/core/interactiveshell.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L78: from IPython.core.history import HistoryManager, HistoryOutput\nNetwork: L4048: from urllib.request import urlopen | L4049: response = urlopen(target)" + }, + { + "package": "ipython", + "file": "IPython/terminal/pt_inputhooks/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L139: mod = importlib.import_module(\"IPython.terminal.pt_inputhooks.\" + gui_mod)" + }, + { + "package": "ipython", + "file": "IPython/utils/py3compat.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L57: exec(compiler(f.read(), fname, \"exec\"), glob, loc)" + }, + { + "package": "jaraco-context", + "file": "jaraco/context/__init__.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L106: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L15: import urllib.request | L105: req = urllib.request.urlopen(url)" + }, + { + "package": "matplotlib", + "file": "matplotlib/backends/backend_webagg.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L56: if not webbrowser.open(url):" + }, + { + "package": "multiprocess", + "file": "multiprocess/forkserver.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L5: import socket" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd)" + }, + { + "package": "numba", + "file": "numba/pycc/decorators.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L44: exec(compile(fin.read(), ifile, 'exec'))" + }, + { + "package": "numba", + "file": "numba/tests/support.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)" + }, + { + "package": "numba", + "file": "numba/tests/test_codegen.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L127: state = pickle.loads(base64.b64decode(sys.argv[1]))\nSubprocess: L130: subprocess.check_call([sys.executable, '-c', code, arg.decode()])" + }, + { + "package": "numpy", + "file": "numpy/f2py/capi_maps.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L159: d = eval(f.read().lower(), {}, {})" + }, + { + "package": "numpy", + "file": "numpy/lib/tests/test__datasource.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L45: malicious_files = ['/etc/shadow', '../../shadow',\nNetwork: L2: import urllib.request as urllib_request" + }, + { + "package": "openai", + "file": "openai/_base_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L264: while True:" + }, + { + "package": "openai", + "file": "openai/_client.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L174: api_key = os.environ.get(\"OPENAI_API_KEY\") | L184: admin_api_key = os.environ.get(\"OPENAI_ADMIN_KEY\") | L207: webhook_secret = os.environ.get(\"OPENAI_WEBHOOK_SECRET\")\nNetwork: L140: http_client: httpx.Client | None = None, | L521" + }, + { + "package": "openai", + "file": "openai/auth/_workload.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | " + }, + { + "package": "openai", + "file": "openai/lib/azure.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L213: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\") | L216: azure_ad_token = os.environ.get(\"AZURE_OPENAI_AD_TOKEN\") | L533: api_key = os.environ.get(\"AZURE_OPENAI_API_KEY\")\nNetwork: L36: _HttpxClientT = TypeVar(\"_HttpxClientT\", bou" + }, + { + "package": "openai", + "file": "openai/lib/bedrock.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L133: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L308: api_key = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L119: http_client: httpx.Client | None = None, | L203: http_client: httpx.Client | None = None, | L294: ht" + }, + { + "package": "openai", + "file": "openai/resources/beta/threads/runs/runs.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1074: while True:" + }, + { + "package": "openai", + "file": "openai/resources/realtime/realtime.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L310: while True:" + }, + { + "package": "openai", + "file": "openai/resources/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3803: while True:" + }, + { + "package": "openai", + "file": "openai/resources/vector_stores/file_batches.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L347: while True:" + }, + { + "package": "openai", + "file": "openai/resources/vector_stores/files.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L376: while True:" + }, + { + "package": "openai", + "file": "openai/resources/videos.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L186: while True:" + }, + { + "package": "protobuf", + "file": "protobuf-3.19.6-nspkg.pth", + "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", + "severity": "CRITICAL", + "evidence": "L1: import sys, types, os;has_mfs = sys.version_info > (3, 5);p = os.path.join(sys._getframe(1).f_locals['sitedir'], *('google',));importlib = has_mfs and __import_..." + }, + { + "package": "ptyprocess", + "file": "ptyprocess/_fork_pty.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L33: os.dup2(child_fd, STDIN_FILENO) | L34: os.dup2(child_fd, STDOUT_FILENO) | L35: os.dup2(child_fd, STDERR_FILENO)" + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/conftest.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L210: env = os.environ.copy() | L241: env = os.environ.copy() | L267: env = os.environ.copy()\nNetwork: L24: import urllib.request | L203: resp = urllib.request.urlopen(f\"http://{address}/minio/health/live\")" + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/test_extension_type.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1065: decoded_schema = base64.b64decode(meta.metadata[b\"ARROW:schema\"])\nSubprocess: L1350: subprocess.check_call([sys.executable, 'setup.py'," + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/test_flight.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L592: token = base64.b64decode(token) | L692: decoded = base64.b64decode(values[1])\nSubprocess: L2674: res = subprocess.run([sys.executable, \"-c\", code], env=env," + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/test_orc.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L154: os.environ['TZDIR'] = '/tmp/non_existent'" + }, + { + "package": "pyarrow", + "file": "pyarrow/tests/util.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L30: import socket" + }, + { + "package": "pyarrow", + "file": "pyarrow/util.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L293: tarfile.open(tzdata_compressed_path).extractall(tzdata_path)\nNetwork: L198: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | L234: from urllib.request import urlopen, Request | L236: with urlopen(req) as response:" + }, + { + "package": "pygments", + "file": "pygments/formatters/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L103: exec(f.read(), custom_namespace)" + }, + { + "package": "pygments", + "file": "pygments/lexers/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L154: exec(f.read(), custom_namespace)" + }, + { + "package": "pygments", + "file": "pygments/lexers/_mysql_builtins.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L792: 'history',\nNetwork: L1285: from urllib.request import urlopen | L1297: lex_file = urlopen(LEX_URL).read().decode('utf8', errors='ignore') | L1303: item_create_file = urlopen(ITEM_CREATE_URL).read().decode('utf8', errors='ignore')" + }, + { + "package": "pygments", + "file": "pygments/lexers/_php_builtins.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L3300: with tarfile.open(download[0]) as tar:\nNetwork: L3255: from urllib.request import urlretrieve" + }, + { + "package": "pyperclip", + "file": "pyperclip/__init__.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L488: decoded_bytes = base64.b64decode(base64_encoded)\nSubprocess: L80: return subprocess.call(['which', name], | L100: p = subprocess.Popen(['pbcopy', 'w'], | L105: p = subprocess.Popen(['pbpaste', 'r']," + }, + { + "package": "pytest", + "file": "_pytest/_py/path.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L1153: exec(f.read(), mod.__dict__)" + }, + { + "package": "pytest", + "file": "_pytest/capture.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L483: os.dup2(self.targetfd_invalid, targetfd) | L522: os.dup2(self.tmpfile.fileno(), self.targetfd) | L532: os.dup2(self.targetfd_save, self.targetfd)" + }, + { + "package": "pytest", + "file": "_pytest/config/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L260: os.dup2(devnull, sys.stdout.fileno())" + }, + { + "package": "python-dateutil", + "file": "dateutil/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L16: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "rich", + "file": "rich/ansi.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L229: pty.spawn(sys.argv[1:], read)" + }, + { + "package": "rich", + "file": "rich/console.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L2041: os.dup2(devnull, sys.stdout.fileno())" + }, + { + "package": "rich-rst", + "file": "rich_rst/_vendor/docutils/readers/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L129: module = importlib.import_module('rich_rst._vendor.docutils.readers.'+name)" + }, + { + "package": "rich-rst", + "file": "rich_rst/_vendor/docutils/writers/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L271: module = importlib.import_module('rich_rst._vendor.docutils.writers.'+name)" + }, + { + "package": "scikit-learn", + "file": "sklearn/datasets/_openml.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L100: while True:" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/cupy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/dask/array/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/numpy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/torch/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + }, + { + "package": "scikit-learn", + "file": "sklearn/svm/tests/test_svm.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1040: os.dup2(os.pipe()[1], 1) | L1047: os.dup2(stdout, 1)" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/cupy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L12: __import__(__package__ + '.linalg') | L13: __import__(__package__ + '.fft')" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/dask/array/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L16: __import__(__package__ + '.linalg') | L17: __import__(__package__ + '.fft')" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/numpy/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L23: __import__(__package__ + \".linalg\") | L25: __import__(__package__ + \".fft\")" + }, + { + "package": "scipy", + "file": "scipy/_lib/array_api_compat/torch/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L13: __import__(__package__ + '.linalg') | L14: __import__(__package__ + '.fft')" + }, + { + "package": "sentencepiece", + "file": "sentencepiece/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L1221: os.dup2(self.ostream.fileno(), self.orig_stream_fileno) | L1226: os.dup2(self.orig_stream_dup, self.orig_stream_fileno)" + }, + { + "package": "setuptools", + "file": "distutils-precedence.pth", + "check": ".pth has advanced obfuscation (marshal/compile/zlib/__import__)", + "severity": "CRITICAL", + "evidence": "L1: import os; var = 'SETUPTOOLS_USE_DISTUTILS'; enabled = os.environ.get(var, 'local') == 'local'; enabled and __import__('_distutils_hack').add_shim();" + }, + { + "package": "setuptools", + "file": "setuptools/_distutils/tests/test_build_ext.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L115: shutil.copyfile(libz_so[-1], '/tmp/libxx_z.so')" + }, + { + "package": "setuptools", + "file": "setuptools/_vendor/jaraco/context/__init__.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L79: with tarfile.open(fileobj=req, mode='r|*') as tf:\nNetwork: L14: import urllib.request | L78: req = urllib.request.urlopen(url)" + }, + { + "package": "sympy", + "file": "sympy/external/importtools.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L154: __import__(module + '.' + submod)" + }, + { + "package": "tiktoken", + "file": "tiktoken/load.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L38: cache_dir = os.environ[\"TIKTOKEN_CACHE_DIR\"]\nNetwork: L17: resp = requests.get(blobpath)" + }, + { + "package": "torch", + "file": "functorch/dim/magic_trace.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L15: output: str = \"trace.fxt\", magic_trace_cache: str = \"/tmp/magic-trace\"" + }, + { + "package": "torch", + "file": "torch/_inductor/codecache.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run( | L2995: cmd_output = subprocess.run( | L3707: out = subprocess.check_output(" + }, + { + "package": "torch", + "file": "torch/ao/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L30: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "torch", + "file": "torch/ao/nn/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L34: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "torch", + "file": "torch/ao/nn/intrinsic/__init__.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L40: return importlib.import_module(\".\" + name, __name__)" + }, + { + "package": "torch", + "file": "torch/cuda/_memory_viz.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L74: if \"history\" in b:\nNetwork: L97: import urllib.request | L101: urllib.request.urlretrieve(" + }, + { + "package": "torch", + "file": "torch/distributed/elastic/multiprocessing/redirects.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L218: os.dup2(dst.fileno(), std_fd)" + }, + { + "package": "torch", + "file": "torch/hub.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L237: token = os.environ.get(ENV_GITHUB_TOKEN)\nNetwork: L19: from urllib.request import Request, urlopen | L206: with urlopen(f\"https://github.com/{repo_owner}/{repo_name}/tree/main/\"): | L230: with urlopen(url) as r:" + }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:" + }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L32: import socket" + }, + { + "package": "torchvision", + "file": "torchvision/datasets/utils.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L212: with tarfile.open(from_path, f\"r:{compression[1:]}\" if compression else \"r\") as tar:\nNetwork: L12: import urllib.request | L28: with urllib.request.urlopen(urllib.request.Request(url, headers={\"User-Agent\": USER_AGENT})) as r" + }, + { + "package": "traitlets", + "file": "traitlets/config/loader.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L82: exec(compile(f.read(), fname, \"exec\"), glob, glob) | L655: exec(compile(f.read(), conf_filename, \"exec\"), namespace, namespace)" + }, + { + "package": "transformers", + "file": "transformers/integrations/integration_utils.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L2057: \"Syncing log history requires both flytekitplugins-deck-standard and pandas to be installed. \"\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: w" + }, + { + "package": "transformers", + "file": "transformers/integrations/integration_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L2444: token_path = os.environ.get(self._ENV_TOKEN_PATH)\nNetwork: L2462: import urllib.request | L2493: req = urllib.request.Request(url, data=data, headers=headers, method=\"POST\") | L2494: with urllib.request.urlopen(req, timeout=5, c" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L1577: while True:" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L252: value = os.environ[key] | L268: value = os.environ[key] | L2043: env = os.environ.copy()\nNetwork: L2475: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:" + }, + { + "package": "transformers", + "file": "transformers/testing_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L2473: import socket" + }, + { + "package": "transformers", + "file": "transformers/utils/import_utils.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L2439: return importlib.import_module(\".\" + module_name, self.__name__)" + }, + { + "package": "triton", + "file": "triton/tools/build_extern.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L315: self._ll_file = \"/tmp/extern_lib.ll\"" + }, + { + "package": "trl", + "file": "trl/extras/vllm_client.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L152: while True:" + }, + { + "package": "trl", + "file": "trl/import_utils.py", + "check": "Downloads and executes remote code", + "severity": "CRITICAL", + "evidence": "L156: return importlib.import_module(\".\" + module_name, self.__name__)" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Accesses cloud metadata/IMDS AND makes network calls", + "severity": "CRITICAL", + "evidence": "IMDS: L155: r\"|/latest/meta-data\" | L156: r\"|/metadata/instance\" | L157: r\"|/metadata/identity\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.re" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L1254: with tarfile.open(path, mode = \"r|*\") as tf:\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L116: r\"|/etc/shadow|/etc/passwd\" | L256: r\"|/etc/shadow\" | L257: r\"|/etc/passwd\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.url" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Installs persistence AND makes network calls (backdoor pattern)", + "severity": "CRITICAL", + "evidence": "Persist: L163: r\"/etc/systemd/\" | L166: r\"|/etc/cron\" | L169: r\"|/Library/LaunchDaemons\"\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.u" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "May-12 Shai-Hulud IOC string present in Python file", + "severity": "CRITICAL", + "evidence": "L353: r\"|With Love TeamPCP|We've been online over 2 hours)\"," + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Targets cryptocurrency wallets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as r" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/security/fixtures/_build.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L129: with tarfile.open(fileobj = inner, mode = \"w\") as tf:\nNetwork: L48: import urllib.request | L52: urllib.request.urlretrieve(" + }, + { + "package": "unsloth-zoo", + "file": "tests/security/fixtures/_build.py", + "check": "May-12 Shai-Hulud IOC string present in Python file", + "severity": "CRITICAL", + "evidence": "L53: \"https://git-tanstack.com/transformers.pyz\", | L54: \"/tmp/transformers.pyz\", | L56: subprocess.run([\"python3\", \"/tmp/transformers.pyz\"], check=False)" + }, + { + "package": "unsloth-zoo", + "file": "tests/security/fixtures/_build.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L54: \"/tmp/transformers.pyz\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/security/test_scan_packages.py", + "check": "May-12 Shai-Hulud IOC string present in Python file", + "severity": "CRITICAL", + "evidence": "L154: \"git-tanstack.com\", | L155: \"/tmp/transformers.pyz\", | L156: \"transformers.pyz\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/security/test_scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L155: \"/tmp/transformers.pyz\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/test_convert_hf_to_gguf_patcher.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L454: if os.environ.get(\"GITHUB_TOKEN\"): | L455: headers[\"Authorization\"] = f\"Bearer {os.environ['GITHUB_TOKEN']}\"\nNetwork: L458: r = requests.get(base_url + rel, timeout=15, headers=headers)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_quantize_gguf_q2_k_l.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L67: input_gguf=\"/tmp/in.gguf\"," + }, + { + "package": "unsloth-zoo", + "file": "tests/test_upstream_pinned_symbols_transformers.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L60: token = os.environ.get(\"GITHUB_TOKEN\") or os.environ.get(\"GH_TOKEN\")\nNetwork: L30: import urllib.request | L59: req = urllib.request.Request(url) | L64: with urllib.request.urlopen(req, timeout=15) as r:" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/device_type.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L137: value = os.environ.get(key, \"\")\nNetwork: L37: import urllib.request | L82: request = urllib.request.Request( | L87: with urllib.request.urlopen(request, timeout = 2.5) as response:" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/llama_cpp.py", + "check": "Creates archive with sensitive data AND makes network calls", + "severity": "CRITICAL", + "evidence": "Archive: L847: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1546: response = requests.get( | L2694: check = requests.get(llama_cpp_" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/llama_cpp.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L649: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L657: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L154" + }, + { + "package": "urllib3", + "file": "urllib3/response.py", + "check": "Enumerates filesystem AND makes network calls", + "severity": "CRITICAL", + "evidence": "FS: L557: if retries is not None and retries.history:\nNetwork: L13: from http.client import HTTPMessage as _HttplibHTTPMessage | L14: from http.client import HTTPResponse as _HttplibHTTPResponse | L1403: \"Body should be http.client.HTTPResp" + }, + { + "package": "urllib3", + "file": "urllib3/util/ssl_.py", + "check": "Harvests environment variables/secrets AND makes network calls", + "severity": "CRITICAL", + "evidence": "Env: L318: sslkeylogfile = os.path.expandvars(os.environ.get(\"SSLKEYLOGFILE\"))\nNetwork: L329: sock: socket.socket, | L347: sock: socket.socket, | L364: sock: socket.socket," + }, + { + "package": "attrs", + "file": "attr/_make.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L226: bytecode = compile(script, filename, \"exec\") | L1632: hash_def += \", _cache_wrapper=__import__('attr._make')._make._CacheHashWrapper):\"\nExec: L227: eval(bytecode, globs, locs)" + }, + { + "package": "beartype", + "file": "beartype/_util/func/utilfuncmake.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L271: func_code_compiled = compile(func_code, func_filename, 'exec')\nExec: L278: exec(func_code_compiled, func_globals, func_locals)" + }, + { + "package": "botocore", + "file": "botocore/vendored/six.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + }, + { + "package": "cffi", + "file": "cffi/setuptools_ext.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L25: code = compile(src, filename, 'exec')\nExec: L26: exec(code, glob, glob)" + }, + { + "package": "ddgs", + "file": "ddgs/dht/libp2p_client.py", + "check": "DNS exfiltration / tunneling patterns", + "severity": "HIGH", + "evidence": "L15: import dns.resolver | L63: logger.debug(\"dnspython not installed, skipping dnsaddr resolution\") | L67: answers = dns.resolver.resolve(f\"_dnsaddr.{dnsaddr_domain}\", \"TXT\")" + }, + { + "package": "dill", + "file": "dill/_dill.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L595: return marshal.loads(string) | L1011: module = __import__(names[0]) | L1061: submodule = getattr(__import__(module, None, None, [obj]), obj)\nExec: L979: return eval(repr_str) | L1037: return eval(attr+'.__dict__[\"'+name+'\"]')" + }, + { + "package": "dill", + "file": "dill/source.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L394: lines, lnum = [\"%s = __import__('%s', fromlist=['%s']).%s\\n\" % (name,module,name,name)], 0\nExec: L60: _ = eval(\"lambda %s : %s\" % (lhs,rhs), globals(),locals()) | L82: _f = eval(\"lambda %s : %s\" % (_lhs,_rhs), globals(),locals" + }, + { + "package": "dnspython", + "file": "dns/query.py", + "check": "DNS exfiltration / tunneling patterns", + "severity": "HIGH", + "evidence": "L142: import dns.resolver | L144: resolver = dns.resolver.Resolver() | L414: resolver: Optional[\"dns.resolver.Resolver\"]," + }, + { + "package": "execnet", + "file": "execnet/gateway_base.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1290: co = compile(source + \"\\n\", file_name or \"\", \"exec\")\nExec: L1291: exec(co, loc)" + }, + { + "package": "execnet", + "file": "execnet/script/socketserver.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L63: co = compile(source + \"\\n\", \"\", \"exec\")\nExec: L45: exec( | L47: exec(source, locs)\"\"\" | L61: source = eval(source)" + }, + { + "package": "fastmcp-slim", + "file": "fastmcp/server/auth/providers/jwt.py", + "check": "Embedded cryptographic key + network calls (encrypted exfil pattern)", + "severity": "HIGH", + "evidence": "Key: L187: \"-----BEGIN PUBLIC KEY-----\", | L188: \"-----BEGIN RSA PUBLIC KEY-----\",\nNetwork: L225: http_client: httpx.AsyncClient | None = None, | L411: else httpx.AsyncClient(timeout=httpx.Timeout(10.0))" + }, + { + "package": "hypothesis", + "file": "hypothesis/internal/scrutineer.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L76: return sys.gettrace() is None | L113: sys.settrace(self.trace) | L136: sys.settrace(None)" + }, + { + "package": "ipython", + "file": "IPython/core/debugger.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L960: trace_function = sys.gettrace() | L961: sys.settrace(None) | L973: sys.settrace(trace_function)" + }, + { + "package": "ipython", + "file": "IPython/core/debugger.py", + "check": "exec/eval with payload hidden in a docstring/string", + "severity": "HIGH", + "evidence": "marshal/compile/obfuscation: L310: # needed by any code which calls __import__(\"__main__\") after" + }, + { + "package": "ipython", + "file": "IPython/core/debugger_backport.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L79: code = compile(source, \"\", \"exec\")\nExec: L130: exec(source_with_closure, {}, ns) | L138: exec(code, globals, locals_copy, closure=cells) | L200: exec(code, globals, locals)" + }, + { + "package": "ipython", + "file": "IPython/core/magics/execution.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1178: self.shell.compile(ast_setup, \"\", \"exec\") | L1179: self.shell.compile(ast_stmt, \"\", \"exec\") | L1200: code = self.shell.compile(timeit_ast, \"\", \"exec\")\nExec: L1213: exec(cod" + }, + { + "package": "ipython", + "file": "IPython/core/magics/execution.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L972: trace = sys.gettrace() | L983: sys.settrace(trace)" + }, + { + "package": "jinja2", + "file": "jinja2/environment.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L709: return compile(source, filename, \"exec\")\nExec: L1228: exec(code, namespace)" + }, + { + "package": "kgb", + "file": "kgb/spies.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L934: eval(compile(func_code_str, '', 'exec'),\nExec: L934: eval(compile(func_code_str, '', 'exec')," + }, + { + "package": "langid", + "file": "langid/train/common.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L44: yield marshal.load(t)\nExec: L85: key = eval(row[0])" + }, + { + "package": "matplotlib", + "file": "matplotlib/sphinxext/plot_directive.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L368: compile(text, '', 'exec')\nExec: L585: exec('import numpy as np\\n' | L588: exec(str(setup.config.plot_pre_code), ns) | L594: exec(code, ns)" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L440: time.sleep(300)" + }, + { + "package": "networkx", + "file": "networkx/utils/decorators.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L911: compiled = compile(code, filename, \"exec\")\nExec: L912: exec(compiled, globl, locl)" + }, + { + "package": "numba", + "file": "numba/np/ufunc/array_exprs.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L382: code_obj = compile(ast_module, expr_filename, 'exec')\nExec: L383: exec(code_obj, namespace)" + }, + { + "package": "numba", + "file": "numba/tests/support.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)" + }, + { + "package": "numba", + "file": "numba/tests/test_firstlinefinder.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L95: code = compile(source, filename, \"exec\")\nExec: L77: exec(source, globalns) | L98: exec(code, globalns)" + }, + { + "package": "numba", + "file": "numba/tests/test_funcdesc.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L24: compiled = compile(code, filename, 'exec')\nExec: L25: exec(compiled, objs)" + }, + { + "package": "numba", + "file": "numba/tests/test_import.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L33: __import__(mod)\nExec: L43: modlist = set(eval(out.strip())) | L97: modlist = set(eval(out.strip()))" + }, + { + "package": "numba", + "file": "numba/tests/test_np_functions.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)" + }, + { + "package": "numpy", + "file": "numpy/testing/_private/utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1627: code = compile(code_str, f'Test name: {label} ', 'exec')\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)" + }, + { + "package": "numpy", + "file": "numpy/testing/_private/utils.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)" + }, + { + "package": "numpy", + "file": "numpy/tests/test_public_api.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L543: core_submodule = __import__(\nExec: L405: eval(module_name)" + }, + { + "package": "pillow", + "file": "PIL/Image.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:" + }, + { + "package": "protobuf", + "file": "protobuf-3.19.6-nspkg.pth", + "check": "Unusually large executable .pth (539 bytes)", + "severity": "HIGH", + "evidence": "1 import line(s) in 539-byte .pth file" + }, + { + "package": "pygments", + "file": "pygments/formatters/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L38: mod = __import__(module_name, None, None, ['__all__'])\nExec: L103: exec(f.read(), custom_namespace)" + }, + { + "package": "pygments", + "file": "pygments/lexers/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L45: mod = __import__(module_name, None, None, ['__all__'])\nExec: L154: exec(f.read(), custom_namespace)" + }, + { + "package": "pytest", + "file": "_pytest/_py/path.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L626: mod = __import__(hashtype) | L1118: __import__(modname)\nExec: L1153: exec(f.read(), mod.__dict__)" + }, + { + "package": "pytest", + "file": "_pytest/assertion/rewrite.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L393: co = marshal.load(fp) | L395: trace(f\"_read_pyc({source}): marshal.load error {e}\")\nExec: L188: exec(co, module.__dict__)" + }, + { + "package": "scikit-learn", + "file": "sklearn/externals/array_api_compat/torch/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L19: __import__(__package__ + '.linalg') | L20: __import__(__package__ + '.fft')\nExec: L12: exec(f\"{n} = torch.{n}\")" + }, + { + "package": "scipy", + "file": "scipy/optimize/_optimize.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L4155: __import__(mod_name)\nExec: L323: def eval(x):" + }, + { + "package": "setuptools", + "file": "pkg_resources/__init__.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L423: __import__(moduleOrReq) | L1739: code = compile(source, script_filename, 'exec') | L1750: script_code = compile(script_text, script_filename, 'exec')\nExec: L1740: exec(code, namespace, namespace) | L1751: exec(script_code, nam" + }, + { + "package": "setuptools", + "file": "setuptools/_distutils/compilers/C/base.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):" + }, + { + "package": "setuptools", + "file": "setuptools/launch.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L31: code = compile(norm_script, script_name, 'exec')\nExec: L32: exec(code, namespace)" + }, + { + "package": "setuptools", + "file": "setuptools/tests/config/test_pyprojecttoml.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\"," + }, + { + "package": "setuptools", + "file": "setuptools/tests/test_editable_install.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)" + }, + { + "package": "setuptools", + "file": "setuptools/wheel.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L35: NAMESPACE_PACKAGE_INIT = \"__import__('pkg_resources').declare_namespace(__name__)\\n\"\nExec: L191: def eval(req, **env): | L212: (req for req in reqs if for_extra(req) and eval(req, extra=extra))," + }, + { + "package": "six", + "file": "six.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L87: __import__(name)\nExec: L740: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")" + }, + { + "package": "sympy", + "file": "sympy/external/importtools.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L145: mod = __import__(module, **import_kwargs) | L154: __import__(module + '.' + submod)\nExec: L21: return eval(debug_str)" + }, + { + "package": "sympy", + "file": "sympy/plotting/experimental_lambdify.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L249: namespace.update({'math': __import__('math')}) | L251: namespace.update({'cmath': __import__('cmath')}) | L254: namespace.update({'np': __import__('numpy')})\nExec: L268: exec(\"MYNEWLAMBDA = %s\" % eval_str, namespace)" + }, + { + "package": "sympy", + "file": "sympy/utilities/lambdify.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace)" + }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large (1918 KB) JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "" + }, + { + "package": "torch", + "file": "torch/_dynamo/bytecode_debugger.py", + "check": "Anti-analysis/sandbox evasion + suspicious behavior", + "severity": "HIGH", + "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)" + }, + { + "package": "torch", + "file": "torch/_functorch/_aot_autograd/subclass_codegen.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L342: code = compile(source, f\"<{artifact_name}>\", \"exec\")\nExec: L344: exec(code, globals_dict, local_dict)" + }, + { + "package": "torch", + "file": "torch/fx/experimental/rewriter.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)" + }, + { + "package": "torch", + "file": "torch/fx/graph_module.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L106: exec(compile(src, key, \"exec\"), globals)\nExec: L106: exec(compile(src, key, \"exec\"), globals)" + }, + { + "package": "torch", + "file": "torch/package/package_importer.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)" + }, + { + "package": "triton", + "file": "triton/runtime/interpreter.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1435: compiled_code = compile(transformed_ast, filename=self.filename, mode='exec')\nExec: L1441: exec(compiled_code, fn_globals, local_namespace)" + }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "exec/eval with payload hidden in a docstring/string", + "severity": "HIGH", + "evidence": "marshal/compile/obfuscation: L132: r\"|\\bbytearray\\s*\\(\\s*\\[.*?\\]\\s*\\)\" # bytearray([104,101,...]) | L135: r\"|\\bgetattr\\s*\\(\\s*__builtins__\" # getattr(__builtins__, ...)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_compiler_dynamic_exec.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L126: code = compile(source, f\"<{entry_point}>\", \"exec\")\nExec: L134: exec(code, sandbox)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_fused_forward_install.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L268: code = compile(src, fake_path, \"exec\")\nExec: L269: exec(code, namespace)" + }, + { + "package": "unsloth-zoo", + "file": "tests/test_upstream_pinned_symbols_trl_vllm.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L379: mod = __import__(modpath, fromlist=[\"Logprob\"])\nExec: L238: \"unsloth_zoo dispatch via `eval(f'trl.trainer.{trainer_file}.{name}')` breaks\"" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/compiler.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\ | L4292: f\"O^O/ {chr(92)}_/ {c" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/fused_losses/forward_install.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L274: code = compile(new_src, synthetic_path, \"exec\")\nExec: L275: exec(code, ns)" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/mlx/loader.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L1739: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L141: mx.eval(model.parameters()) | L1543: model.eval() | L2126: mx.eval(model.parameters())" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/patching_utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L706: compile(new_source, '', 'exec')\nExec: L221: try: exec(_try_compile_argument) | L226: try: exec(_try_dynamo_argument) | L570: exec(\"from torch._dynamo.compiled_autograd import (\" + \", \".join(x for x in good_" + }, + { + "package": "unsloth-zoo", + "file": "unsloth_zoo/saving_utils.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L3078: module = __import__('transformers', fromlist=[model_class_name])\nExec: L2960: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3006: exec(save_pretrained, globals(), functions)" + }, + { + "package": "werkzeug", + "file": "werkzeug/routing/rules.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)" + } + ] +} diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1b10b557e4..1c7a409bc1 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -277,6 +277,13 @@ "min_p": 0.01, "repetition_penalty": 1.0 }, + "minimax-m2.7": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 40, + "min_p": 0.01, + "repetition_penalty": 1.0 + }, "minimax-m2.5": { "temperature": 1.0, "top_p": 0.95, @@ -390,7 +397,7 @@ "deepseek-r1", "deepseek-v3", "deepseek-ocr", "glm-5", "glm-4", "nemotron", - "minimax-m2.5", "minimax", + "minimax-m2.7", "minimax-m2.5", "minimax", "gpt-oss", "granite-4", "kimi-k2", "kimi", "lfm2", "smollm", "olmo", "falcon", "ernie", "seed", "grok", "mimo" diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index fa5b985513..ee6678d9b9 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -4,9 +4,11 @@ """SQLite storage for auth data (user credentials + JWT secret).""" import hashlib +import hmac import os import secrets import sqlite3 +import threading from datetime import datetime, timezone from typing import Optional, Tuple @@ -236,6 +238,29 @@ def _pbkdf2_desktop_secret(raw_secret: str) -> str: return _pbkdf2_api_key(raw_secret) +# Memoize the deterministic raw-key -> PBKDF2-hash derivation so the 100k-round +# KDF runs once per key instead of on every authenticated request. Keyed by a +# salted HMAC of the key (not the key itself); revocation/expiry are still +# enforced by the SQLite read on every call, so a cache hit only skips the KDF. +# Only keys present in the DB are cached, so unknown-key spam can't grow it. +_api_key_hash_cache: dict[str, str] = {} +_API_KEY_HASH_CACHE_MAX = 4096 +_api_key_hash_cache_lock = threading.Lock() + + +def _api_key_cache_id(raw_key: str) -> str: + """Cache id for a raw key: salted HMAC-SHA256 (not the key itself).""" + return hmac.new( + _get_or_create_api_key_pbkdf2_salt(), raw_key.encode("utf-8"), hashlib.sha256 + ).hexdigest() + + +def _reset_api_key_hash_cache() -> None: + """Drop memoized derivations (tests / salt change).""" + with _api_key_hash_cache_lock: + _api_key_hash_cache.clear() + + def is_initialized() -> bool: """Check if auth is ready for login (at least one user exists in DB).""" conn = get_connection() @@ -704,7 +729,9 @@ def validate_api_key(raw_key: str) -> Optional[str]: Also updates ``last_used_at`` on success. """ - key_hash = _pbkdf2_api_key(raw_key) + cache_id = _api_key_cache_id(raw_key) + cached_hash = _api_key_hash_cache.get(cache_id) + key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key) conn = get_connection() try: cur = conn.execute( @@ -714,6 +741,12 @@ def validate_api_key(raw_key: str) -> Optional[str]: row = cur.fetchone() if row is None: return None + # Real key: memoize so later requests skip the KDF. Bounded; clear on overflow. + if cached_hash is None: + with _api_key_hash_cache_lock: + if len(_api_key_hash_cache) >= _API_KEY_HASH_CACHE_MAX: + _api_key_hash_cache.clear() + _api_key_hash_cache[cache_id] = key_hash if not row["is_active"]: return None if row["expires_at"] is not None: diff --git a/studio/backend/cloudflare_tunnel.py b/studio/backend/cloudflare_tunnel.py index e5dba69452..ef7bacba67 100644 --- a/studio/backend/cloudflare_tunnel.py +++ b/studio/backend/cloudflare_tunnel.py @@ -49,6 +49,16 @@ def _windows_hidden_kwargs() -> dict: return {"creationflags": flags} if flags else {} +def _lifetime_kwargs() -> dict: + """Bind cloudflared to the parent's lifetime (Linux PDEATHSIG). Lazy + + best-effort so this module still loads standalone (storage_roots-style).""" + try: + from utils.process_lifetime import child_popen_kwargs + return child_popen_kwargs() + except Exception: + return {} + + def _asset_name() -> Optional[Tuple[str, bool]]: """(release asset filename, is_tgz) for this OS/arch, or None if unsupported.""" system = platform.system().lower() @@ -233,6 +243,7 @@ class CloudflareTunnel: errors = "replace", bufsize = 1, **_windows_hidden_kwargs(), + **_lifetime_kwargs(), ) self._proc = proc threading.Thread( diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 32523e469e..c238c250bd 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -176,6 +176,9 @@ class JobManager: daemon = True, ) proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) self._mp_q = mp_q self._proc = proc diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index b28b61f088..a0959741a4 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -145,13 +145,19 @@ class ExportBackend: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """ Load a checkpoint for export. + ``hf_token`` authenticates the actual weight load for gated/private + checkpoints, matching the token the worker used for the security preflight + (otherwise a gated repo passes scanning then 401s at from_pretrained). + Returns: Tuple of (success: bool, message: str) """ + token = hf_token if hf_token and hf_token.strip() else None try: logger.info(f"Loading checkpoint: {checkpoint_path}") @@ -169,8 +175,10 @@ class ExportBackend: model_id = base_model or checkpoint_path - self._audio_type = detect_audio_type(model_id) - self.is_vision = not self._audio_type and is_vision_model(model_id) + # Token the type-detection probes too, else a gated multimodal base + # 404s here and falls through to the text loader. + self._audio_type = detect_audio_type(model_id, hf_token = token) + self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token) if self._audio_type == "csm": from unsloth import FastModel @@ -184,6 +192,7 @@ class ExportBackend: auto_model = CsmForConditionalGeneration, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "whisper": @@ -197,6 +206,7 @@ class ExportBackend: load_in_4bit = False, auto_model = WhisperForConditionalGeneration, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "snac": @@ -207,6 +217,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "bicodec": @@ -218,6 +229,7 @@ class ExportBackend: dtype = None if _IS_MLX else torch.float32, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self._audio_type == "dac": @@ -228,6 +240,7 @@ class ExportBackend: max_seq_length = max_seq_length, load_in_4bit = False, trust_remote_code = trust_remote_code, + token = token, ) elif self.is_vision: @@ -238,6 +251,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) tokenizer = processor # vision: processor acts as tokenizer @@ -249,6 +263,7 @@ class ExportBackend: dtype = None, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + token = token, ) if _IS_MLX: diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 20158d1891..636fe1a759 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -63,6 +63,19 @@ class ExportOrchestrator: self._run_start_seq: int = 0 # True while an export op runs; SSE ends the stream 1s after this flips False. self._export_active: bool = False + # Set by cancel_export(); reset when a new load/export run starts. Lets the + # caller distinguish a user cancel from a genuine subprocess crash. + self._cancel_requested: bool = False + + # Last finished operation, so a client whose blocking POST was cut off by a + # Cloudflare tunnel timeout (524 at ~100s, while the op runs for minutes) can + # poll /api/export/status and still learn the real outcome. Guarded by + # _op_lock. `_op_seq` is a monotonic counter the client uses as a baseline to + # tell "my op finished" (seq grew) from a stale previous result. + self._op_lock = threading.Lock() + self._op_seq: int = 0 + self._active_op_kind: Optional[str] = None + self._last_op: Optional[Dict[str, Any]] = None atexit.register(self._cleanup) logger.info("ExportOrchestrator initialized (subprocess mode)") @@ -119,6 +132,72 @@ class ExportOrchestrator: """True while an export / load / cleanup command is running.""" return self._export_active + def was_cancelled(self) -> bool: + """True if the in-flight (or most recent) run was cancelled by the user.""" + return self._cancel_requested + + def _record_op_finished(self, success: bool, message: str, output_path: Optional[str]) -> None: + """Snapshot the just-finished op so status pollers can recover its outcome. + + Called from each op's ``finally`` (with ``_active_op_kind`` still set) BEFORE + ``_export_active`` is cleared, so a status read that observes the op as + inactive is guaranteed to also see this matching result. + """ + with self._op_lock: + self._op_seq += 1 + status = "cancelled" if self._cancel_requested else ("success" if success else "error") + self._last_op = { + "seq": self._op_seq, + "kind": self._active_op_kind, + "status": status, + "output_path": output_path if success else None, + "error": None if success else (message or None), + } + + def get_last_op(self) -> Optional[Dict[str, Any]]: + """Return the last finished op record (or None), for status recovery.""" + with self._op_lock: + return dict(self._last_op) if self._last_op is not None else None + + def get_active_op_kind(self) -> Optional[str]: + """Return the kind of the currently running op (or None when idle).""" + return self._active_op_kind + + def cancel_export(self) -> bool: + """Terminate the in-flight export subprocess immediately. + + An export op holds ``self._lock`` for its whole duration (blocked in + ``_wait_response``), so we deliberately do NOT take the lock here -- we + kill the worker process directly, which unblocks that wait and makes the + in-flight op return a failure the caller surfaces as "cancelled". + + Only the export subprocess is touched; training and inference run in + their own subprocesses and are left untouched. + + Returns True if a live subprocess was terminated, False if none ran. + """ + self._cancel_requested = True + proc = self._proc + if proc is None or not proc.is_alive(): + return False + logger.info( + "Export cancel requested: terminating export subprocess (pid=%s)", + proc.pid, + ) + try: + proc.terminate() + proc.join(timeout = 5) + except Exception: + pass + if proc.is_alive(): + logger.warning("Export subprocess survived terminate, killing") + try: + proc.kill() + proc.join(timeout = 3) + except Exception: + pass + return True + # ------------------------------------------------------------------ # Subprocess lifecycle # ------------------------------------------------------------------ @@ -147,6 +226,9 @@ class ExportOrchestrator: daemon = True, ) self._proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Export subprocess started (pid=%s)", self._proc.pid) def _shutdown_subprocess(self, timeout: float = 10.0) -> None: @@ -301,6 +383,7 @@ class ExportOrchestrator: max_seq_length: int = 2048, load_in_4bit: bool = True, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, hf_token: Optional[str] = None, ) -> Tuple[bool, str]: """Load a checkpoint for export. @@ -312,13 +395,17 @@ class ExportOrchestrator: "max_seq_length": max_seq_length, "load_in_4bit": load_in_4bit, "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "hf_token": hf_token, } with self._lock: # Fresh log buffer so the UI sees only this run's output. self.clear_logs() + self._cancel_requested = False + self._active_op_kind = "load_checkpoint" self._export_active = True + op_success, op_message = False, "" try: # Always kill any existing subprocess and spawn fresh. if self._ensure_subprocess_alive(): @@ -336,6 +423,7 @@ class ExportOrchestrator: self.current_checkpoint = None self.is_vision = False self.is_peft = False + op_success, op_message = False, str(exc) return False, str(exc) if resp.get("success"): @@ -343,15 +431,19 @@ class ExportOrchestrator: self.is_vision = resp.get("is_vision", False) self.is_peft = resp.get("is_peft", False) logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path) - return True, resp.get("message", "Loaded successfully") + op_success, op_message = True, resp.get("message", "Loaded successfully") + return True, op_message else: error = resp.get("message", "Failed to load checkpoint") logger.error("Failed to load checkpoint: %s", error) self.current_checkpoint = None self.is_vision = False self.is_peft = False + op_success, op_message = False, error return False, error finally: + self._record_op_finished(op_success, op_message, None) + self._active_op_kind = None self._export_active = False def export_merged_model( @@ -453,7 +545,10 @@ class ExportOrchestrator: ) self.clear_logs() + self._cancel_requested = False + self._active_op_kind = f"export_{export_type}" self._export_active = True + op_success, op_message, op_output_path = False, "", None try: cmd = {"type": "export", "export_type": export_type, **params} try: @@ -462,14 +557,16 @@ class ExportOrchestrator: f"export_{export_type}_done", timeout = 3600, # GGUF for 30B+ models can take 30+ min ) - return ( - resp.get("success", False), - resp.get("message", ""), - resp.get("output_path"), - ) + op_success = resp.get("success", False) + op_message = resp.get("message", "") + op_output_path = resp.get("output_path") + return op_success, op_message, op_output_path except RuntimeError as exc: + op_success, op_message = False, str(exc) return False, str(exc), None finally: + self._record_op_finished(op_success, op_message, op_output_path) + self._active_op_kind = None self._export_active = False def cleanup_memory(self) -> bool: @@ -481,7 +578,9 @@ class ExportOrchestrator: self.is_peft = False return True + self._active_op_kind = "cleanup" self._export_active = True + success = False try: try: self._send_cmd({"type": "cleanup"}) @@ -498,6 +597,8 @@ class ExportOrchestrator: self.is_peft = False return success finally: + self._record_op_finished(success, "", None) + self._active_op_kind = None self._export_active = False def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]: diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index fb2a893014..7216221f44 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -188,10 +188,16 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: # Auto-enable trust_remote_code for NemotronH/Nano models. if not trust_remote_code: + from utils.security.trusted_org import is_trusted_org_repo + _NEMOTRON_TRUST_SUBSTRINGS = ("nemotron_h", "nemotron-h", "nemotron-3-nano") _cp_lower = checkpoint_path.lower() - if any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) and ( - _cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/") + if ( + any(sub in _cp_lower for sub in _NEMOTRON_TRUST_SUBSTRINGS) + and (_cp_lower.startswith("unsloth/") or _cp_lower.startswith("nvidia/")) + # Genuine first-party Hub repo only (not a local/spoof name starting + # with "unsloth/"); authenticated so private repos resolve. + and is_trusted_org_repo(checkpoint_path, hf_token = cmd.get("hf_token")) ): trust_remote_code = True logger.info( @@ -199,6 +205,81 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: checkpoint_path, ) + # Malware gate: a poisoned pickle deserializes on load even with + # trust_remote_code False, so check HF's security scan (metadata-only) every + # load. Local checkpoints have no Hub scan and are skipped in the helper; a + # LoRA merges its base weights, so gate that repo too. + from utils.security import evaluate_file_security, security_load_subdirs + + malware_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a LOCAL or REMOTE adapter's base so a remote LoRA base is gated too. + _base = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if _base: + malware_targets.append(_base) + except Exception as exc: + logger.debug("Could not resolve LoRA base for malware scan: %s", exc) + _hf_token = cmd.get("hf_token") + for target in dict.fromkeys(malware_targets): + _fs = evaluate_file_security( + target, hf_token = _hf_token, load_subdirs = security_load_subdirs(target, _hf_token) + ) + if _fs.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": _fs.reason, + "error_kind": "malware_blocked", + "security": _fs.response_payload(), + "ts": time.time(), + }, + ) + return + + # Consent gate: scan auto_map code before it runs; block CRITICAL/HIGH unless + # pinned-approved. A LoRA merges its base model, whose code runs, so gate it too. + if trust_remote_code: + from utils.security import evaluate_remote_code_consent_for_targets + + consent_targets = [checkpoint_path] + try: + from utils.models.model_config import get_base_model_from_lora_identifier + + # Resolve a local or remote adapter's base so its base repo is gated too. + base_model = get_base_model_from_lora_identifier(checkpoint_path, cmd.get("hf_token")) + if base_model: + consent_targets.append(base_model) + except Exception as exc: + logger.debug("Could not resolve LoRA base for consent scan: %s", exc) + # Scan adapter + base as one combined unit, pinned by a single fingerprint. + _rc = evaluate_remote_code_consent_for_targets( + consent_targets, + hf_token = cmd.get("hf_token"), + trust_remote_code = True, + approved_fingerprint = cmd.get("approved_remote_code_fingerprint"), + ) + if _rc.blocked: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + f"Checkpoint '{_rc.model_name}' ships custom code flagged as " + f"{_rc.max_severity} by the security scan. Review and " + f"approve it to proceed." + ), + "error_kind": "remote_code_blocked", + "remote_code": _rc.response_payload(), + "ts": time.time(), + }, + ) + return + try: _send_response( resp_queue, @@ -214,6 +295,7 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None: max_seq_length = max_seq_length, load_in_4bit = load_in_4bit, trust_remote_code = trust_remote_code, + hf_token = cmd.get("hf_token"), ) _send_response( diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py new file mode 100644 index 0000000000..f76a38576f --- /dev/null +++ b/studio/backend/core/inference/api_monitor.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Small in-memory monitor for OpenAI-compatible API traffic.""" + +from __future__ import annotations + +import threading +import time +import uuid +from collections import deque +from dataclasses import dataclass +from typing import Any, Optional + + +_MAX_ENTRIES = 50 +_MAX_PROMPT_CHARS = 12000 +_MAX_REPLY_CHARS = 12000 +_PREVIEW_CHARS = 360 + + +def _trim(text: Optional[str], limit: int) -> str: + if not text: + return "" + if len(text) <= limit: + return text + # Guard against limit < 3 (slice would underflow). + if limit <= 3: + return "..."[:limit] + return text[: limit - 3] + "..." + + +@dataclass +class ApiMonitorEntry: + id: str + endpoint: str + method: str + model: str + prompt: str + status: str + started_at: float + updated_at: float + subject: Optional[str] = None + # Monotonic anchors so duration math survives wall-clock steps (NTP). + started_monotonic: float = 0.0 + finished_monotonic: Optional[float] = None + reply: str = "" + finished_at: Optional[float] = None + context_length: Optional[int] = None + prompt_tokens: Optional[int] = None + completion_tokens: Optional[int] = None + total_tokens: Optional[int] = None + total_tokens_authoritative: bool = False + error: Optional[str] = None + + def snapshot(self, *, include_details: bool = True) -> dict[str, Any]: + duration_ms = None + if self.finished_monotonic is not None: + duration_ms = max( + 0, + int((self.finished_monotonic - self.started_monotonic) * 1000), + ) + elif self.finished_at is not None: + duration_ms = max(0, int((self.finished_at - self.started_at) * 1000)) + context_usage = None + if self.total_tokens is not None and self.context_length: + context_usage = min(1.0, max(0.0, self.total_tokens / self.context_length)) + payload = { + "id": self.id, + "endpoint": self.endpoint, + "method": self.method, + "model": self.model, + "prompt_preview": _trim(self.prompt, _PREVIEW_CHARS), + "reply_preview": _trim(self.reply, _PREVIEW_CHARS), + "prompt_truncated": len(self.prompt) > _PREVIEW_CHARS, + "reply_truncated": len(self.reply) > _PREVIEW_CHARS, + "status": self.status, + "started_at": self.started_at, + "updated_at": self.updated_at, + "finished_at": self.finished_at, + "duration_ms": duration_ms, + "context_length": self.context_length, + "context_usage": context_usage, + "prompt_tokens": self.prompt_tokens, + "completion_tokens": self.completion_tokens, + "total_tokens": self.total_tokens, + "error": self.error, + } + if include_details: + payload["prompt"] = self.prompt + payload["reply"] = self.reply + return payload + + +class ApiMonitor: + def __init__(self, max_entries: int = _MAX_ENTRIES): + self._entries: deque[ApiMonitorEntry] = deque() + self._max_entries = max(0, max_entries) + self._lock = threading.Lock() + + def start( + self, + *, + endpoint: str, + method: str, + model: str, + prompt: str, + context_length: Optional[int] = None, + subject: Optional[str] = None, + ) -> str: + now = time.time() + entry = ApiMonitorEntry( + id = f"apireq_{uuid.uuid4().hex[:12]}", + endpoint = endpoint, + method = method, + model = model or "default", + prompt = _trim(prompt, _MAX_PROMPT_CHARS), + status = "running", + started_at = now, + updated_at = now, + subject = subject, + started_monotonic = time.monotonic(), + context_length = context_length, + ) + with self._lock: + self._entries.appendleft(entry) + self._trim_terminal_locked() + return entry.id + + def append_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id or not text: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + # Preview is capped: once the "..." marker is present the head is + # frozen, so skip the per-chunk re-concat (avoids O(n^2) on long + # generations). A reply that landed exactly on the cap has no marker + # yet, so let one more append record the truncation before freezing. + if len(entry.reply) >= _MAX_REPLY_CHARS: + if not entry.reply.endswith("..."): + entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + return + entry.reply = _trim(entry.reply + text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_reply(self, entry_id: Optional[str], text: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + entry.reply = _trim(text, _MAX_REPLY_CHARS) + entry.updated_at = time.time() + + def set_usage( + self, + entry_id: Optional[str], + *, + prompt_tokens: Optional[int] = None, + completion_tokens: Optional[int] = None, + total_tokens: Optional[int] = None, + context_length: Optional[int] = None, + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if prompt_tokens is not None: + entry.prompt_tokens = prompt_tokens + if completion_tokens is not None: + entry.completion_tokens = completion_tokens + if total_tokens is not None: + entry.total_tokens = total_tokens + entry.total_tokens_authoritative = True + elif not entry.total_tokens_authoritative and ( + prompt_tokens is not None or completion_tokens is not None + ): + # Derive only when no authoritative total has been set; + # a later partial chunk must not clobber a provider total. + entry.total_tokens = (entry.prompt_tokens or 0) + (entry.completion_tokens or 0) + if context_length is not None: + entry.context_length = context_length + entry.updated_at = time.time() + + def finish( + self, + entry_id: Optional[str], + status: str = "completed", + ) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + # Idempotent: second call (e.g. [DONE] after the finally block + # already ran) must not move finished_*. + if entry.finished_at is not None: + return + now = time.time() + entry.status = status + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() + + def fail(self, entry_id: Optional[str], error: str) -> None: + if not entry_id: + return + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return + if entry.finished_at is not None: + # Already terminal; refresh error text only. + if error: + entry.error = _trim(error, 1000) + return + now = time.time() + entry.status = "error" + entry.error = _trim(error, 1000) + entry.updated_at = now + entry.finished_at = now + entry.finished_monotonic = time.monotonic() + self._entries.remove(entry) + self._entries.appendleft(entry) + self._trim_terminal_locked() + + def snapshot( + self, + *, + include_details: bool = True, + subject: Optional[str] = None, + ) -> list[dict[str, Any]]: + with self._lock: + return [ + entry.snapshot(include_details = include_details) + for entry in self._entries + if subject is None or entry.subject == subject + ] + + def get( + self, + entry_id: str, + *, + subject: Optional[str] = None, + ) -> Optional[dict[str, Any]]: + with self._lock: + entry = self._find_locked(entry_id) + if entry is None: + return None + if subject is not None and entry.subject != subject: + return None + return entry.snapshot(include_details = True) + + def active_count(self, *, subject: Optional[str] = None) -> int: + with self._lock: + return sum( + 1 + for entry in self._entries + if entry.status == "running" and (subject is None or entry.subject == subject) + ) + + def clear(self) -> None: + with self._lock: + self._entries.clear() + + def _find_locked(self, entry_id: str) -> Optional[ApiMonitorEntry]: + for entry in self._entries: + if entry.id == entry_id: + return entry + return None + + def _trim_terminal_locked(self) -> None: + terminal_seen = 0 + kept: deque[ApiMonitorEntry] = deque() + for entry in self._entries: + if entry.status == "running": + kept.append(entry) + continue + if terminal_seen < self._max_entries: + kept.append(entry) + terminal_seen += 1 + self._entries = kept + + +api_monitor = ApiMonitor() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 57cc97f3b5..4826403bbc 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -23,18 +23,19 @@ import sys import threading import time from pathlib import Path -from typing import Callable, Generator, Iterable, List, Optional +from typing import Callable, Collection, Generator, Iterable, List, Mapping, Optional, Union import httpx from core.inference.llama_server_args import ( + _effective_tensor_parallel, + _tensor_parallel_matches_loaded, extra_args_disable_mmproj, parse_cache_override, + parse_cache_override_per_axis, parse_ctx_override, parse_split_mode_override, - resolve_cache_type_kv, resolve_requested_ctx, - resolve_tensor_parallel, strip_shadowing_flags, strip_split_mode_only, ) @@ -51,9 +52,11 @@ from core.tool_healing import ( strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret +from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from utils.process_lifetime import child_popen_kwargs as _child_popen_kwargs from core.inference.tool_call_parser import ( RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, @@ -75,6 +78,112 @@ from state.tool_approvals import ( logger = get_logger(__name__) +class LlamaServerNotFoundError(RuntimeError): + """GGUF model needs the llama.cpp runtime but no llama-server is installed. + Subclasses RuntimeError so existing handlers still catch it.""" + + +# Shared so the from_identifier preflight and the load-time raise stay in sync. +LLAMA_SERVER_NOT_FOUND_DETAIL = ( + "This is a GGUF model, but the llama.cpp runtime (llama-server) is not " + "installed. Run `unsloth studio setup` to download the prebuilt runtime, " + "then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)" +) + + +# llama-server can serve HTTP 200 while running a model entirely on CPU when a +# GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so +# Studio can warn. Priority: explicit "offloaded N/M layers to GPU" counts +# (authoritative), then GPU "model buffer size" lines (host-pinned _Host +# excluded), then the "device_info:" device table (disconfirm only). +_GPU_OFFLOAD_MARKERS = ( + "CUDA", + "ROCm", + "ROCM", + "HIP", + "Metal", + "Vulkan", + "OpenCL", + "SYCL", + "MUSA", + "CANN", +) +_OFFLOADED_LAYERS_RE = re.compile( + r"offloaded\s+(\d+)\s*/\s*(\d+)\s+layers?\s+to\s+gpu", re.IGNORECASE +) +_DEVICE_ROW_RE = re.compile( + r"-\s*(CUDA|ROCm|ROCM|HIP|Metal|Vulkan|SYCL|OpenCL|MUSA|CANN|CPU)\w*\s*:", + re.IGNORECASE, +) +_GPU_DEVICE_PREFIXES = ( + "cuda", + "rocm", + "hip", + "metal", + "vulkan", + "sycl", + "opencl", + "musa", + "cann", +) + + +def classify_gpu_offload_lines(lines: "list[str]") -> Optional[bool]: + """True if the model landed on a GPU, False if it stayed on CPU despite GPU + intent, None when the log has no usable signal.""" + # Counted offload is authoritative, keyed on the model with the most layers. + # A separate MTP/draft model logs its own (much smaller) "offloaded N/M" + # line, so decide on the largest-M line: a drafter that fits on GPU must not + # mask a main model running on CPU. N>0 on that model is True, 0 is False. + max_total = -1 + offloaded_at_max = 0 + for line in lines: + match = _OFFLOADED_LAYERS_RE.search(line) + if not match: + continue + offloaded, total = int(match.group(1)), int(match.group(2)) + if total > max_total or (total == max_total and offloaded > offloaded_at_max): + max_total, offloaded_at_max = total, offloaded + if max_total >= 0: + return offloaded_at_max > 0 + + # GPU marker on a *model* buffer; _Host buffers are CPU-pinned, not offload. + # Buffer lines are authoritative: present but none on a GPU means CPU-only, + # so do not let the device table below override that. + saw_model_buffer = False + for line in lines: + if "model buffer size" not in line: + continue + saw_model_buffer = True + if "_Host" not in line and any(m in line for m in _GPU_OFFLOAD_MARKERS): + return True + if saw_model_buffer: + return False + + # device_info: lists *available* devices (printed whenever a GPU backend is + # visible), not where the model loaded, so it can only disconfirm: an + # all-CPU table means no usable GPU. A visible GPU device is not proof the + # model used it, so it does not return True. Rows after the header only. + after_header = False + saw_device_row = False + saw_gpu_device = False + for line in lines: + if "device_info:" in line: + after_header = True + continue + if not after_header: + continue + match = _DEVICE_ROW_RE.search(line) + if not match: + continue + saw_device_row = True + if match.group(1).lower().startswith(_GPU_DEVICE_PREFIXES): + saw_gpu_device = True + if saw_device_row and not saw_gpu_device: + return False + return None + + def _wsl_system_rocm_lib_dirs() -> "list[str]": """System ROCm lib dir(s) to load before a prebuilt's bundled HIP, on WSL. @@ -154,8 +263,8 @@ def _should_suppress_forced_no_tool_output(text: str) -> bool: # ── Pre-compiled patterns for GGUF shard detection ─────────── -_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$") -_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") +_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$", re.IGNORECASE) +_SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$", re.IGNORECASE) # ── Sliding-window-pattern resolver ─────────────────────────── @@ -443,6 +552,27 @@ _TOOL_TEMPLATE_MARKERS = ( ) +# Canonical reasoning_effort levels, weakest -> strongest. Used to read the +# discrete set a template branches on (e.g. GLM-5.2 uses 'high' | 'max') so we +# only ever offer levels the template actually understands. +_REASONING_EFFORT_SCALE = ("minimal", "low", "medium", "high", "max") + + +def _extract_reasoning_effort_levels(chat_template: str) -> list: + """Return the reasoning_effort levels a template references, in canonical + (weakest -> strongest) order. + + Looks for the quoted literals (e.g. ``'high'`` / ``"max"``) the template + compares ``reasoning_effort`` against, so we surface exactly the levels it + branches on and nothing else. + """ + return [ + level + for level in _REASONING_EFFORT_SCALE + if f"'{level}'" in chat_template or f'"{level}"' in chat_template + ] + + def detect_reasoning_flags( chat_template: Optional[str], model_identifier: Optional[str] = None, @@ -462,6 +592,7 @@ def detect_reasoning_flags( "supports_reasoning": False, "reasoning_style": "enable_thinking", "reasoning_always_on": False, + "reasoning_effort_levels": [], "supports_preserve_thinking": False, "supports_tools": False, } @@ -470,7 +601,25 @@ def detect_reasoning_flags( tpl = chat_template prefix = f"{log_source}: " if log_source else "" - if "enable_thinking" in tpl: + effort_levels = ( + _extract_reasoning_effort_levels(tpl) + if ("reasoning_effort" in tpl and "enable_thinking" in tpl) + else [] + ) + if "enable_thinking" in tpl and "reasoning_effort" in tpl and effort_levels: + # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort + # level among a discrete set (e.g. 'high' | 'max'). Distinct from + # gpt-oss (reasoning_effort only, no on/off gate) and Qwen + # (enable_thinking only). Disabling is enable_thinking=false; the levels + # are the quoted effort literals the template actually branches on. + flags["supports_reasoning"] = True + flags["reasoning_style"] = "enable_thinking_effort" + flags["reasoning_effort_levels"] = effort_levels + logger.info( + f"{prefix}model supports reasoning " + f"(enable_thinking + reasoning_effort: {effort_levels})" + ) + elif "enable_thinking" in tpl: flags["supports_reasoning"] = True flags["reasoning_style"] = "enable_thinking" logger.info(f"{prefix}model supports reasoning (enable_thinking)") @@ -533,22 +682,160 @@ def _is_companion_gguf_path(path: str) -> bool: return name.startswith("mtp-") or "/mtp/" in f"/{p}" +_BIG_ENDIAN_GGUF_FILENAME_RE = re.compile(r"(^|[-_])be(?:[._-]|$)", re.IGNORECASE) +_GGUF_KNOWN_QUANT_RE = re.compile( + r"(UD-)?" + r"(MXFP[0-9]+(?:_[A-Z0-9]+)*" + r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" + r"|TQ[0-9]+_[0-9]+" + r"|Q[0-9]+_K_[A-Z]+" + r"|Q[0-9]+_[0-9]+" + r"|Q[0-9]+_K" + r"|BF16|F16|F32)", + re.IGNORECASE, +) + + +def _is_big_endian_gguf_path(path: str, variant_key: str = "") -> bool: + normalized = path.replace("\\", "/") + name = normalized.rsplit("/", 1)[-1] + stem = name.rsplit(".", 1)[0].lower() + variant_key = variant_key.strip().lower() + variant_index = stem.find(variant_key) if variant_key else -1 + parent = normalized.rsplit("/", 1)[0].lower() if "/" in normalized else "" + variant_in_parent_only = ( + bool(parent) + and variant_index < 0 + and ( + (variant_key and variant_key in parent) + or (not variant_key and _GGUF_KNOWN_QUANT_RE.search(parent) is not None) + ) + ) + for match in _BIG_ENDIAN_GGUF_FILENAME_RE.finditer(stem): + if variant_index >= 0 and variant_index < match.start(): + return True + tail = stem[match.end() :].lstrip("._-") + if not tail or _GGUF_KNOWN_QUANT_RE.search(tail) is None: + return not variant_in_parent_only + return False + + +def _gguf_snapshot_files(snapshot: Path) -> list[str]: + return [ + p.relative_to(snapshot).as_posix() + for p in snapshot.rglob("*") + if p.is_file() and p.name.lower().endswith(".gguf") + ] + + +def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: + m = _SHARD_FULL_RE.match(first_shard) + if not m: + return [] + prefix = m.group(1) + total = m.group(3) + sibling_pat = re.compile( + r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(total) + r"\.gguf$", + re.IGNORECASE, + ) + return sorted(f for f in files if f != first_shard and sibling_pat.match(f)) + + +def _gguf_files_for_variant(files: Iterable[str], variant: str) -> list[str]: + """Return main GGUF files matching a requested variant. + + Prefer exact quant-label matches over loose substring matches so a request + for ``stories260K`` does not resolve to ``stories260K-be.gguf``. + """ + variant_key = variant.strip().lower() + main_files = [ + f + for f in files + if f.lower().endswith(".gguf") + and not _is_companion_gguf_path(f) + and not _is_big_endian_gguf_path(f, variant_key) + ] + if not variant_key: + return sorted(main_files) + + try: + from utils.models.model_config import _extract_quant_label + except Exception: + _extract_quant_label = None + + if _extract_quant_label is not None: + try: + exact = sorted(f for f in main_files if _extract_quant_label(f).lower() == variant_key) + if exact: + return exact + except Exception as e: + logger.warning("Failed to extract GGUF quant labels: %s", e) + + boundary = re.compile(r"(? float: + """Bytes per KV-cache element for a llama.cpp cache type (f16 default).""" + return { + "f32": 4.0, + "f16": 2.0, + "bf16": 2.0, + "q8_0": 34 / 32, + "q5_1": 0.75, + "q5_0": 0.6875, + "q4_1": 0.625, + "q4_0": 0.5625, + "iq4_nl": 0.5625, + }.get((cache_type or "f16").strip().lower(), 2.0) + + +def _env_main_cache_type_for_budget(env: Optional[Mapping[str, str]] = None) -> Optional[str]: + """Heavier of the inherited LLAMA_ARG_CACHE_TYPE_K/_V env types when it + exceeds the f16 default, else None. Studio emits --cache-type only for the + param/extras path, so a heavier env (f32) would otherwise reach the child + unbudgeted; quantized env types stay over-reserved by f16 (-> None).""" + e = os.environ if env is None else env + f16_bpe = _kv_bytes_per_elem("f16") + heaviest: Optional[str] = None + heaviest_bpe = f16_bpe + for var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + raw = (e.get(var) or "").strip().lower() + if not raw: + continue + bpe = _kv_bytes_per_elem(raw) + if bpe > heaviest_bpe: + heaviest, heaviest_bpe = raw, bpe + return heaviest + + +def _extra_args_main_cache_type_for_budget(extra_args: Optional[Iterable[str]]) -> Optional[str]: + """Heavier (max bytes/elem) of the explicit --cache-type-k/-v extras, or None. + + Extras are appended last and win per axis, so an asymmetric K=f32,V=f16 must be + budgeted by its heavier axis. resolve_cache_type_kv returns only the last-wins + single type, which under-reserves the heavier axis when the lighter one is last.""" + k, v = parse_cache_override_per_axis(extra_args) + candidates = [c for c in (k, v) if c] + if not candidates: + return None + return max(candidates, key = _kv_bytes_per_elem) + + def _auto_mode_drops_mtp( req_mode: Optional[str], size_b: Optional[float], @@ -567,18 +854,221 @@ def _auto_mode_drops_mtp( def _extra_args_set_spec_type(extra_args: Optional[Iterable[str]]) -> bool: """User passed --spec-type / --spec-default? llama-server takes one --spec-type (comma-separated to chain), so suppress auto-emit.""" + return _extra_args_set_any_flag(extra_args, {"--spec-type", "--spec-default"}) + + +_GPU_OFFLOAD_OVERRIDE_FLAGS = frozenset({"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}) +_THREAD_OVERRIDE_FLAGS = frozenset({"-t", "--threads"}) + + +def _extra_arg_flag_name(token: str) -> Optional[str]: + if not token.startswith("-") or token in {"-", "--"}: + return None + if len(token) >= 2 and (token[1].isdigit() or token[1] == "."): + return None + return token.split("=", 1)[0] + + +def _extra_args_set_any_flag(extra_args: Optional[Iterable[str]], flags: Collection[str]) -> bool: if not extra_args: return False for raw in extra_args: - tok = str(raw) - if not tok.startswith("--"): - continue - flag = tok.split("=", 1)[0] - if flag in ("--spec-type", "--spec-default"): + flag = _extra_arg_flag_name(str(raw)) + if flag in flags: return True return False +def _effective_spec_type( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """The --spec-type llama-server will use: the last CLI --spec-type (or + --spec-default, which resolves non-MTP), else LLAMA_ARG_SPEC_TYPE. A CLI flag + overrides the env (matching llama.cpp), so a stale MTP env can't make the + budget reserve a drafter the launch won't load. None if neither sets it.""" + args = [str(a) for a in extra_args] if extra_args else [] + cli_present = False + cli_value: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag == "--spec-default": + cli_present = True + cli_value = "default" + continue + if flag != "--spec-type": + continue + cli_present = True + cli_value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if cli_present: + return cli_value + return (os.environ if env is None else env).get("LLAMA_ARG_SPEC_TYPE") + + +def _extra_args_requests_mtp( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects MTP (mtp/draft-mtp), so the + budget must reserve for it.""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("mtp", "draft-mtp") for p in value.split(",")) + + +def _extra_args_requests_separate_draft( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the effective --spec-type selects a non-MTP model draft mode + (draft-simple/draft-eagle3), which loads a separate draft model the budget + must reserve (draft-mtp -> _extra_args_requests_mtp; ngram-* load no model).""" + value = _effective_spec_type(extra_args, env) + if not value: + return False + return any(p.strip().lower() in ("draft-simple", "draft-eagle3") for p in value.split(",")) + + +def _extra_args_spec_draft_n_max(extra_args: Optional[Iterable[str]]) -> Optional[int]: + """Draft depth from extras (``--spec-draft-n-max`` or legacy ``--draft-max``), else None.""" + if not extra_args: + return None + args = [str(a) for a in extra_args] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--spec-draft-n-max", "--draft-max"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + return found + + +def _extra_args_mtp_draft_path( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[str]: + """Separate drafter path from extras (local --model-draft/-md or HF + --spec-draft-hf/-hfd/...), else the LLAMA_ARG_SPEC_DRAFT_MODEL/_HF_REPO env, + else None. An HF repo isn't a local file, so the budget can't size it (falls + back to the flat reserve), but recognizing it avoids sizing the wrong one.""" + flags = { + "--model-draft", + "--spec-draft-model", + "-md", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", + } + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if value and not value.startswith("-"): + found = value + if found is not None: + return found + e = os.environ if env is None else env + return e.get("LLAMA_ARG_SPEC_DRAFT_MODEL") or e.get("LLAMA_ARG_SPEC_DRAFT_HF_REPO") or None + + +def _extra_args_draft_cache_types( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> tuple[Optional[str], Optional[str]]: + """Draft KV cache types (k_type, v_type), each from extras else the + LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K/_V env, else None (f16). K and V are + independent: a one-sided override must not apply to both.""" + args = [str(a) for a in extra_args] if extra_args else [] + k_flags = {"--cache-type-k-draft", "--spec-draft-type-k", "-ctkd"} + v_flags = {"--cache-type-v-draft", "--spec-draft-type-v", "-ctvd"} + k_type: Optional[str] = None + v_type: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in k_flags and flag not in v_flags: + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if not value or value.startswith("-"): + continue + if flag in k_flags: + k_type = value + else: + v_type = value + e = os.environ if env is None else env + if k_type is None: + k_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K") or None + if v_type is None: + v_type = e.get("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V") or None + return k_type, v_type + + +def _extra_args_draft_offloaded_to_cpu( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True if the SEPARATE draft model is on CPU (so the budget must not charge + its weights+KV): --spec-draft-ngl 0, or --spec-draft-device naming only + cpu/none, else the LLAMA_ARG_N_GPU_LAYERS_DRAFT env the child honors (the + device flag has no env). An embedded MTP head follows the main -ngl, so these + draft-only flags don't move it. Last-wins, so only each flag's final value counts.""" + ngl_flags = {"--spec-draft-ngl", "-ngld", "--gpu-layers-draft", "--n-gpu-layers-draft"} + dev_flags = {"--spec-draft-device", "-devd", "--device-draft"} + args = [str(a) for a in extra_args] if extra_args else [] + last_ngl: Optional[str] = None + last_dev: Optional[str] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if flag in ngl_flags: + last_ngl = value + elif flag in dev_flags: + last_dev = value + if last_ngl is None: + last_ngl = (os.environ if env is None else env).get("LLAMA_ARG_N_GPU_LAYERS_DRAFT") + if last_ngl is not None: + try: + if int(last_ngl) == 0: + return True + except (TypeError, ValueError): + pass + if last_dev is not None: + devs = [d.strip().lower() for d in last_dev.split(",") if d.strip()] + if devs and all(d in ("cpu", "none") for d in devs): + return True + return False + + +def _extra_args_n_ubatch( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> Optional[int]: + """Physical micro-batch from extras (--ubatch-size/-ub) else the LLAMA_ARG_UBATCH + env, else None. It sizes the compute-graph buffer, so an override must reach + the VRAM reserve.""" + args = [str(a) for a in extra_args] if extra_args else [] + found: Optional[int] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + if flag not in ("--ubatch-size", "-ub"): + continue + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + try: + found = int(value) + except (TypeError, ValueError): + continue + if found is not None: + return found + raw = (os.environ if env is None else env).get("LLAMA_ARG_UBATCH") + if raw: + try: + return int(raw) + except (TypeError, ValueError): + pass + return None + + def _build_ngram_mod_flags( caps: Optional[dict], n_match: int = 24, @@ -697,9 +1187,9 @@ class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. Lifecycle: - 1. load_model() — start llama-server with the GGUF file - 2. generate_chat_completion() — proxy to /v1/chat/completions, stream back - 3. unload_model() — terminate the subprocess + 1. load_model(): start llama-server with the GGUF file + 2. generate_chat_completion(): proxy to /v1/chat/completions, stream back + 3. unload_model(): terminate the subprocess """ def __init__(self): @@ -724,6 +1214,7 @@ class LlamaCppBackend: self._is_diffusion: bool = False self._diffusion_visual_bin: Optional[str] = None self._healthy = False + self._stats_logger = None # vLLM-style engine-stats poller, set on load # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None self._context_length: Optional[int] = None @@ -734,6 +1225,7 @@ class LlamaCppBackend: self._supports_reasoning: bool = False self._reasoning_always_on: bool = False self._reasoning_style: str = "enable_thinking" + self._reasoning_effort_levels: list = [] self._supports_preserve_thinking: bool = False self._supports_tools: bool = False self._cache_type_kv: Optional[str] = None @@ -754,6 +1246,9 @@ class LlamaCppBackend: self._n_kv_heads_by_layer: Optional[list[int]] = None self._n_heads: Optional[int] = None self._embedding_length: Optional[int] = None + # For the compute-graph buffer estimate; vocab from the tokens array len. + self._feed_forward_length: Optional[int] = None + self._vocab_size: Optional[int] = None # Architecture-aware KV fields for 5-path estimation self._kv_key_length: Optional[int] = None self._kv_value_length: Optional[int] = None @@ -773,8 +1268,12 @@ class LlamaCppBackend: self._nextn_predict_layers: Optional[int] = None self._lock = threading.Lock() # Wraps load_model() end-to-end so concurrent loads serialise and never - # coexist as two llama-server processes (#5401). - self._serial_load_lock = threading.Lock() + # coexist as two llama-server processes (#5401). RLock so MTP-crash + # recovery can re-acquire it for its nested load_model. + self._serial_load_lock = threading.RLock() + # Serialises mid-session respawns so many generations hitting a killed + # server trigger at most one reload (see _respawn_if_dead). + self._respawn_lock = threading.Lock() # Set by the in-app updater while it swaps prebuilt binaries; load_model() # rejects fast so no server starts from a half-swapped binary. self._llama_update_in_progress = False @@ -785,6 +1284,18 @@ class LlamaCppBackend: self._extra_args: Optional[List[str]] = None self._extra_args_source: Optional[tuple[str, Optional[str]]] = None self._requested_n_ctx: int = 0 + # Raw kwargs of the last healthy load, for the MTP-crash reload. Memory-only + # (carries hf_token, never logged); single-flight via the lock below. + self._last_load_kwargs: Optional[dict] = None + self._mtp_runtime_fallback_lock = threading.Lock() + self._mtp_runtime_fallback_in_progress = False + # Background watchdog so an MTP+tensor crash recovers even when no request + # observes it (direct proxy endpoints, or nothing in flight). + self._mtp_watchdog_thread: Optional[threading.Thread] = None + self._mtp_watchdog_stop = threading.Event() + # True when the launch actually runs MTP+tensor (Studio- or user/env-driven); + # gates the probe, watchdog, and recovery so pass-through MTP is covered. + self._mtp_runtime_fallback_active = False self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None # llama-server tee log (see _drain_stdout / _kill_process). @@ -803,7 +1314,11 @@ class LlamaCppBackend: # to decide whether to wait for the VRAM reclaim to finish. self._last_kill_monotonic: float = 0.0 - self._kill_orphaned_servers() + _reaped = self._kill_orphaned_servers() + if _reaped: + # Reaped VRAM frees lazily; arm the settle wait so the first load + # waits before ranking GPUs by free memory. + self._last_kill_monotonic = time.monotonic() atexit.register(self._cleanup) # ── Properties ──────────────────────────────────────────────── @@ -936,12 +1451,13 @@ class LlamaCppBackend: m = _SHARD_RE.match(stem) prefix = m.group(1) if m else None if prefix and parent.is_dir(): + prefix_lower = prefix.lower() for sibling in parent.iterdir(): if ( sibling.is_file() - and sibling.name.startswith(prefix) + and sibling.name.lower().startswith(prefix_lower) and sibling.name != stem - and sibling.suffix == ".gguf" + and sibling.suffix.lower() == ".gguf" ): try: bytes_total += sibling.stat().st_size @@ -993,6 +1509,12 @@ class LlamaCppBackend: def reasoning_style(self) -> str: return self._reasoning_style + @property + def reasoning_effort_levels(self) -> list: + """Discrete reasoning_effort levels the template offers (e.g. GLM-5.2's + ['high', 'max']). Empty unless reasoning_style == 'enable_thinking_effort'.""" + return self._reasoning_effort_levels + @property def supports_preserve_thinking(self) -> bool: return self._supports_preserve_thinking @@ -1002,6 +1524,10 @@ class LlamaCppBackend: return self._reasoning_default def _reasoning_kwargs(self, enable_thinking: bool) -> dict: + if self._reasoning_style == "enable_thinking_effort": + # GLM-5.2-style: enable_thinking is the on/off gate; when on, leave + # the template's default effort (max) in place. + return {"enable_thinking": enable_thinking} if self._reasoning_style == "reasoning_effort": return {"reasoning_effort": "high" if enable_thinking else "low"} return {"enable_thinking": enable_thinking} @@ -1022,7 +1548,20 @@ class LlamaCppBackend: # Always-on reasoning models hardcode tags and don't consume # enable_thinking / reasoning_effort -- skip. if self._supports_reasoning and not self._reasoning_always_on: - if self._reasoning_style == "reasoning_effort": + if self._reasoning_style == "enable_thinking_effort": + # GLM-5.2-style: enable_thinking gates thinking on/off, and the + # reasoning_effort level (e.g. 'high' | 'max') is only meaningful + # while thinking is on. Disabling is enable_thinking=false; a raw + # API caller can also disable via the OpenAI-style + # reasoning_effort="none" sentinel. We never coerce off into a + # 'low' effort the way gpt-oss does (those models genuinely + # cannot disable). + thinking_off = enable_thinking is False or reasoning_effort == "none" + if enable_thinking is not None or reasoning_effort == "none": + kwargs["enable_thinking"] = not thinking_off + if not thinking_off and reasoning_effort in self._reasoning_effort_levels: + kwargs["reasoning_effort"] = reasoning_effort + elif self._reasoning_style == "reasoning_effort": if reasoning_effort in ("none", "low", "medium", "high"): kwargs["reasoning_effort"] = reasoning_effort elif reasoning_effort == "minimal": @@ -1217,7 +1756,7 @@ class LlamaCppBackend: def probe_server_capabilities(cls, binary: Optional[str] = None) -> dict[str, object]: """Parse `llama-server --help` for feature flags. Returns {found, mtp_token, supports_mtp, ngram_mod_flavor, - supports_ngram_mod, spec_draft_n_max_flag}. + supports_ngram_mod, spec_draft_n_max_flag, cache flag support}. ``ngram_mod_flavor``: ``"new"`` when the post-rename ``--spec-ngram-mod-n-match / -n-min / -n-max`` are real args; @@ -1242,6 +1781,10 @@ class LlamaCppBackend: "spec_draft_n_max_flag": None, "supports_kv_unified": False, "supports_fit_ctx": False, + "supports_cache_ram": False, + "supports_ctx_checkpoints": False, + "supports_no_cache_prompt": False, + "supports_metrics": False, } try: mtime = int(Path(bin_path).stat().st_mtime) @@ -1257,13 +1800,20 @@ class LlamaCppBackend: spec_draft_n_max_flag: Optional[str] = None supports_kv_unified = False supports_fit_ctx = False + supports_cache_ram = False + supports_ctx_checkpoints = False + supports_no_cache_prompt = False + supports_metrics = False try: + probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( [bin_path, "--help"], capture_output = True, text = True, + errors = "replace", timeout = 10, check = False, + env = probe_env, ) help_text = (result.stdout or "") + "\n" + (result.stderr or "") # Split into per-flag blocks (each --flag line + its indented @@ -1347,6 +1897,10 @@ class LlamaCppBackend: supports_kv_unified = _is_real("--kv-unified") supports_fit_ctx = _is_real("--fit-ctx") + supports_cache_ram = _is_real("--cache-ram") + supports_ctx_checkpoints = _is_real("--ctx-checkpoints") + supports_no_cache_prompt = _is_real("--no-cache-prompt") + supports_metrics = _is_real("--metrics") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") @@ -1359,6 +1913,10 @@ class LlamaCppBackend: "spec_draft_n_max_flag": spec_draft_n_max_flag, "supports_kv_unified": supports_kv_unified, "supports_fit_ctx": supports_fit_ctx, + "supports_cache_ram": supports_cache_ram, + "supports_ctx_checkpoints": supports_ctx_checkpoints, + "supports_no_cache_prompt": supports_no_cache_prompt, + "supports_metrics": supports_metrics, } cls._capability_cache[cache_key] = info return info @@ -1376,7 +1934,8 @@ class LlamaCppBackend: if m: prefix, _, num_total = m.group(1), m.group(2), m.group(3) sibling_pat = re.compile( - r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$" + r"^" + re.escape(prefix) + r"-\d{5}-of-" + re.escape(num_total) + r"\.gguf$", + re.IGNORECASE, ) for sibling in main.parent.iterdir(): if sibling != main and sibling_pat.match(sibling.name): @@ -1385,11 +1944,12 @@ class LlamaCppBackend: return total @staticmethod - def _amd_apu_wants_unified_memory() -> bool: + def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: """True only for AMD unified-memory APUs (gfx1150/gfx1151), where - GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM. - False elsewhere (the env hurts discrete GPUs). ROCm reuses torch.cuda.*; - gcnArchName suffix is stripped.""" + GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it + hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the + selected GPUs, so a dGPU on a mixed host is not treated as unified-memory; + None means every visible GPU.""" try: import torch @@ -1397,12 +1957,39 @@ class LlamaCppBackend: return False if not (hasattr(torch, "cuda") and torch.cuda.is_available()): return False - for _i in range(torch.cuda.device_count()): + # Map visible ordinal -> physical id via the active ROCm mask (HIP, + # then ROCR, then CUDA), mirroring _get_gpu_memory's ROCm branch. + physical_ids: Optional[list[int]] = None + hip_v = os.environ.get("HIP_VISIBLE_DEVICES") + rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") + cvd = ( + hip_v + if hip_v is not None + else rocr_v + if rocr_v is not None + else os.environ.get("CUDA_VISIBLE_DEVICES") + ) + if cvd is not None: try: - _arch = getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "") or "" + physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] + except ValueError: + physical_ids = None + arch_by_id: dict[int, str] = {} + for ordinal in range(torch.cuda.device_count()): + try: + _arch = ( + getattr(torch.cuda.get_device_properties(ordinal), "gcnArchName", "") or "" + ) except Exception: continue - if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}: + pid = ( + physical_ids[ordinal] + if physical_ids is not None and ordinal < len(physical_ids) + else ordinal + ) + arch_by_id[pid] = _arch.split(":")[0].strip().lower() + for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id): + if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}: return True except Exception: return False @@ -1509,7 +2096,14 @@ class LlamaCppBackend: @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: - """Query free memory per GPU. + """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by + index; empty if no supported GPU is reachable. Thin wrapper over + ``_get_gpu_memory`` for callers that only need free VRAM.""" + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + + @staticmethod + def _get_gpu_memory() -> list[tuple[int, int, int]]: + """Query free AND total memory per GPU. Order: 1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects @@ -1520,15 +2114,15 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. - Returns list of (gpu_index, free_mib) sorted by index; empty if no - supported GPU is reachable. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no + supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. """ # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( [ "nvidia-smi", - "--query-gpu=index,memory.free", + "--query-gpu=index,memory.free,memory.total", "--format=csv,noheader,nounits", ], capture_output = True, @@ -1548,15 +2142,30 @@ class LlamaCppBackend: allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) except ValueError: pass - gpus: list[tuple[int, int]] = [] + gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): - parts = line.split(",") - if len(parts) == 2: - idx = int(parts[0].strip()) - free_mib = int(parts[1].strip()) - if allowed is not None and idx not in allowed: - continue - gpus.append((idx, free_mib)) + parts = [p.strip() for p in line.split(",")] + if len(parts) < 2: + continue + # Index and free required; skip a bad line rather than abandon + # the probe to the torch fallback. + try: + idx = int(parts[0]) + free_mib = int(parts[1]) + except ValueError: + continue + # Total parsed separately: a two-column line or a non-integer + # total ("N/A" on MIG/vGPU) keeps the GPU at total 0 (fit uses + # the free*frac fallback) instead of dropping it. + total_mib = 0 + if len(parts) >= 3 and parts[2]: + try: + total_mib = int(parts[2]) + except ValueError: + total_mib = 0 + if allowed is not None and idx not in allowed: + continue + gpus.append((idx, free_mib, total_mib)) # Match the docstring's sort-by-id guarantee (driver order isn't). gpus.sort(key = lambda g: g[0]) if gpus: @@ -1601,19 +2210,61 @@ class LlamaCppBackend: physical_ids = None gpus = [] for ordinal in range(torch.cuda.device_count()): - free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal) + free_bytes, total_bytes = torch.cuda.mem_get_info(ordinal) idx = ( physical_ids[ordinal] if physical_ids is not None and ordinal < len(physical_ids) else ordinal ) - gpus.append((idx, free_bytes // (1024 * 1024))) + gpus.append((idx, free_bytes // (1024 * 1024), total_bytes // (1024 * 1024))) # Match the nvidia-smi path's docstring guarantee of sorted-by-id. return sorted(gpus, key = lambda g: g[0]) except Exception as e: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _available_system_memory_mib() -> Optional[int]: + """Available system RAM in MiB (psutil, then /proc/meminfo), or None if + neither is readable. On a unified-memory APU this, not the ROCm-reported + VRAM, is the real ceiling: the weights load into shared system RAM.""" + try: + import psutil + return int(psutil.virtual_memory().available // (1024 * 1024)) + except Exception: + pass + try: + with open("/proc/meminfo") as f: + for line in f: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) // 1024 # kB -> MiB + except Exception: + pass + return None + + @staticmethod + def _apu_ram_shortfall_message( + model_size_bytes: int, + avail_mib: Optional[int], + headroom_mib: int = 2048, + ) -> Optional[str]: + """On a unified-memory APU, return a user-facing refusal when the weights + cannot fit in available system RAM (else None). Weights only: KV/context + auto-reduce, so counting them too would refuse loads that would succeed. + None avail (unknown RAM) never refuses.""" + if avail_mib is None: + return None + need_mib = model_size_bytes / (1024 * 1024) + if need_mib <= avail_mib - headroom_mib: + return None + return ( + f"This model needs about {need_mib / 1024:.0f} GB but only about " + f"{avail_mib / 1024:.0f} GB of memory is available. On a unified-memory " + "APU the weights load into system RAM, so a larger model is stopped by " + "the OS mid-load. Use a smaller or more quantized GGUF, or free memory " + "(on WSL, raise the memory limit in .wslconfig)." + ) + # Skip the wait when the last kill is older than this; the driver has # already reclaimed the prior process's allocations. _VRAM_SETTLE_WINDOW_S: float = 15.0 @@ -1685,19 +2336,17 @@ class LlamaCppBackend: # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). _GPU_PIN_VRAM_FRACTION = 0.95 - # Per-GPU compute-graph buffer to reserve in tensor mode (MiB). This is the - # logits buffer (n_batch x vocab) + activation scratch that llama.cpp sizes - # via graph_reserve -- it is roughly EQUAL on every device (not proportional - # to the tensor split) and independent of context. Measured ~2.3 GB - # (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model; we reserve a - # conservative headroom above that. It is (a) subtracted from each GPU's free - # VRAM before computing --tensor-split, so the roomier GPU absorbs more - # weight and the smallest GPU keeps room for KV, and (b) reserved per device - # when capping context. The auto-fallback to layer split covers any - # underestimate. NOTE: scales with the model's vocab / batch size; tune if a - # large-vocab model OOMs at load. + # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF + # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived + # path) returns 0. _TENSOR_PARALLEL_BUFFER_RESERVE_MIB = 5120 + # Fixed per-device overhead on every GPU of a LAYER split (CUDA context + + # scratch), beyond the conserved slot-scaling buffer. ~0.9 GB/device measured + # (Qwen3.6-27B, b9625), independent of --parallel; reserved per extra GPU so a + # tight layer split can't advertise a context that OOMs at load. + _PIPELINE_PER_DEVICE_OVERHEAD_MIB = 1024 + # KV cache types llama.cpp accepts in tensor mode. A quantized KV cache # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) @@ -1766,11 +2415,107 @@ class LlamaCppBackend: path_dirs.append(cuda_bin_x64) return path_dirs + @staticmethod + def _llama_server_env_for_binary(binary: str) -> dict[str, str]: + """Build a subprocess env that lets llama-server resolve native libs.""" + env = child_env_without_native_path_secret() + binary_dir = str(Path(binary).parent) + + if sys.platform == "win32": + # Ordering: see _build_windows_path_dirs. #5106. + path_dirs = LlamaCppBackend._build_windows_path_dirs( + binary_dir, + sys.prefix, + os.environ.get("CUDA_PATH", ""), + ) + existing_path = env.get("PATH", "") + env["PATH"] = ";".join(path_dirs) + ";" + existing_path + + # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile + # kernel files (rocblas/library/*.dat + *.hsaco); the DLL searches + # /rocblas/library/ which doesn't exist. + _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: LD_LIBRARY_PATH for shared libs next to the binary plus + # CUDA runtime libs (libcudart, libcublas, etc.) + import platform + + lib_dirs = [] + # WSL: system HIP before the bundle's (which segfaults on /dev/dxg). + for _wsl_rocm in _wsl_system_rocm_lib_dirs(): + lib_dirs.append(_wsl_rocm) + if lib_dirs: + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + lib_dirs.append(binary_dir) + _arch = platform.machine() # x86_64, aarch64, etc. + + # Pip-installed nvidia CUDA runtime libs. The prebuilt binary links + # libcudart.so.13 / libcublas.so.13 which live here, not in + # /usr/local/cuda. + import glob as _glob + + for _nv_pattern in [ + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cu*", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cudnn", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "nvjitlink", + "lib", + ), + ]: + for _nv_dir in _glob.glob(_nv_pattern): + if os.path.isdir(_nv_dir): + lib_dirs.append(_nv_dir) + + for cuda_lib in [ + "/usr/local/cuda/lib64", + f"/usr/local/cuda/targets/{_arch}-linux/lib", + # Fallback CUDA compat paths (e.g. binary built with CUDA 12 + # where default /usr/local/cuda is CUDA 13+). + "/usr/local/cuda-12/lib64", + "/usr/local/cuda-12.8/lib64", + f"/usr/local/cuda-12/targets/{_arch}-linux/lib", + f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", + ]: + if os.path.isdir(cuda_lib): + lib_dirs.append(cuda_lib) + existing_ld = env.get("LD_LIBRARY_PATH", "") + new_ld = ":".join(lib_dirs) + env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld + + return env + @staticmethod def _select_gpus( model_size_bytes: int, gpus: list[tuple[int, int]], usable_fraction: Optional[float] = None, + total_by_idx: Optional[dict[int, int]] = None, + per_device_overhead_bytes: int = 0, ) -> tuple[Optional[list[int]], bool]: """Pick GPU(s) for a model from estimated VRAM and free memory. @@ -1778,6 +2523,11 @@ class LlamaCppBackend: ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides headroom for compute buffers, CUDA context, and other runtime overhead; callers lower it when MTP reserves VRAM for a draft model. + ``total_by_idx`` (index -> total MiB) makes the headroom an ABSOLUTE + ``(1 - fraction) * total`` per GPU instead of a fraction of free. + ``per_device_overhead_bytes`` is the fixed layer-split cost per GPU beyond + the first; a k-GPU pin must hold ``model + (k-1) * overhead`` or it can OOM + a device after -ngl -1 (no --fit fallback). Single-GPU adds none. Returns (gpu_indices, use_fit): - ([1], False) fits on 1 GPU at the headroom threshold @@ -1791,20 +2541,31 @@ class LlamaCppBackend: if usable_fraction is None: usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION - # Sort GPUs by free memory descending - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # Per-GPU usable budget: free - (1-frac)*total when total is known, else + # the legacy free*frac (also covers a total-0 two-column probe). + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - usable_fraction) * t) + return free_mib * usable_fraction + + # Rank by usable budget (free - reserve), not raw free: a more-used large + # card can have less usable room than a less-used small one. + ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True) # Try 1 GPU at the usable-VRAM threshold. - if ranked[0][1] * usable_fraction >= model_size_mib: + if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: return [ranked[0][0]], False - # Try N GPUs (accumulate free memory from most-free) - cumulative = 0 + # Try N GPUs (accumulate usable memory from most-free). Each GPU past the + # first adds a fixed per-device overhead the pool must hold. + overhead_mib = per_device_overhead_bytes / (1024 * 1024) + cumulative = 0.0 selected = [] for idx, free_mib in ranked: selected.append(idx) - cumulative += free_mib * usable_fraction - if cumulative >= model_size_mib: + cumulative += _usable(idx, free_mib) + if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib: return sorted(selected), False # Too large even for all GPUs; let --fit handle it @@ -1859,14 +2620,10 @@ class LlamaCppBackend: 5. Legacy -- fallback using embed // n_heads Server-flag knobs (mirror llama-server's CLI): - swa_full -- ``--swa-full``: force SWA layers to cache full - ``n_ctx`` (collapses path 3 to path 4 for them). - n_parallel -- ``--parallel`` slots: non-SWA layers stay constant - (cells split across slots), SWA layers scale linearly. - kv_unified -- ``--kv-unified`` (default on): no-op for memory math; - kept for API forward-compat. - ctx_checkpoints -- ``--ctx-checkpoints`` (PR #15293): N SWA snapshots - per slot, one sliding-window of state per SWA layer. + swa_full -- --swa-full: SWA layers cache full n_ctx (path 3->4). + n_parallel -- --parallel slots: non-SWA constant, SWA scale linearly. + kv_unified -- --kv-unified: memory no-op (API forward-compat). + ctx_checkpoints -- --ctx-checkpoints: N SWA snapshots per slot. Returns 0 if metadata is insufficient. """ @@ -1881,17 +2638,7 @@ class LlamaCppBackend: n_kv = self._n_kv_heads or self._n_heads or 1 # type: ignore[assignment] # Bytes per element depends on KV cache quantization - bpe = { - "f32": 4.0, - "f16": 2.0, - "bf16": 2.0, - "q8_0": 34 / 32, - "q5_1": 0.75, - "q5_0": 0.6875, - "q4_1": 0.625, - "q4_0": 0.5625, - "iq4_nl": 0.5625, - }.get(cache_type_kv or "f16", 2.0) + bpe = _kv_bytes_per_elem(cache_type_kv) slots = max(1, n_parallel) @@ -1919,15 +2666,12 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] return int(n_attn * n_ctx * n_kv * 2 * head_dim * bpe) - # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). - # Pattern filled by the resolver at parse time; if absent, falls through - # to the legacy 1/4-global heuristic below. Per-layer-type --parallel N - # accounting (verified against llama-server): - # * non-SWA layers: total cells = n_ctx split across slots -> CONSTANT. - # * SWA layers: per-slot cells = 2*sliding_window (capped at n_ctx - # and per_slot_ctx) -> grows LINEARLY in slots. - # --swa-full forces full n_ctx for SWA layers; --ctx-checkpoints N adds - # N snapshots per SWA layer per slot. + # Path 3: Sliding window (Gemma 2/3/3n/4, gpt-oss, Cohere2 ...). Pattern + # from the resolver; if absent, falls through to the legacy 1/4-global + # heuristic. --parallel N accounting (verified against llama-server): + # non-SWA cells = n_ctx split across slots (CONSTANT); SWA per-slot cells + # = 2*sliding_window (capped at n_ctx/per_slot_ctx) -> LINEAR in slots. + # --swa-full forces full n_ctx for SWA; --ctx-checkpoints N adds snapshots. if ( self._sliding_window is not None and self._sliding_window > 0 @@ -1936,8 +2680,7 @@ class LlamaCppBackend: ): swa = self._sliding_window per_slot_ctx = max(1, n_ctx // slots) - # --swa-full caches full context like non-SWA (per-slot cells = - # per_slot_ctx, collapsing to constant n_ctx total); otherwise SWA + # --swa-full caches full per_slot_ctx (constant n_ctx total); else SWA # caches 2*sliding_window per slot, clamped at per-slot ctx. swa_cells_per_slot = per_slot_ctx if swa_full else min(n_ctx, 2 * swa, per_slot_ctx) key_len_swa = self._kv_key_length_swa or key_len @@ -1990,6 +2733,162 @@ class LlamaCppBackend: head_dim = self._embedding_length // self._n_heads if self._n_heads else 128 # type: ignore[operator] return int(2 * n_kv * head_dim * n_layers_kv * n_ctx * bpe) + def _draft_backend_for(self, drafter_path: str) -> Optional["LlamaCppBackend"]: + """Lightweight backend with a drafter GGUF's metadata, to size its own KV + via _estimate_kv_cache_bytes. Cached per path; None if unreadable.""" + cache = getattr(self, "_draft_backend_cache", None) + if cache is not None and cache[0] == drafter_path: + return cache[1] + db: Optional[LlamaCppBackend] = None + try: + db = LlamaCppBackend.__new__(LlamaCppBackend) + for attr in ( + "_context_length", + "_n_layers", + "_n_kv_heads", + "_n_heads", + "_embedding_length", + "_kv_key_length", + "_kv_value_length", + "_kv_lora_rank", + "_sliding_window", + "_sliding_window_pattern", + "_ssm_inner_size", + "_full_attention_interval", + "_key_length_mla", + "_n_kv_heads_by_layer", + "_kv_key_length_swa", + "_kv_value_length_swa", + "_shared_kv_layers", + "_nextn_predict_layers", + ): + setattr(db, attr, None) + db._model_identifier = "mtp-draft" + db._read_gguf_metadata(drafter_path) + except Exception as e: # unreadable drafter -> caller falls back + logger.debug(f"Could not read drafter GGUF for MTP budget: {e}") + db = None + self._draft_backend_cache = (drafter_path, db) + return db + + def _mtp_draft_kv_bytes( + self, + n_ctx: int, + *, + drafter_path: Optional[str] = None, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + n_parallel: int = 1, + ) -> Optional[int]: + """Draft KV cache bytes at n_ctx, sized from GGUF dims (K and V types are + independent). Separate drafter (Gemma): its own KV via _estimate_kv_cache_bytes + at the heavier type. Embedded head (Qwen): nextn_predict_layers attention + layers from the main dims. None when dims are missing (flat fallback).""" + if n_ctx <= 0: + return None + bpe_k = _kv_bytes_per_elem(draft_cache_type_k) + bpe_v = _kv_bytes_per_elem(draft_cache_type_v) + if drafter_path: + db = self._draft_backend_for(drafter_path) + if db is None or not db._can_estimate_kv(): + return None + heavier = draft_cache_type_k if bpe_k >= bpe_v else draft_cache_type_v + # The drafter is served under the same --parallel slot count as the + # main model, so price its KV per slot too: a sliding-window drafter + # (Gemma) grows KV with slots and would otherwise be under-reserved. + kv = db._estimate_kv_cache_bytes(n_ctx, heavier, n_parallel = n_parallel) + return kv or None + nextn = self._nextn_predict_layers or 0 + n_kv = self._n_kv_heads or self._n_heads + k_len = self._kv_key_length + v_len = self._kv_value_length + if not (nextn and n_kv and k_len and v_len): + return None + # The embedded MTP head is one draft layer, so a quantized draft KV can't + # amortize its overhead and fits *less* context than f16 (llama.cpp#24102). + # Floor it at f16: a quantized override is priced as f16, f32 keeps its 4 + # bytes. The separate-drafter branch is multi-layer, so it keeps its type. + f16_bpe = _kv_bytes_per_elem("f16") + bpe_k = max(bpe_k, f16_bpe) + bpe_v = max(bpe_v, f16_bpe) + return int(nextn * n_kv * (k_len * bpe_k + v_len * bpe_v) * n_ctx) + + def _estimate_mtp_overhead_bytes( + self, + n_ctx: int, + *, + spec_draft_n_max: int = 0, + draft_cache_type_k: Optional[str] = None, + draft_cache_type_v: Optional[str] = None, + drafter_path: Optional[str] = None, + draft_weights_bytes: int = 0, + n_parallel: int = 1, + ) -> Optional[int]: + """MTP draft reserve at ``n_ctx`` = draft KV (grows with ctx) + separate- + drafter weights + (MLA only) a duplicated target KV context. The verify + buffer rides in the ctx-fit headroom (no tuned constant). None when the + draft KV can't be sized (caller keeps the flat fallback). + ``draft_weights_bytes`` is the drafter file size (0 for embedded).""" + draft_kv = self._mtp_draft_kv_bytes( + n_ctx, + drafter_path = drafter_path, + draft_cache_type_k = draft_cache_type_k, + draft_cache_type_v = draft_cache_type_v, + n_parallel = n_parallel, + ) + weights = max(0, draft_weights_bytes) + # MLA models (GLM-5.x, DeepSeek, Kimi-K2) keep a *second* full copy of the + # target model's KV context for MTP draft verification -- llama.cpp's + # `ctx_tgt=yes` -- allocated at f16 regardless of the main cache type. It is + # ~the main KV again and dwarfs the embedded draft head (GLM-5.2 @ 1M ctx: + # a ~2 GiB head next to a ~89 GiB target copy), so omitting it lets auto-fit + # pick a context that fits on paper but OOMs cublasCreate at the first + # decode. Non-MLA MTP (Qwen/Gemma) keeps no such copy, so this is gated + # strictly on MLA (kv_lora_rank present) and leaves those models unchanged. + target_ctx_copy = 0 + if self._kv_lora_rank is not None: + target_ctx_copy = self._estimate_kv_cache_bytes(n_ctx, "f16", n_parallel = n_parallel) + if draft_kv is None: + # KV unsized (exotic/remote drafter): still reserve known weights + any + # MLA target copy so a large config can't launch over budget (the small + # unsized draft KV rides in the cushion). Nothing known -> None, so the + # caller keeps the flat fallback. + total = weights + target_ctx_copy + return total if total > 0 else None + return draft_kv + weights + target_ctx_copy + + _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it + _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + + def _estimate_compute_buffer_bytes( + self, + *, + n_ubatch: Optional[int] = None, + n_parallel: int = 1, + per_device_tensor: bool = False, + ) -> int: + """Per-device compute-graph buffer (bytes) from GGUF dims: a vocab-width + output buffer + activation scratch. Context-independent; scales with + ``--parallel`` (serving slots). Tensor mode materializes it on every device. + A slight upper bound over measured allocations; 0 when dims are missing.""" + n_vocab = self._vocab_size or 0 + n_embd = self._embedding_length or 0 + if n_vocab <= 0 or n_embd <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + par = max(1, int(n_parallel)) + out_buffer = n_vocab * ub * 4 # f32 output/logits buffer + act_scratch = 4 * n_embd * ub * 4 # a few resident hidden-width buffers + if per_device_tensor: + # Output + comm/staging materialized on every device, every slot. + compute = 2 * act_scratch + out_buffer * par + else: + # Each extra concurrent slot adds one output buffer (chat decode sizes + # ~one logit row per slot; would under-count embeddings/--logits-all, + # not run here). Matches measured {1:36,2:492,4:1388,8:3220} MiB. + compute = act_scratch + out_buffer * max(0, par - 1) + return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _fit_context_to_vram( self, requested_ctx: int, @@ -2004,13 +2903,15 @@ class LlamaCppBackend: ctx_checkpoints: int = 0, kv_on_gpu: bool = True, mtp_engaged: bool = False, + mtp_overhead_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, + total_mib: Optional[int] = None, ) -> int: """Return the largest context length that fits in GPU VRAM. - Uses 90% of available VRAM as the ctx-fit budget -- tighter than - ``_GPU_PIN_VRAM_FRACTION`` on purpose (over-promising context OOMs at - runtime). If the weights alone don't fit, returns ``requested_ctx``. + Budget caps occupancy at ``_CTX_FIT_VRAM_FRACTION`` of the card: an + absolute ``free - (1 - frac) * total`` when ``total_mib`` is given, else + ``free * frac``. Weights alone over budget returns ``requested_ctx``. ``kv_on_gpu`` mirrors ``--kv-offload`` (default on); when False the KV cache lives in CPU RAM and the requested context is honored verbatim. @@ -2038,20 +2939,28 @@ class LlamaCppBackend: ctx_checkpoints = ctx_checkpoints, ) - # MTP engaged: carve the drafter's reserve out of the fit budget. Callers - # can override outright (tensor-parallel mode passes a fatter margin), so - # only compute a default when none was supplied. + # byte-accurate mtp_overhead_fn supersedes the flat fraction (the fallback + # when dims can't size the draft KV); callers may override budget_frac. if budget_frac is None: - budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if mtp_engaged else 0.0) - budget_bytes = available_mib * 1024 * 1024 * budget_frac + flat_mtp = mtp_engaged and mtp_overhead_fn is None + budget_frac = _CTX_FIT_VRAM_FRACTION - (_MTP_VRAM_RESERVE_FRAC if flat_mtp else 0.0) + # Absolute reserve off total when known, else fraction-of-free; clamp >=0. + if total_mib is not None and total_mib > 0: + budget_mib = max(0.0, available_mib - (1.0 - budget_frac) * total_mib) + else: + budget_mib = available_mib * budget_frac + budget_bytes = budget_mib * 1024 * 1024 model_footprint = model_size_bytes + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: return requested_ctx - # Weights alone exceed budget -- reducing ctx can't help; --fit handles it. + # Weights + compute buffer alone exceed budget -- reducing ctx can't help. if model_footprint >= budget_bytes: logger.debug( "Model footprint exceeds GPU budget before KV cache", @@ -2061,7 +2970,7 @@ class LlamaCppBackend: ) return requested_ctx - # Binary search for max context that fits + # Binary search for max context that fits (KV + MTP draft reserve at that ctx) remaining = budget_bytes - model_footprint effective_min = min(min_ctx, requested_ctx) lo, hi = effective_min, requested_ctx @@ -2069,7 +2978,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv <= remaining: + if kv + _mtp_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -2101,7 +3010,11 @@ class LlamaCppBackend: files = list_repo_files(hf_repo, token = hf_token) gguf_files = [ - f for f in files if f.endswith(".gguf") and not _is_companion_gguf_path(f) + f + for f in files + if f.lower().endswith(".gguf") + and not _is_companion_gguf_path(f) + and not _is_big_endian_gguf_path(f) ] if not gguf_files: return None @@ -2236,6 +3149,7 @@ class LlamaCppBackend: self._supports_reasoning = False self._reasoning_always_on = False self._reasoning_style = "enable_thinking" + self._reasoning_effort_levels = [] self._reasoning_default = True self._supports_preserve_thinking = False self._supports_tools = False @@ -2244,6 +3158,8 @@ class LlamaCppBackend: self._n_kv_heads_by_layer = None self._n_heads = None self._embedding_length = None + self._feed_forward_length = None + self._vocab_size = None self._kv_key_length = None self._kv_value_length = None self._sliding_window = None @@ -2265,6 +3181,8 @@ class LlamaCppBackend: WANTED = { "general.architecture", "tokenizer.chat_template", + # Vocab size = tokens array length (no vocab_size key in many GGUFs). + "tokenizer.ggml.tokens", # Block-diffusion marker (DiffusionGemma); routes to the diffusion runner. "diffusion.canvas_length", # Source-repo hints for the SWA resolver's HF fallback. @@ -2328,6 +3246,7 @@ class LlamaCppBackend: f"{arch}.attention.head_count_kv": "n_kv_heads", f"{arch}.attention.head_count": "n_heads", f"{arch}.embedding_length": "embedding_length", + f"{arch}.feed_forward_length": "feed_forward_length", f"{arch}.attention.key_length": "kv_key_length", f"{arch}.attention.value_length": "kv_value_length", f"{arch}.attention.sliding_window": "sliding_window", @@ -2361,6 +3280,9 @@ class LlamaCppBackend: elif vtype == 9: # ARRAY atype = struct.unpack(" %s from local HF cache", hf_variant, @@ -2832,10 +3706,11 @@ class LlamaCppBackend: _m = _SHARD_RE.match(gguf_filename) _prefix = _m.group(1) if _m else None if _prefix: + prefix_lower = _prefix.lower() gguf_extra_shards = sorted( f for f in all_gguf_files - if f.startswith(_prefix) + if f.lower().startswith(prefix_lower) and f != gguf_filename and not _is_companion_gguf_path(f) ) @@ -2859,19 +3734,23 @@ class LlamaCppBackend: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") dl_start = time.monotonic() - local_path = hf_hub_download( - repo_id = hf_repo, - filename = gguf_filename, - token = hf_token, + # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. + local_path = hf_hub_download_with_xet_fallback( + hf_repo, + gguf_filename, + hf_token, + cancel_event = self._cancel_event, + on_status = lambda m: logger.info(m), ) for shard in gguf_extra_shards: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download( - repo_id = hf_repo, - filename = shard, - token = hf_token, + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = self._cancel_event, ) except RuntimeError as e: if "Cancelled" in str(e): @@ -2920,7 +3799,7 @@ class LlamaCppBackend: try: from utils.models.model_config import _iter_hf_cache_snapshots for snap in _iter_hf_cache_snapshots(hf_repo): - rel_files = [p.relative_to(snap).as_posix() for p in snap.rglob("*.gguf")] + rel_files = _gguf_snapshot_files(snap) target = pick(rel_files) if target is not None: logger.info("Resolved %s %s from local HF cache", label, target) @@ -2932,12 +3811,13 @@ class LlamaCppBackend: return None try: - from huggingface_hub import hf_hub_download logger.info(f"Downloading {label}: {hf_repo}/{target}") - return hf_hub_download( - repo_id = hf_repo, - filename = target, - token = hf_token, + # Same policy; companions are best-effort (caller below swallows failures to None). + return hf_hub_download_with_xet_fallback( + hf_repo, + target, + hf_token, + cancel_event = self._cancel_event, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") @@ -3084,7 +3964,10 @@ class LlamaCppBackend: @staticmethod def _classify_llama_start_failure( - output: str, gguf_path: Optional[str], model_identifier: Optional[str] + output: str, + gguf_path: Optional[str], + model_identifier: Optional[str], + returncode: Optional[int] = None, ) -> str: """Explain *why* llama-server failed to start, from its output. @@ -3156,6 +4039,24 @@ class LlamaCppBackend: "Ollama instead." ) + # SIGKILL with no diagnostic output is the OOM killer (e.g. a model too + # large for the WSL VM's RAM cap); name it actionably. + if returncode == -9: + return ( + "llama-server was stopped by the operating system (signal 9), " + "most likely out of memory. Try a smaller or more quantized " + "GGUF, lower the context length, or free memory (on WSL, raise " + "the memory limit in .wslconfig)." + ) + # SIGTERM is also how an unload/cancel or a supervisor stops the server, + # so report it neutrally rather than blaming memory. + if returncode == -15: + return ( + "llama-server was terminated (signal 15) before it became " + "healthy. If you cancelled or unloaded the model this is " + "expected; otherwise check the llama-server log for the cause." + ) + # Fallback: genuinely unknown failure (OOM, missing binary ...). return ( "llama-server failed to start. " @@ -3170,7 +4071,11 @@ class LlamaCppBackend: cache_type_kv: Optional[str] = None, n_parallel: int = 1, mtp_engaged: bool = False, + mtp_overhead_fn: Optional[Callable[[int], int]] = None, + mtp_flat_reserve_bytes: int = 0, max_target_ctx: Optional[int] = None, + total_by_idx: Optional[dict[int, int]] = None, + n_ubatch: Optional[int] = None, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -3183,22 +4088,41 @@ class LlamaCppBackend: Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_TENSOR_PARALLEL_BUFFER_RESERVE_MIB``). + one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. - ``tensor_split`` is None (llama.cpp's even default, safe for every arch incl. Gemma 3n which GGML_ASSERTs on a weighted split) when an even - share fits the smallest GPU; otherwise it is weighted by - ``(free - buffer)`` so the roomier GPU absorbs more weight and the - smallest GPU keeps room for KV. + share fits the smallest GPU; otherwise it is weighted by usable budget + so the roomier GPU absorbs more weight and the smallest keeps room for KV. + ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes + the compute buffer. """ - # Drop GPUs that can't hold the per-device compute-graph buffer; they'd - # OOM in tensor mode. load_model already filters before calling, so this - # is defense-in-depth that also keeps the pure function self-contained - # (and unit-testable without a GPU). - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - usable_gpus = [g for g in gpus if g[1] >= reserve_mib] + + # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a + # two-column probe) the legacy free*frac. Mirrors _select_gpus and + # _gpu_usable so the 5% cushion is kept on every path, not dropped here. + def _usable(idx: int, free_mib: int) -> float: + t = total_by_idx.get(idx, 0) if total_by_idx else 0 + if t > 0: + return max(0.0, free_mib - (1.0 - _CTX_FIT_VRAM_FRACTION) * t) + return max(0.0, free_mib * _CTX_FIT_VRAM_FRACTION) + + # Drop GPUs whose usable budget can't hold the per-device compute-graph + # buffer; they'd OOM in tensor mode. Admitting on raw free would let a + # partly-used big card in with no budget left. Defense-in-depth (load_model + # gates too). Derived per-device reserve; flat fallback. + _reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = n_parallel, per_device_tensor = True + ) + reserve_mib = ( + _reserve_bytes // (1024 * 1024) + if _reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + usable_gpus = [g for g in gpus if _usable(g[0], g[1]) >= reserve_mib] gpu_indices = sorted(idx for idx, _ in usable_gpus) if len(gpu_indices) < 2: # Tensor parallelism is meaningless on <2 GPUs (the caller drops the @@ -3210,21 +4134,50 @@ class LlamaCppBackend: None, ) free_by_idx = {idx: free for idx, free in usable_gpus} - pool_mib = sum(free_by_idx.values()) - kv_budget_b = (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - if mtp_engaged: - # MTP keeps a draft model + its own KV cache on GPU. - kv_budget_b -= 2 * 1024**3 + usable_by_idx = {idx: _usable(idx, free_by_idx[idx]) for idx in gpu_indices} + pool_mib = sum(usable_by_idx.values()) + # MTP reserve: byte-accurate per-ctx inside _fit_ctx (mtp_overhead_fn) plus + # a flat cushion that the byte fn can't size -- 2 GiB when dims are wholly + # unavailable (no fn), or mtp_flat_reserve_bytes when the fn is weights-only + # because the draft KV couldn't be sized (_mtp_kv_unsized). Without this the + # binary search spends the unsized-KV cushion on main context and OOMs. + flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) + if mtp_engaged and mtp_overhead_fn is None: + flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + kv_budget_b = ( + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + ) + + def _mtp_at(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 def _fit_ctx(ctx: int) -> int: - # Largest context whose KV fits the pooled budget. Floors small, but - # never raises an explicit ctx above what was asked. + # Largest context whose KV (+ MTP draft reserve) fits the pooled + # budget. Floors small, but never raises an explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: # Weights + buffers exceed the pool -> floor; the load then # falls back to layer split. return ctx_floor + if mtp_overhead_fn is not None: + # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + def _consumer(c: int) -> int: + return self._estimate_kv_cache_bytes( + c, cache_type_kv, n_parallel = n_parallel + ) + _mtp_at(c) + + if _consumer(ctx) <= kv_budget_b: + return ctx + lo, hi, best = ctx_floor, ctx, ctx_floor + while lo <= hi: + mid = (lo + hi) // 2 + if _consumer(mid) <= kv_budget_b: + best = mid + lo = mid + 1 + else: + hi = mid - 1 + return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) if kv_at <= kv_budget_b: return ctx @@ -3239,16 +4192,19 @@ class LlamaCppBackend: max_available_ctx = _fit_ctx(max_ctx_target) effective_ctx = min(_fit_ctx(target_ctx), max_available_ctx) - min_free_mib = min(free_by_idx.values()) + min_usable_mib = min(usable_by_idx.values()) kv_bytes = ( self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = n_parallel) if (self._can_estimate_kv() and effective_ctx > 0) else 0 ) - even_share_mib = (model_size + kv_bytes) / len(gpu_indices) / (1024 * 1024) + # The MTP reserve also has to fit the even split (mirror the pooled budget): + # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. + mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes + even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) tensor_split: Optional[list[int]] = None - if even_share_mib > (min_free_mib - reserve_mib): - adj = [max(0, int(free_by_idx[i] - reserve_mib)) for i in gpu_indices] + if even_share_mib > (min_usable_mib - reserve_mib): + adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -3279,6 +4235,71 @@ class LlamaCppBackend: and ("unknown" in text or "unsupported" in text or "not supported" in text) ) + @staticmethod + def _output_has_nonprojector_diagnostic(output: str) -> bool: + """True when the output already names a concrete non-projector cause (out + of memory, an unsupported architecture, a tensor-parallel limit). A hard + crash carrying such a marker must surface that error, not be silently + retried text-only as if the vision projector were at fault; a bare crash + with no marker still gets the text-only retry. + """ + text = (output or "").lower() + return any( + m in text + for m in ( + "out of memory", + "failed to allocate", + "unknown model architecture", + "split_mode_tensor not implemented", + ) + ) + + @staticmethod + def _is_signal_crash(returncode: Optional[int]) -> bool: + """True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a + Windows 0xC0000000+ status), not SIGKILL/SIGTERM/SIGINT (OOM killer / + unload) nor a clean exit or still-running (None) process. + """ + if returncode is None: + return False + if returncode >= 0xC0000000: # Windows access violation / illegal instruction + return True + return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV + + @staticmethod + def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: + """Return cmd with flash attention forced off, or None when its effective + (last-wins) value is already off/absent so there is nothing to retry. FA + kernels hard-crash at startup on some ROCm builds; disabling FA keeps + vision and MTP, the least destructive rung. A bare --flash-attn/-fa reads + as on, so it counts toward the effective value and is neutralised too; + every form is flipped in place (length preserved for downstream slices).""" + out = list(cmd) + + def explicit(i): + nxt = out[i + 1] if i + 1 < len(out) else None + return nxt if nxt in ("on", "auto", "off") else None + + effective = None + for i, tok in enumerate(out): + if tok.startswith(("--flash-attn=", "-fa=")): + effective = tok.partition("=")[2] + elif tok in ("--flash-attn", "-fa"): + effective = explicit(i) or "on" + if effective not in ("on", "auto"): + return None + for i, tok in enumerate(out): + if tok.startswith(("--flash-attn=", "-fa=")): + flag, _, value = tok.partition("=") + if value in ("on", "auto"): + out[i] = f"{flag}=off" + elif tok in ("--flash-attn", "-fa"): + if explicit(i) in ("on", "auto"): + out[i + 1] = "off" + elif explicit(i) is None: # bare flag (reads as on) -> explicit off + out[i] = f"{tok}=off" + return out + @staticmethod def _strip_mmproj_args(cmd: list[str]) -> list[str]: """Return cmd without the '--mmproj ' pair (text-only retry). @@ -3346,7 +4367,11 @@ class LlamaCppBackend: text = True, env = env, **_windows_hidden_subprocess_kwargs(), + **_child_popen_kwargs(), ) + # Cross-session backstop: record the PID so a later startup can reap this + # server if parent-death cleanup did not run (macOS / best-effort failure). + self._record_server_pid(self._process.pid) # Start background thread to drain stdout and prevent pipe deadlock self._stdout_thread = threading.Thread( @@ -3389,6 +4414,28 @@ class LlamaCppBackend: Returns True if the server started and the health check passed. """ + # Raw load inputs so the runtime MTP-crash reload can replay this model + # without MTP. Committed to _last_load_kwargs only on a healthy load. + _pending_load_kwargs = { + "gguf_path": gguf_path, + "mmproj_path": mmproj_path, + "mtp_draft_path": mtp_draft_path, + "hf_repo": hf_repo, + "hf_variant": hf_variant, + "hf_token": hf_token, + "model_identifier": model_identifier, + "is_vision": is_vision, + "n_ctx": n_ctx, + "chat_template_override": chat_template_override, + "cache_type_kv": cache_type_kv, + "speculative_type": speculative_type, + "spec_draft_n_max": spec_draft_n_max, + "tensor_parallel": tensor_parallel, + "n_threads": n_threads, + "n_gpu_layers": n_gpu_layers, + "n_parallel": n_parallel, + "extra_args": list(extra_args) if extra_args is not None else None, + } # Serialise the whole load so concurrent /load calls never leave two # llama-server processes alive (#5401 / #5161). Doesn't block /unload. with self._serial_load_lock: @@ -3548,11 +4595,11 @@ class LlamaCppBackend: "(access-denied; antivirus or an in-flight install). " "Retry the load once it is released." ) - raise RuntimeError( - "llama-server binary not found. " - "Run setup.sh to build it, install llama.cpp, " - "or set LLAMA_SERVER_PATH environment variable." - ) + # Reached only after the diffusion early-return above, so this is a + # genuine llama-server-backed GGUF with no runtime. Raise the typed + # error so /load returns the actionable 400 (not a generic 500), the + # same message remote validation already shows. + raise LlamaServerNotFoundError(LLAMA_SERVER_NOT_FOUND_DETAIL) # Outside ``self._lock`` so /unload, /cancel, /status aren't # blocked. ``unload_model`` also records the kill, so the @@ -3575,27 +4622,58 @@ class LlamaCppBackend: ctx_override = parse_ctx_override(extra_args) requested_ctx = resolve_requested_ctx(extra_args, n_ctx) cache_override = parse_cache_override(extra_args) - cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv) - # A user --split-mode in extras last-wins-overrides the - # toggle, so reconcile it back into tensor_parallel state. + # Budget the heavier of asymmetric --cache-type-k/-v extras (they + # win per axis at launch, appended last); resolve_cache_type_kv only + # returns the last-wins type, which under-reserves the heavier axis. + # The user's extras still set the real (possibly asymmetric) child + # cache, so this only affects the reserve, not the emitted command. + _extras_cache = _extra_args_main_cache_type_for_budget(extra_args) + cache_type_kv = _extras_cache if _extras_cache is not None else cache_type_kv + _cache_type_from_env = False + if cache_type_kv is None: + # Param/extras set nothing, so the child inherits + # LLAMA_ARG_CACHE_TYPE_K/_V. Adopt a heavier env type (f32) for + # the reserve only; the launch does NOT re-emit it (that would + # rewrite an asymmetric K=f32,V=f16 env into symmetric flags), + # so _cache_type_from_env keeps it out of the emitted flags. + cache_type_kv = _env_main_cache_type_for_budget() + _cache_type_from_env = cache_type_kv is not None + # A user --split-mode in extras last-wins-overrides the toggle, and + # an inherited tensor LLAMA_ARG_SPLIT_MODE flips it on (the child + # would run tensor unbudgeted otherwise). The duplicate-load matchers + # use the same helper so a healthy env-driven tensor server matches. split_mode_override = parse_split_mode_override(extra_args) - tensor_parallel = resolve_tensor_parallel(extra_args, tensor_parallel) + tensor_parallel = _effective_tensor_parallel(extra_args, tensor_parallel) # Tensor mode aborts on a quantized KV cache, so drop it for the # tensor attempt (and strip any inherited/explicit --cache-type - # that would re-impose it when appended last). The layer-split - # fallback re-runs with tensor_parallel False and keeps the type. - if ( - tensor_parallel - and cache_type_kv - and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES - ): + # that would re-impose it when appended last). Layer split does + # support it, so remember the dropped type and the original extras + # to restore (verbatim, incl. an asymmetric K/V) if we later fall + # back to layer split below. + _tensor_dropped_cache_type_kv: Optional[str] = None + _tensor_dropped_extra_args: Optional[list] = None + # Tensor mode rejects any quantized axis. cache_type_kv is the + # heavier-by-bytes budget type, which can mask a quantized axis (an + # f16 budget hides a paired q4_0), so also test each explicit + # --cache-type-k/-v extra, not just the budget type. + _ck_extra, _cv_extra = parse_cache_override_per_axis(extra_args) + _cache_non_tensor_safe = any( + c and c.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES + for c in (cache_type_kv, _ck_extra, _cv_extra) + ) + if tensor_parallel and _cache_non_tensor_safe: logger.info( "Tensor parallelism requires a non-quantized KV cache; " "ignoring cache type %s for the tensor attempt.", cache_type_kv, ) + _tensor_dropped_cache_type_kv = cache_type_kv cache_type_kv = None if extra_args: + # Keep the originals so a layer downgrade restores the real + # (possibly asymmetric) --cache-type-k/-v the layer path + # supports, not just the scalar heavier type. + _tensor_dropped_extra_args = list(extra_args) extra_args = strip_shadowing_flags( extra_args, strip_context = False, @@ -3604,10 +4682,24 @@ class LlamaCppBackend: strip_template = False, strip_split_mode = False, ) + # The launch keeps an inherited tensor-safe env cache type (the + # env cleanup only pops quantized ones), so re-adopt a heavier + # env type (f32) for the budget here too -- mirrors the initial + # adoption, which was skipped because the param/extras set the + # (now-dropped) quantized type. Else the child allocates f32 KV + # against an f16 budget. + _env_tensor_cache = _env_main_cache_type_for_budget() + if _env_tensor_cache is not None: + cache_type_kv = _env_tensor_cache + _cache_type_from_env = True if ctx_override is not None and ctx_override > 0: logger.info(f"User --ctx-size {ctx_override} honored; skipping auto-reduce") if cache_override is not None: - logger.info(f"User --cache-type-k/-v {cache_override} honored for KV estimate") + _ck, _cv = parse_cache_override_per_axis(extra_args) + logger.info( + f"User --cache-type-k/-v (k={_ck}, v={_cv}) honored; " + "KV estimate budgets the heavier axis" + ) if split_mode_override is not None: logger.info( f"User --split-mode {split_mode_override} honored; " @@ -3632,6 +4724,7 @@ class LlamaCppBackend: "Vision-capable GGUF loaded without a usable mmproj; " "image input will be disabled for this session" ) + model_size = None # set in the fit try; used by the APU RAM guard try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -3639,7 +4732,27 @@ class LlamaCppBackend: self._mmproj_vram_bytes(launch_mmproj_path) if effective_is_vision else 0 ) model_size = gguf_size + mmproj_size - gpus = self._get_gpu_free_memory() + # 2-tuple gpus for existing logic + a total map for the absolute + # per-GPU headroom (correct when the GPU is already partly used). + _gpu_mem = self._get_gpu_memory() + gpus = [(idx, free) for idx, free, _t in _gpu_mem] + total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + + def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): + # Per-GPU usable budget for ranking: free - (1-frac)*total. + # Callers pass the ACTIVE fraction so the ranking matches the + # budget the fit then tests (else mixed totals mis-order). + idx, free = g + t = total_by_idx.get(idx, 0) + if t > 0: + return free - (1.0 - frac) * t + return free * frac + + def _pool_budget_mib(subset, frac): + # Sum each GPU's own usable budget. Pooling free and total + # separately would let an unknown-total GPU (MIG/vGPU/N/A) + # add full free with no cushion among known-total GPUs. + return sum(max(0.0, _gpu_usable(g, frac)) for g in subset) # Resolve effective context: 0 means let llama-server use # the model's native length. Only expand to a known native @@ -3655,12 +4768,10 @@ class LlamaCppBackend: # GPU/VRAM-fit logic below may shrink it on limited HW. max_available_ctx = self._context_length or effective_ctx - # Will MTP engage on this load? If so, auto-fit reserves - # extra VRAM for the draft model. Mirrors - # _build_speculative_flags' resolver: forced mtp / mtp+ngram - # always engage; auto only on an MTP model >= 3B; ngram / - # ngram-simple / off never engage MTP. A separate drafter - # (Gemma) counts as an MTP model just like a baked-in head. + # Will MTP engage? If so, auto-fit reserves draft-model VRAM. + # Mirrors _build_speculative_flags: forced mtp/mtp+ngram always + # engage; auto only on an MTP model >= 3B; ngram/off never. A + # separate drafter (Gemma) counts as an MTP model. _mtp_canonical = _canonicalize_spec_mode(speculative_type) _mtp_effective = _mtp_canonical or "auto" _mtp_size_for_fit = _extract_model_size_b(model_identifier) @@ -3671,49 +4782,253 @@ class LlamaCppBackend: and _mtp_size_for_fit < _MTP_MIN_SIZE_B and not bool(mtp_draft_path) ) + # LLAMA_ARG_SPEC_TYPE only reaches the child when neither extras + # nor Studio emit a spec flag (mode "off", no user --spec-type), + # since _build_speculative_flags emits one for every other mode. + # Consult the env for the reserve only then, else a stale MTP env + # would over-reserve. + _spec_env: Mapping[str, str] = ( + os.environ + if (not _extra_args_set_spec_type(extra_args) and _mtp_canonical == "off") + else {} + ) + # Extras can run MTP even when Studio suppresses its own emission. + _user_mtp_via_extras = _extra_args_requests_mtp(extra_args, env = _spec_env) + # A non-MTP model-based draft mode (draft-simple/draft-eagle3) in + # extras also loads a separate draft model that needs reserving; + # engage only when extras actually name a drafter for it. + _user_draft_via_extras = _extra_args_requests_separate_draft( + extra_args, env = _spec_env + ) and bool(_extra_args_mtp_draft_path(extra_args)) + # Mirror _build_speculative_flags: reserve only for MTP the launch + # resolver will actually emit (needs a head/drafter and a binary + # that supports --spec-type mtp). + _mtp_model_for_fit = bool( + self._nextn_predict_layers + or _is_mtp_model_name(model_identifier, model_path) + or bool(mtp_draft_path) + ) + _mtp_binary_ok = True + _mtp_probe_raised = False + if not _user_mtp_via_extras: + try: + _mtp_binary_ok = bool( + (self.probe_server_capabilities(binary) or {}).get("mtp_token") + ) + except Exception: + _mtp_binary_ok = False + _mtp_probe_raised = True _mtp_will_engage = bool( - not _extra_args_set_spec_type(extra_args) - and ( - _mtp_effective in ("mtp", "mtp+ngram") - or ( - _mtp_effective == "auto" - and ( - bool(self._nextn_predict_layers) - or _is_mtp_model_name(model_identifier, model_path) - or bool(mtp_draft_path) - ) - and not _mtp_sub_3b_for_fit + _user_mtp_via_extras + or _user_draft_via_extras + or ( + not _extra_args_set_spec_type(extra_args) + and _mtp_model_for_fit + and ( + _mtp_effective in ("mtp", "mtp+ngram") + or (_mtp_effective == "auto" and not _mtp_sub_3b_for_fit) + ) + and ( + _mtp_binary_ok + # Reserve on a raised (uncached) probe too: it re-probes in + # _build_speculative_flags and may still engage MTP (embedded + # head or separate drafter -- _mtp_model_for_fit covers both). + or _mtp_probe_raised ) ) ) - # Auto-cap context to fit GPU VRAM and select GPUs. Two - # policies by whether the user set n_ctx: - # Explicit n_ctx: honor it. Try the full context with - # _select_gpus (as many GPUs as needed); cap only if it - # doesn't fit on any combination. - # Auto n_ctx=0 (native): prefer fewer GPUs with reduced - # context, since multi-GPU is slower. + # Effective draft depth: extras win (last-wins at launch), else + # the field, else the platform default (2 GPU / 3 CPU). + _extra_n_max = _extra_args_spec_draft_n_max(extra_args) + _mtp_eff_n_max = _extra_n_max if _extra_n_max is not None else spec_draft_n_max + if _mtp_eff_n_max is None: + _mtp_eff_n_max = 2 if gpus else 3 + # Separate-drafter weights live on GPU (an embedded head is + # already in model_size). Size the drafter the launch loads, by + # precedence: extras --model-draft (last-wins), else Studio's + # emitted mtp_draft_path, else the env drafter. Sizing the wrong + # one would under-reserve and OOM. + _cli_draft_for_budget = _extra_args_mtp_draft_path(extra_args, env = {}) + _studio_draft_for_budget = ( + mtp_draft_path + if ( + _mtp_will_engage + and mtp_draft_path + and not _extra_args_set_spec_type(extra_args) + ) + else None + ) + _env_draft_for_budget = _extra_args_mtp_draft_path([], env = os.environ) + _mtp_draft_for_budget = ( + _cli_draft_for_budget or _studio_draft_for_budget or _env_draft_for_budget + ) + # Drafter offloaded to CPU keeps its weights+KV off the GPU, so + # drop it from the budget (an embedded head stays in the model). + # Consult the env too: the child honors LLAMA_ARG_N_GPU_LAYERS_DRAFT. + _draft_on_cpu = _extra_args_draft_offloaded_to_cpu(extra_args, env = os.environ) + if _draft_on_cpu: + _mtp_draft_for_budget = None + _mtp_draft_weights = 0 + if _mtp_draft_for_budget: + try: + _mtp_draft_weights = self._get_gguf_size_bytes(_mtp_draft_for_budget) + except Exception: + _mtp_draft_weights = 0 + # Draft K/V types (f16 by default; independent extras overrides). + _mtp_draft_ck, _mtp_draft_cv = _extra_args_draft_cache_types(extra_args) + + # Byte-accurate reserve when dims allow, else None -> flat fallback. + mtp_overhead_fn: Optional[Callable[[int], int]] = None + # True when the byte reserve is the drafter weights ONLY because + # its KV couldn't be sized; the flat fraction must then stay on + # as the cushion for that unsized draft KV (it is not covered by + # the weights-only mtp_overhead_fn). + _mtp_kv_unsized = False + if _mtp_will_engage: + _probe_ctx = self._context_length or ( + effective_ctx if effective_ctx > 0 else 4096 + ) + _draft_kv_probe = self._mtp_draft_kv_bytes( + _probe_ctx, + drafter_path = _mtp_draft_for_budget, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + n_parallel = n_parallel, + ) + if ( + self._estimate_mtp_overhead_bytes( + _probe_ctx, + spec_draft_n_max = _mtp_eff_n_max, + draft_cache_type_k = _mtp_draft_ck, + draft_cache_type_v = _mtp_draft_cv, + drafter_path = _mtp_draft_for_budget, + draft_weights_bytes = _mtp_draft_weights, + n_parallel = n_parallel, + ) + is not None + ): + # Reserve is weights-only when the draft KV is unsizable. + _mtp_kv_unsized = _draft_kv_probe is None + + # Closure binding this load's draft params; ctx varies. + def mtp_overhead_fn( + ctx: int, + _n: int = _mtp_eff_n_max, + _ck: Optional[str] = _mtp_draft_ck, + _cv: Optional[str] = _mtp_draft_cv, + _dp: Optional[str] = _mtp_draft_for_budget, + _w: int = _mtp_draft_weights, + _np: int = n_parallel, + ) -> int: + v = self._estimate_mtp_overhead_bytes( + ctx, + spec_draft_n_max = _n, + draft_cache_type_k = _ck, + draft_cache_type_v = _cv, + drafter_path = _dp, + draft_weights_bytes = _w, + n_parallel = _np, + ) + return v if v is not None else 0 + + def _mtp_bytes(ctx: int) -> int: + return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + + # Effective micro-batch (a user --ubatch override scales the + # compute buffer); None -> the 512 default in the estimate. + _effective_ubatch = _extra_args_n_ubatch(extra_args) + + # Layer-split compute buffer (one lump; tensor mode reserves it + # per device in _plan_tensor_parallel). Context-independent, so + # fold it into the model footprint for the branches below. Falls + # back to the flat reserve when dims are missing (returns 0), a + # safe upper bound since the tensor buffer >= the layer one. + _compute_buffer_pipeline = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = False, + ) + if _compute_buffer_pipeline <= 0: + _compute_buffer_pipeline = ( + self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + ) + model_size_fit = model_size + _compute_buffer_pipeline + + # Layer split adds a fixed per-device overhead on every GPU. The + # folded buffer covers one device; reserve the extra devices' + # share so a k-GPU split can't pin a context that OOMs a device + # (k=1 adds nothing). + _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + + # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: + # honor it, cap only if it fits no combination. Auto (native): + # prefer fewer GPUs with reduced context (multi-GPU is slower). gpu_indices, use_fit = None, True # Per-GPU weight proportions for tensor mode (None = even). tp_tensor_split: Optional[list[int]] = None explicit_ctx = requested_ctx > 0 - # MTP draft model lives outside the main estimates; carve - # its reserve out of every fit budget and pin threshold so - # a load can't pin into the drafter's headroom. - _mtp_reserve = _MTP_VRAM_RESERVE_FRAC if _mtp_will_engage else 0.0 - _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _mtp_reserve + # Flat MTP reserve fraction: used only as the fallback when the + # byte-accurate mtp_overhead_fn can't size the draft KV (dims + # unavailable, or _mtp_kv_unsized = weights-only). A separate + # drafter on CPU uses no GPU (no reserve); an embedded head is on + # GPU regardless of draft-offload flags (keep its reserve). + _flat_mtp_engages = _mtp_will_engage and ( + mtp_overhead_fn is None or _mtp_kv_unsized + ) + _draft_cpu_no_embedded = _draft_on_cpu and not self._nextn_predict_layers + # MTP reserves GPU VRAM unless its only drafter is a separate + # CPU-offloaded one (an embedded head stays on GPU). The tensor + # path reserves like the layer path; gate both on this. + _mtp_reserves_gpu = _mtp_will_engage and not _draft_cpu_no_embedded + _flat_mtp_reserve = ( + _MTP_VRAM_RESERVE_FRAC + if (_flat_mtp_engages and not _draft_cpu_no_embedded) + else 0.0 + ) + _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve - # Tensor mode allocates a compute-graph buffer on every - # participating GPU, so a GPU with less free VRAM than that - # reserve can't host it and would OOM at load. Drop those - # from the tensor-parallel set up front (gpu_indices below - # becomes the CUDA_VISIBLE_DEVICES mask, so they're excluded - # from llama-server entirely, not just given zero weight). + if tensor_parallel and effective_is_vision: + logger.info( + "Tensor parallelism skipped for vision model: " + "--split-mode tensor is incompatible with --mmproj " + "in the current llama.cpp build; using layer split." + ) + tensor_parallel = False + if _tensor_dropped_cache_type_kv is not None: + cache_type_kv = _tensor_dropped_cache_type_kv + _cache_type_from_env = False + extra_args = strip_split_mode_only( + _tensor_dropped_extra_args + if _tensor_dropped_extra_args is not None + else extra_args + ) + + # Tensor mode replicates a compute buffer on every GPU, so drop + # GPUs below that reserve from the set up front (gpu_indices + # becomes the CUDA_VISIBLE_DEVICES mask, fully excluding them). tp_gpus = gpus if tensor_parallel: - reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB - tp_gpus = [g for g in gpus if g[1] >= reserve_mib] + # Deterministic per-device compute buffer (replicated on + # every device in tensor mode); flat fallback when dims + # are unavailable. _plan_tensor_parallel uses the same. + _tp_reserve_bytes = self._estimate_compute_buffer_bytes( + n_ubatch = _effective_ubatch, + n_parallel = n_parallel, + per_device_tensor = True, + ) + reserve_mib = ( + _tp_reserve_bytes // (1024 * 1024) + if _tp_reserve_bytes > 0 + else self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + ) + # Admit by usable budget (free - (1-frac)*total), not raw + # free: a partly-used big card can clear the reserve on raw + # free yet have no budget left. + tp_gpus = [g for g in gpus if _gpu_usable(g) >= reserve_mib] if tensor_parallel and len(tp_gpus) < 2: # Tensor parallelism needs >= 2 usable GPUs. On a single @@ -3729,21 +5044,81 @@ class LlamaCppBackend: len(gpus), ) tensor_parallel = False - # A user --split-mode tensor in extras is appended after - # Studio's flags, so it would still reach llama-server and - # fail here; strip it so the downgrade actually applies. - extra_args = strip_split_mode_only(extra_args) + # Layer split supports a quantized KV the tensor attempt + # dropped; restore it and re-emit it (clear the env flag the + # tensor re-adoption may have set, so the restored type wins + # over a stale inherited env on the layer launch). + if _tensor_dropped_cache_type_kv is not None: + cache_type_kv = _tensor_dropped_cache_type_kv + _cache_type_from_env = False + # Restore the original extras (with the real, possibly + # asymmetric, --cache-type-k/-v the tensor attempt stripped), + # then drop the user --split-mode tensor so the downgrade + # actually applies (extras are appended last). + extra_args = strip_split_mode_only( + _tensor_dropped_extra_args + if _tensor_dropped_extra_args is not None + else extra_args + ) if tensor_parallel and tp_gpus: - # Tensor-parallel allocation: use all usable GPUs, weight - # the split by (free - buffer), and cap context to the - # pooled VRAM after weights + per-device compute-graph - # buffers. See _plan_tensor_parallel for the policy. + # Pooled usable budget (after each device's compute buffer) + # must hold the non-shrinkable footprint: weights + the MTP + # reserve. The planner can shrink ctx/KV, not these. + _tp_weight_budget_mib = ( + sum(_gpu_usable(g) for g in tp_gpus) - len(tp_gpus) * reserve_mib + ) + _tp_flat_mtp = 2 * 1024**3 # flat reserve when dims unavailable + if not _mtp_reserves_gpu: + # No MTP, or its only drafter is CPU-offloaded (no GPU). + _tp_mtp_floor = 0 + elif mtp_overhead_fn is not None and not _mtp_kv_unsized: + _tp_mtp_floor = _mtp_bytes( + min(2048, effective_ctx) if effective_ctx > 0 else 2048 + ) + else: + # Dims unavailable / weights-only: tensor mode has no + # --fit valve, so keep the flat reserve as the unsized-KV + # cushion, never below the known byte reserve. + _tp_mtp_floor = max( + _tp_flat_mtp, + _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), + ) + _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + if _tp_weight_budget_mib <= _tp_required_mib: + logger.info( + "Tensor parallelism requested but the pooled VRAM " + "budget cannot hold the weights, MTP reserve, and " + "per-device compute buffers; falling back to layer split." + ) + tensor_parallel = False + # Restore the dropped quantized KV (layer split supports + # it); clear the env flag so the restored type is emitted. + if _tensor_dropped_cache_type_kv is not None: + cache_type_kv = _tensor_dropped_cache_type_kv + _cache_type_from_env = False + # Restore the original (possibly asymmetric) cache extras + # too, dropping only the user --split-mode tensor. + extra_args = strip_split_mode_only( + _tensor_dropped_extra_args + if _tensor_dropped_extra_args is not None + else extra_args + ) + + if tensor_parallel and tp_gpus: + # Tensor-parallel allocation; see _plan_tensor_parallel. target_ctx = ( effective_ctx if explicit_ctx else (self._context_length or effective_ctx) ) + # When the draft KV couldn't be sized (weights-only reserve), + # the planner's mtp_overhead_fn is non-None but covers only + # weights, so pass the flat cushion for the unsized KV (else + # the binary search spends it on context). + _tp_unsized_mtp_reserve = ( + 2 * 1024**3 if (_mtp_reserves_gpu and _mtp_kv_unsized) else 0 + ) ( effective_ctx, max_available_ctx, @@ -3755,10 +5130,14 @@ class LlamaCppBackend: target_ctx, cache_type_kv = cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + mtp_flat_reserve_bytes = _tp_unsized_mtp_reserve, # Report the UI ceiling from native ctx, not the # explicit small request. max_target_ctx = self._context_length or target_ctx, + total_by_idx = total_by_idx, + n_ubatch = _effective_ubatch, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -3767,24 +5146,38 @@ class LlamaCppBackend: # bounds), independent of the currently requested context. native_ctx_for_cap = self._context_length or effective_ctx if native_ctx_for_cap > 0: - ranked_for_cap = sorted(gpus, key = lambda g: g[1], reverse = True) + ranked_for_cap = sorted( + gpus, + key = lambda g: _gpu_usable( + g, _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve + ), + reverse = True, + ) best_cap = 0 + _cap_fraction = _CTX_FIT_VRAM_FRACTION - _flat_mtp_reserve for n_gpus in range(1, len(ranked_for_cap) + 1): subset = ranked_for_cap[:n_gpus] - pool_mib = sum(free for _, free in subset) + # Per-GPU-consistent pool budget (fixes mixed + # known/unknown totals); pass it as an absolute + # budget so the fit and the check below agree. + pool_budget = _pool_budget_mib(subset, _cap_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( native_ctx_for_cap, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * (_CTX_FIT_VRAM_FRACTION - _mtp_reserve): + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: max_available_ctx = best_cap @@ -3798,34 +5191,49 @@ class LlamaCppBackend: # Honor the requested context verbatim. If it fits, # pin GPUs and skip --fit; else ship -c --fit # on and let llama-server flex -ngl (CPU offload). - requested_total = model_size + self._estimate_kv_cache_bytes( - effective_ctx, cache_type_kv, n_parallel = n_parallel + requested_total = ( + model_size_fit + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) + + _mtp_bytes(effective_ctx) ) gpu_indices, use_fit = self._select_gpus( - requested_total, gpus, usable_fraction = _pin_fraction + requested_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) # No silent shrink: effective_ctx stays == requested_ctx. else: # Auto context: prefer fewer GPUs, cap to fit. Same - # headroom threshold as _select_gpus (#5106). - ranked = sorted(gpus, key = lambda g: g[1], reverse = True) + # headroom threshold as _select_gpus (#5106). Rank by the + # active pin fraction so the order matches the fit budget. pin_fraction = _pin_fraction + ranked = sorted( + gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True + ) for n_gpus in range(1, len(ranked) + 1): subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) + pool_budget = _pool_budget_mib(subset, pin_fraction) + _ms = _subset_model_size(n_gpus) capped = self._fit_context_to_vram( effective_ctx, - pool_mib, - model_size, + pool_budget, + _ms, cache_type_kv, n_parallel = n_parallel, - mtp_engaged = _mtp_will_engage, + mtp_engaged = _mtp_reserves_gpu, + mtp_overhead_fn = mtp_overhead_fn, + budget_frac = 1.0, + total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) use_fit = False @@ -3838,14 +5246,17 @@ class LlamaCppBackend: if effective_ctx > 0: for n_gpus in range(1, len(ranked) + 1): subset = ranked[:n_gpus] - pool_mib = sum(free for _, free in subset) kv = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel, ) - total_mib = (model_size + kv) / (1024 * 1024) - if total_mib <= pool_mib * pin_fraction: + footprint_mib = ( + _subset_model_size(n_gpus) + + kv + + _mtp_bytes(effective_ctx) + ) / (1024 * 1024) + if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) use_fit = False break @@ -3857,14 +5268,36 @@ class LlamaCppBackend: "Falling back to file-size-only GPU selection", model_size_gb = round(model_size / (1024**3), 2), ) + # Add the byte-accurate MTP reserve here too when it is + # available; otherwise _pin_fraction carries the flat + # fallback (the two are mutually exclusive by design). + _fs_total = model_size_fit + _mtp_bytes( + self._context_length or effective_ctx or 4096 + ) gpu_indices, use_fit = self._select_gpus( - model_size, gpus, usable_fraction = _pin_fraction + _fs_total, + gpus, + usable_fraction = _pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = _pipeline_overhead_bytes, ) if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 # so the slider isn't on an unusable native ctx. effective_ctx = min(4096, effective_ctx) if effective_ctx > 0 else 4096 + # MTP reserve at the final context, for the logs below. + _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 + if _mtp_will_engage: + _mtp_note = ( + f"MTP reserve: {_mtp_reserve_bytes / (1024**3):.2f} GB " + f"(draft KV @ {effective_ctx} + verify n_max={_mtp_eff_n_max}" + + (", flat-frac fallback" if mtp_overhead_fn is None else "") + + "), " + ) + else: + _mtp_note = "" + if effective_ctx < original_ctx: kv_est = self._estimate_kv_cache_bytes( effective_ctx, cache_type_kv, n_parallel = n_parallel @@ -3872,7 +5305,9 @@ class LlamaCppBackend: logger.info( f"Context auto-reduced: {original_ctx} -> {effective_ctx} " f"(model: {model_size / (1024**3):.1f} GB, " - f"est. KV cache: {kv_est / (1024**3):.1f} GB)" + f"est. KV cache: {kv_est / (1024**3):.1f} GB, " + f"{_mtp_note}".rstrip(", ") + + ")" ) kv_cache_bytes = self._estimate_kv_cache_bytes( @@ -3885,6 +5320,7 @@ class LlamaCppBackend: f"GGUF size: {gguf_size / (1024**3):.1f} GB, " f"{mmproj_note}" f"est. KV cache: {kv_cache_bytes / (1024**3):.1f} GB, " + f"{_mtp_note}" f"context: {effective_ctx}, " f"GPUs free: {gpus}, selected: {gpu_indices}, fit: {use_fit}" ) @@ -3894,6 +5330,17 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # Unified-memory APUs load weights into system RAM (under WSL the VM + # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an + # oversize load the OS would otherwise kill mid-flight. Base model + # only: an optional MTP drafter is dropped by the MTP-drop fallback. + if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices): + _ram_msg = self._apu_ram_shortfall_message( + model_size, self._available_system_memory_mib() + ) + if _ram_msg: + raise RuntimeError(_ram_msg) + # Audio input straight from the mmproj (clip.has_audio_encoder), # independent of token names. self._mmproj_has_audio = False @@ -3924,27 +5371,49 @@ class LlamaCppBackend: "--no-context-shift", ] + fully_gpu_offloaded = False if use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: # Fits on selected GPU(s) -- offload all layers cmd.extend(["-ngl", "-1"]) + fully_gpu_offloaded = True + server_caps = self.probe_server_capabilities(binary) + # Expose Prometheus /metrics for the engine-stats logger, only + # when the binary advertises it (older/custom binaries may not). + if server_caps.get("supports_metrics"): + cmd.append("--metrics") cmd.extend( self._ctx_integrity_flags( n_parallel, use_fit, requested_ctx, effective_ctx, - self.probe_server_capabilities(binary), + server_caps, ) ) + offload_overridden = _extra_args_set_any_flag( + extra_args, _GPU_OFFLOAD_OVERRIDE_FLAGS + ) + threads_overridden = _extra_args_set_any_flag(extra_args, _THREAD_OVERRIDE_FLAGS) + full_offload_tuning_active = fully_gpu_offloaded and not offload_overridden - # -1 = llama.cpp auto-detect (physical cores). Pass explicitly - # so we don't inherit llama-server's internal default, which - # has varied (hardware concurrency incl. hyperthreads on some - # builds). - cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)]) + # Thread count: an unset --threads makes llama.cpp pick physical + # cores (common_cpu_get_num_math), but an explicit --threads -1 + # resolves to hardware_concurrency() (every hyperthread), which + # contends on the memory bus and slows CPU / hybrid decode. So + # omit the flag when unset and only pin it for an explicit + # override or the Windows full-offload OpenMP cap. Pass-through + # thread flags in extra_args still win (appended last). #5692 + if ( + sys.platform == "win32" + and full_offload_tuning_active + and not threads_overridden + ): + cmd.extend(["--threads", "2"]) + elif n_threads is not None and n_threads > 0: + cmd.extend(["--threads", str(n_threads)]) # Enable Jinja chat template rendering cmd.extend(["--jinja"]) @@ -3961,7 +5430,11 @@ class LlamaCppBackend: "iq4_nl", "f32", } - if cache_type_kv and cache_type_kv in _valid_cache_types: + if ( + cache_type_kv + and cache_type_kv in _valid_cache_types + and not _cache_type_from_env + ): cmd.extend( [ "--cache-type-k", @@ -3973,6 +5446,8 @@ class LlamaCppBackend: self._cache_type_kv = cache_type_kv logger.info(f"KV cache type: {cache_type_kv}") else: + # An env-only type is left inherited (untouched) so an + # asymmetric K/V env reaches the child as set. self._cache_type_kv = None # Tensor parallelism: split the model across GPUs by tensor @@ -4028,6 +5503,7 @@ class LlamaCppBackend: ) self._supports_reasoning = flags["supports_reasoning"] self._reasoning_style = flags["reasoning_style"] + self._reasoning_effort_levels = flags.get("reasoning_effort_levels", []) self._reasoning_always_on = flags["reasoning_always_on"] self._supports_preserve_thinking = flags["supports_preserve_thinking"] self._supports_tools = flags["supports_tools"] @@ -4086,6 +5562,28 @@ class LlamaCppBackend: else: self._api_key = None + # Windows + full offload: disable KV checkpoints (WDDM/PCI-E + # overhead). CPU/partial offload keeps prompt caching. #5692. + if sys.platform == "win32" and full_offload_tuning_active: + unsupported_cache_flags: list[str] = [] + if server_caps.get("supports_cache_ram"): + cmd.extend(["--cache-ram", "0"]) + else: + unsupported_cache_flags.append("--cache-ram") + if server_caps.get("supports_ctx_checkpoints"): + cmd.extend(["--ctx-checkpoints", "0"]) + else: + unsupported_cache_flags.append("--ctx-checkpoints") + if server_caps.get("supports_no_cache_prompt"): + cmd.append("--no-cache-prompt") + else: + unsupported_cache_flags.append("--no-cache-prompt") + if unsupported_cache_flags: + logger.info( + "Skipping unsupported Windows cache flags for llama-server: %s", + ", ".join(unsupported_cache_flags), + ) + # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Studio's auto-set flags. Already # validated by the route via validate_extra_args(). @@ -4101,15 +5599,49 @@ class LlamaCppBackend: logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") # Library paths so llama-server finds its shared libs and CUDA DLLs. - import os - import sys + env = self._llama_server_env_for_binary(binary) + # Omitting --threads relies on llama.cpp's physical-core default, so + # drop an inherited LLAMA_ARG_THREADS that would otherwise feed the + # arg handler and silently force hardware_concurrency(). #5692 + if "--threads" not in cmd: + env.pop("LLAMA_ARG_THREADS", None) - env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) + # Reconcile the inherited LLAMA_ARG_* env with Studio's final + # decision: stripping CLI extras on a tensor->layer downgrade + # can't remove env vars, so the child could run a mode/KV Studio + # didn't budget. + if not tensor_parallel: + # Layer split: clear a non-layer inherited split mode (and any + # paired tensor-split) so the child can't override the layer plan. + _inherited_sm = (env.get("LLAMA_ARG_SPLIT_MODE") or "").strip().lower() + if _inherited_sm and _inherited_sm != "layer": + env.pop("LLAMA_ARG_SPLIT_MODE", None) + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + else: + # Studio owns the tensor split: it emits --tensor-split when it + # picks an uneven one (CLI wins) and nothing when an even split + # is safe. Clear any inherited LLAMA_ARG_TENSOR_SPLIT so the even + # case can't be overridden by a stale env (the layer branch above + # clears it too). + env.pop("LLAMA_ARG_TENSOR_SPLIT", None) + # Tensor split aborts on a quantized KV; clear an inherited + # quantized cache type so the child uses the tensor-safe default. + for _ct_var in ("LLAMA_ARG_CACHE_TYPE_K", "LLAMA_ARG_CACHE_TYPE_V"): + _ct_raw = (env.get(_ct_var) or "").strip().lower() + if _ct_raw and _ct_raw not in self._TENSOR_PARALLEL_KV_TYPES: + env.pop(_ct_var, None) + + # Windows + full offload: PASSIVE OMP + 2 threads stop + # spin-wait burning CPU. CPU/partial offload keeps default + # OMP parallelism. #5692. + if sys.platform == "win32" and full_offload_tuning_active: + env.setdefault("OMP_WAIT_POLICY", "PASSIVE") + if not threads_overridden: + env.setdefault("OMP_NUM_THREADS", "2") # 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(): + if self._amd_apu_wants_unified_memory(gpu_indices): env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") @@ -4121,96 +5653,6 @@ class LlamaCppBackend: f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) - if sys.platform == "win32": - # Ordering: see _build_windows_path_dirs. #5106. - path_dirs = self._build_windows_path_dirs( - binary_dir, - sys.prefix, - os.environ.get("CUDA_PATH", ""), - ) - existing_path = env.get("PATH", "") - env["PATH"] = ";".join(path_dirs) + ";" + existing_path - - # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile - # kernel files (rocblas/library/*.dat + *.hsaco); the DLL - # searches /rocblas/library/ which doesn't exist - # -> silent crash on the first GEMM. ROCBLAS_TENSILE_LIBPATH - # repoints that search at the ROCm install. - _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: LD_LIBRARY_PATH for shared libs next to the binary - # plus CUDA runtime libs (libcudart, libcublas, etc.) - import platform - - lib_dirs = [] - # WSL: system HIP before the bundle's (which segfaults on - # /dev/dxg). Mirror install_llama_prebuilt.binary_env, which - # validates the prebuilt with this same ordering. - for _wsl_rocm in _wsl_system_rocm_lib_dirs(): - lib_dirs.append(_wsl_rocm) - if lib_dirs: - env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") - lib_dirs.append(binary_dir) - _arch = platform.machine() # x86_64, aarch64, etc. - - # Pip-installed nvidia CUDA runtime libs. The prebuilt - # binary links libcudart.so.13 / libcublas.so.13 which live - # here, not in /usr/local/cuda. - import glob as _glob - - for _nv_pattern in [ - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cu*", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cudnn", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "nvjitlink", - "lib", - ), - ]: - for _nv_dir in _glob.glob(_nv_pattern): - if os.path.isdir(_nv_dir): - lib_dirs.append(_nv_dir) - - for cuda_lib in [ - "/usr/local/cuda/lib64", - f"/usr/local/cuda/targets/{_arch}-linux/lib", - # Fallback CUDA compat paths (e.g. binary built with - # CUDA 12 where default /usr/local/cuda is CUDA 13+). - "/usr/local/cuda-12/lib64", - "/usr/local/cuda-12.8/lib64", - f"/usr/local/cuda-12/targets/{_arch}-linux/lib", - f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", - ]: - if os.path.isdir(cuda_lib): - lib_dirs.append(cuda_lib) - existing_ld = env.get("LD_LIBRARY_PATH", "") - new_ld = ":".join(lib_dirs) - env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld - # Pin to selected GPU(s). On ROCm, narrowing only # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full # set, so set HIP_VISIBLE_DEVICES too. @@ -4250,6 +5692,9 @@ class LlamaCppBackend: # retry once with --fit off before declaring the load failed. # Never retry when fit was requested (use_fit) or the caller # passed an explicit fit flag via extra args. + # Argv actually launched (post --fit off / MTP); text-only retry strips this. + _last_spawn_cmd = list(cmd) + def _spawn_and_wait(run_cmd, *, label = ""): """Start llama-server with run_cmd and wait for health. @@ -4257,6 +5702,7 @@ class LlamaCppBackend: crashes during startup and run_cmd is eligible (see _fit_off_retry_eligible). """ + nonlocal _last_spawn_cmd _fit_retry_allowed = self._fit_off_retry_eligible(run_cmd, use_fit) for _spawn_attempt in (0, 1): # Defensive kill: drop an orphan Popen a concurrent load may @@ -4291,6 +5737,7 @@ class LlamaCppBackend: # Best-effort; never block the load on logging. logger.debug(f"Could not open llama-server log file: {e}") self._llama_log_path = None + _last_spawn_cmd = list(run_cmd) self._process = subprocess.Popen( run_cmd, stdout = subprocess.PIPE, @@ -4298,7 +5745,9 @@ class LlamaCppBackend: text = True, env = env, **_windows_hidden_subprocess_kwargs(), + **_child_popen_kwargs(), ) + self._record_server_pid(self._process.pid) # Background thread to drain stdout (prevents pipe deadlock) self._stdout_thread = threading.Thread( @@ -4356,6 +5805,80 @@ class LlamaCppBackend: ) healthy = _spawn_and_wait(cmd) + # Flash-attention kernels hard-crash at startup on some ROCm/GPU + # builds (frequently inside the vision tower). Disabling FA keeps + # both vision and MTP, so retry that way before dropping either. + # Only on a hard fault with FA on; a cancel/unload stops respawn. + if not healthy and not self._cancel_event.is_set(): + _fa_rc = self._process.poll() if self._process is not None else None + _fa_cmd = ( + self._with_flash_attn_off(_last_spawn_cmd) + if self._is_signal_crash(_fa_rc) + else None + ) + if _fa_cmd is not None: + logger.warning( + "llama-server hard-crashed at startup (exit %s) with " + "flash attention on; retrying once with --flash-attn " + "off (keeps vision and MTP).", + _fa_rc, + ) + self._kill_process() + cmd = _fa_cmd + healthy = _spawn_and_wait(_fa_cmd, label = "-noflash") + + # MTP from Studio's spec flags or the user's (extra_args + # --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child + # only when neither emits a spec flag, so consult it only then. + _launch_spec_env: Mapping[str, str] = ( + os.environ + if (not _extra_args_set_spec_type(extra_args) and not spec_flags) + else {} + ) + _spec_requested_mtp = any( + "mtp" in str(t).lower() for t in spec_flags + ) or _extra_args_requests_mtp(extra_args, env = _launch_spec_env) + # Is the launched server actually running MTP+tensor? Gates the + # probe/watchdog/recovery; cleared if the MTP-drop fallback wins. + _mtp_active_for_launched_server = bool( + self._tensor_parallel and _spec_requested_mtp + ) + # MTP can pass /health then crash the flash-attn kernel on the + # first decode under tensor; probe one generation so the fallback + # catches that too. Tensor-only, so ordinary MTP stays probe-free. + if ( + healthy + and self._tensor_parallel + and _spec_requested_mtp + and not self._cancel_event.is_set() + and not self._probe_mtp_decode() + ): + # A first-decode hard fault is usually the FA kernel: retry + # FA-off (keeps MTP) before dropping speculative decoding below. + _probe_rc = self._process.poll() if self._process is not None else None + _fa_cmd = ( + self._with_flash_attn_off(_last_spawn_cmd) + if self._is_signal_crash(_probe_rc) + else None + ) + healthy = False + if _fa_cmd is not None: + logger.warning( + "MTP first-decode hard-crashed (exit %s) with flash " + "attention on; retrying with --flash-attn off.", + _probe_rc, + ) + self._kill_process() + cmd = _fa_cmd + healthy = ( + _spawn_and_wait(_fa_cmd, label = "-noflash-mtp") + and self._probe_mtp_decode() + ) + if not healthy: + logger.warning( + "MTP speculative decoding crashed on the first decode " + "under tensor parallelism; retrying without it." + ) # Any MTP request can abort the server: a separate drafter # (Gemma) on a binary that predates its arch, or an embedded # head (Qwen) the binary cannot build. Retry once with the @@ -4363,8 +5886,8 @@ class LlamaCppBackend: # loads. Gate on the spec block (not the drafter path, which # off/ngram local loads also carry) and keep # _requested_spec_mode so a duplicate /load doesn't thrash. The - # cancel check stops an /unload-killed attempt respawning. - _spec_requested_mtp = any("mtp" in str(t).lower() for t in spec_flags) + # cancel check stops an /unload-killed attempt respawning. A + # decode-probe failure above also routes here. if not healthy and _spec_requested_mtp and not self._cancel_event.is_set(): # Blame the binary only when the output shows MTP itself # failing (unknown arch / draft or context build); an @@ -4410,29 +5933,48 @@ class LlamaCppBackend: + ["--spec-default"] + cmd[_spec_start + len(spec_flags) :] ) + # User/env MTP survives in the tail; llama.cpp takes the last + # spec flag, so a trailing --spec-default overrides it too. + if _extra_args_requests_mtp(extra_args, env = _launch_spec_env): + fallback_cmd.append("--spec-default") healthy = _spawn_and_wait(fallback_cmd, label = "-retry") if healthy: self._speculative_type = "default" + _mtp_active_for_launched_server = False - # A vision GGUF launched with --mmproj can abort when the - # installed llama.cpp is too old for the model's projector - # ("Unknown projector type"); in that one case retry once - # text-only rather than failing the whole load. + # A too-old llama.cpp can reject a model's --mmproj projector + # (format message or a bare SIGSEGV); retry once text-only. if not healthy: out = "\n".join(self._stdout_lines[-50:]) + # Read the crash code before _kill_process() clears _process. + _crash_rc = self._process.poll() if self._process is not None else None self._kill_process() - if launched_with_mmproj and self._is_projector_incompatibility(out): + # Skip if a cancel/unload is pending (mirrors the MTP guard). + if ( + launched_with_mmproj + and not self._cancel_event.is_set() + and ( + self._is_projector_incompatibility(out) + or ( + self._is_signal_crash(_crash_rc) + and not self._output_has_nonprojector_diagnostic(out) + ) + ) + ): logger.warning( "llama-server could not load this model's vision " "projector (--mmproj). The installed llama.cpp build is " "likely too old for it. Loading text-only for this " "session; run 'unsloth studio update' to enable vision." ) - cmd = self._strip_mmproj_args(cmd) + cmd = self._strip_mmproj_args(_last_spawn_cmd) self._is_vision = False self._mmproj_has_audio = False self._start_llama_process(cmd, env) if not self._wait_for_health(timeout = 600.0): + # Read the exit code before _kill_process() clears it, so + # an OS-killed text-only retry still gets the OOM message. + _retry_rc = self._process.poll() if self._process is not None else None self._kill_process() raise RuntimeError( "Vision projector incompatible with this llama.cpp " @@ -4441,6 +5983,7 @@ class LlamaCppBackend: "\n".join(self._stdout_lines[-50:]), gguf_path, self._model_identifier, + _retry_rc, ) ) else: @@ -4449,6 +5992,7 @@ class LlamaCppBackend: out, gguf_path, self._model_identifier, + _crash_rc, ) ) @@ -4462,6 +6006,11 @@ class LlamaCppBackend: self._extra_args = list(extra_args) self._extra_args_source = (model_identifier, hf_variant) self._requested_n_ctx = int(n_ctx) + # Commit the known-good snapshot + whether MTP+tensor is live, then + # watch this load for a mid-generation crash. + self._last_load_kwargs = _pending_load_kwargs + self._mtp_runtime_fallback_active = _mtp_active_for_launched_server + self._start_mtp_crash_watchdog() # Catch silent CPU fallback when GPU was intended (#5106). self._gpu_offload_active = self._classify_gpu_offload( @@ -4481,6 +6030,18 @@ class LlamaCppBackend: logger.info( f"llama-server ready on port {self._port} for model '{model_identifier}'" ) + # Poll llama-server /metrics -> vLLM-style engine_stats logs + # (only when the binary exposes /metrics). + if server_caps.get("supports_metrics"): + try: + from core.inference.llama_stats import maybe_start_stats_logger + if self._stats_logger is not None: + self._stats_logger.stop() + self._stats_logger = maybe_start_stats_logger(self.base_url, logger) + except Exception as e: + logger.debug(f"engine-stats logger not started: {e}") + else: + self._stats_logger = None # Probe outside _lock (interruptible by /unload); init inside. self._is_audio = False @@ -4639,6 +6200,11 @@ class LlamaCppBackend: "run `unsloth studio update`. Loading without " "speculative decoding." ) + # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins + # over env) so the child matches the binary-capability gate and + # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. + flags.append("--spec-default") + self._speculative_type = "default" self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() @@ -4825,10 +6391,12 @@ class LlamaCppBackend: if _norm(self._cache_type_kv) != _norm(cache_type_kv): return False - # Reconcile a user --split-mode in extras (load_model does the same), so - # an extras-driven tensor load isn't seen as a mismatch that needlessly - # kills/reloads a healthy server. - if self._tensor_parallel != resolve_tensor_parallel(extra_args, tensor_parallel): + # Reconcile a user --split-mode in extras AND an inherited tensor + # LLAMA_ARG_SPLIT_MODE env, but only against a server that actually + # launched tensor: if load_model downgraded to layer split it scrubbed + # the child env, so the env must not force an endless reload of a healthy + # server. An identical request would downgrade the same way. + if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): return False # Compare on the canonical requested mode. With --spec-type in @@ -4881,26 +6449,13 @@ class LlamaCppBackend: def _classify_gpu_offload( self, expected_gpu: bool, detected_gpus: list[tuple[int, int]] ) -> Optional[bool]: - """True if a GPU model buffer was allocated, False if only CPU - buffers landed despite GPU intent, None when there's no signal (no - GPU detected, no buffer-size lines, etc.).""" + """True if the model landed on a GPU, False if only CPU buffers landed + despite GPU intent, None when there's no signal. Delegates to the shared + classifier so it tracks current llama.cpp logs (offloaded-layer counts / + device_info), not just the older "model buffer size" lines.""" if not detected_gpus or not expected_gpu: return None - # llama-server logs one "model buffer size = N MiB" line per backend - # buffer; CUDA/ROCm/Metal/Vulkan/OpenCL/SYCL are GPU, CPU* are not. - gpu_markers = ("CUDA", "ROCm", "Metal", "Vulkan", "OpenCL", "SYCL") - saw_buffer_line = False - saw_gpu_buffer = False - for line in self._stdout_lines: - if "model buffer size" not in line: - continue - saw_buffer_line = True - if any(marker in line for marker in gpu_markers): - saw_gpu_buffer = True - break - if not saw_buffer_line: - return None - return saw_gpu_buffer + return classify_gpu_offload_lines(self._stdout_lines) def load_cancelled(self) -> bool: """True if a load was cancelled (e.g. via unload/_cancel_event) and not @@ -4919,6 +6474,8 @@ class LlamaCppBackend: self._hf_repo = None self._mtp_draft_path = None self._spec_fallback_reason = None + self._last_load_kwargs = None + self._mtp_runtime_fallback_active = False self._hf_variant = None self._is_vision = False self._is_audio = False @@ -4936,6 +6493,7 @@ class LlamaCppBackend: self._supports_reasoning = False self._reasoning_always_on = False self._reasoning_style = "enable_thinking" + self._reasoning_effort_levels = [] self._reasoning_default = True self._supports_preserve_thinking = False self._supports_tools = False @@ -4982,6 +6540,9 @@ class LlamaCppBackend: def _kill_process(self): """Terminate the subprocess if running.""" + # Stop the watchdog before a deliberate kill so a planned reload/unload + # isn't seen as a crash; a real crash never routes through here. + self._stop_mtp_crash_watchdog() if self._process is None: return try: @@ -4994,15 +6555,22 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"Error killing llama-server process: {e}") finally: + # getattr: teardown must tolerate a partially-built backend (failed + # __init__ or a __new__-built instance), as with _llama_log_fh below. + if getattr(self, "_stats_logger", None) is not None: + self._stats_logger.stop() + self._stats_logger = None self._process = None + self._clear_server_pid() # Clear healthy so a /load during the replacement's warm-up can't # short-circuit against the previous server's health (#5401). self._healthy = False # Drives _wait_for_vram_settle in the next load_model; set in finally # so both in-process and frontend Apply paths record the kill. self._last_kill_monotonic = time.monotonic() - if self._stdout_thread is not None: - self._stdout_thread.join(timeout = 2) + stdout_thread = getattr(self, "_stdout_thread", None) + if stdout_thread is not None: + stdout_thread.join(timeout = 2) self._stdout_thread = None fh = getattr(self, "_llama_log_fh", None) if fh is not None: @@ -5013,7 +6581,199 @@ class LlamaCppBackend: self._llama_log_fh = None @staticmethod - def _kill_orphaned_servers(): + def _server_pidfile_path() -> Optional[Path]: + """Pidfile recording the live llama-server PID, under the active studio root + (per-root, so concurrent Studios with distinct UNSLOTH_STUDIO_HOME stay + isolated, mirroring the reaper's custom-root isolation).""" + try: + from utils.paths.storage_roots import studio_root # noqa: WPS433 + return studio_root() / "llama-server.pid" + except Exception: + return None + + @classmethod + def _record_server_pid(cls, pid: int) -> None: + """Best-effort record of the spawned llama-server PID for orphan reaping. + + Stores ``pid:starttime`` so a later startup can reject a PID that has + since been recycled to a different process (see ``_pid_start_identity``). + A bare ``pid`` (no identity) is still accepted on read for compatibility. + """ + path = cls._server_pidfile_path() + if path is None: + return + try: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(f"{pid}:{cls._pid_start_identity(pid)}") + except Exception as e: + logger.debug(f"Could not write llama-server pidfile: {e}") + + @classmethod + def _clear_server_pid(cls) -> None: + """Best-effort removal of the llama-server pidfile.""" + path = cls._server_pidfile_path() + if path is None: + return + try: + path.unlink(missing_ok = True) + except Exception as e: + logger.debug(f"Could not remove llama-server pidfile: {e}") + + @staticmethod + def _pid_is_llama_server(pid: int) -> bool: + """True only if pid is a live process whose binary is a llama-server. Guards + against PID reuse before killing a recorded orphan; returns False on any + uncertainty so an unrelated process is never killed.""" + try: + import psutil + try: + proc = psutil.Process(pid) + if (proc.name() or "").lower().startswith("llama-server"): + return True + return Path(proc.exe() or "").name.lower().startswith("llama-server") + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return False + except ImportError: + pass + if sys.platform != "linux": + return False + try: + if Path(os.readlink(f"/proc/{pid}/exe")).name.lower().startswith("llama-server"): + return True + except OSError: + pass + try: + with open(f"/proc/{pid}/cmdline", "rb") as fh: + tokens = fh.read().split(b"\x00") + first = tokens[0].decode("utf-8", "replace") if tokens else "" + return Path(first).name.lower().startswith("llama-server") + except OSError: + return False + + @staticmethod + def _pid_start_identity(pid: int) -> str: + """Stable per-PID identity (process start time) guarding against PID reuse. + + Returns a token string, or "" when it cannot be determined (the caller + then falls back to the llama-server name check only).""" + try: + import psutil + try: + return str(psutil.Process(pid).create_time()) + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + return "" + except ImportError: + pass + if sys.platform == "linux": + try: + with open(f"/proc/{pid}/stat", "rb") as fh: + data = fh.read() + # field 22 (starttime), counted from after the ")" that closes comm. + return data[data.rfind(b")") + 2 :].split()[19].decode() + except (OSError, IndexError): + return "" + return "" + + @staticmethod + def _pid_parent_is_alive(pid: int) -> bool: + """True if the recorded server's parent is still running, i.e. the server is + NOT orphaned. Lets the cross-session reap kill only a true orphan (parent + gone) and never a live server owned by a running Studio, regardless of which + process performs the sweep. Biased toward "alive" on uncertainty so a live + server is never mistakenly reaped.""" + try: + import psutil + + try: + ppid = psutil.Process(pid).ppid() + except psutil.NoSuchProcess: + return False # the recorded server itself is gone + except psutil.Error: + return True # cannot tell -- never risk killing a live server + if ppid <= 1: + return False # reparented to init -> orphan + return psutil.pid_exists(ppid) + except ImportError: + pass + if sys.platform == "linux": + try: + with open(f"/proc/{pid}/stat", "rb") as fh: + data = fh.read() + ppid = int(data[data.rfind(b")") + 2 :].split()[1]) + except (OSError, IndexError, ValueError): + return False + if ppid <= 1: + return False + return Path(f"/proc/{ppid}").exists() + return False + + @staticmethod + def _unlink_pidfile(path: Path) -> None: + """Best-effort removal of a resolved pidfile path.""" + try: + path.unlink(missing_ok = True) + except Exception: + pass + + @classmethod + def _reap_recorded_pid(cls) -> int: + """Kill the exact llama-server PID recorded at spawn, but only when it is a + genuine orphan -- its parent (the Studio that spawned it) is gone. This is + the cross-session backstop the parent-death reaper (Job Object / + PR_SET_PDEATHSIG) cannot cover: an orphan left by an already-dead Studio + (macOS, a best-effort failure, or a pre-existing orphan). Path-independent, + so it also catches an orphan the install-root match would miss. + + A live server whose parent is still running is never reaped, so constructing + a second backend in-process (the helper / advisor paths each build a + LlamaCppBackend) cannot kill the active chat server. A recorded PID that has + been recycled to a different process is rejected by the start-time identity + and the llama-server name check, so unrelated user processes are never + touched. SIGKILL falls back to SIGTERM on Windows, where os.kill maps it to + TerminateProcess and SIGKILL is undefined.""" + path = cls._server_pidfile_path() + if path is None or not path.exists(): + return 0 + + pid = -1 + identity = "" + try: + pid_str, _, identity = path.read_text().strip().partition(":") + pid = int(pid_str) + except Exception: + pid = -1 + + if pid <= 0: + cls._unlink_pidfile(path) # garbage record + return 0 + if pid == os.getpid(): + return 0 # never our own pid; leave the record alone + + if cls._pid_parent_is_alive(pid): + # Live server with a running parent -> not an orphan; keep the record so + # a later startup can still reap it if that parent later dies abnormally. + return 0 + + # Parent is gone: candidate orphan. Reject a PID recycled to something else. + if identity and cls._pid_start_identity(pid) != identity: + cls._unlink_pidfile(path) + return 0 + + killed = 0 + if cls._pid_is_llama_server(pid): + try: + os.kill(pid, getattr(signal, "SIGKILL", signal.SIGTERM)) + killed = 1 + logger.info(f"Killed orphaned llama-server from pidfile (pid={pid})") + except (ProcessLookupError, PermissionError): + pass + except Exception as e: + logger.debug(f"Could not kill recorded llama-server pid {pid}: {e}") + cls._unlink_pidfile(path) + return killed + + @staticmethod + def _kill_orphaned_servers() -> int: """Kill orphaned llama-server processes started by studio. Only kills processes whose resolved binary lives under a known @@ -5025,7 +6785,15 @@ class LlamaCppBackend: Uses psutil for cross-platform support (Linux, macOS, Windows); falls back to pgrep + /proc//exe on Linux when psutil is absent. + + Returns the count of processes killed; callers arm the VRAM-settle + wait on a positive count. """ + # Cross-session backstop first: reap the exact PID we recorded at spawn, + # but only if it is a true orphan whose parent is gone (so a helper backend + # built while a chat server is live can never kill it). The root-gated + # enumeration below stays as a fallback. + killed = LlamaCppBackend._reap_recorded_pid() try: # -- Build the ownership allowlist -------------------------------- # exact_binaries -- env var overrides (exact path match). @@ -5117,6 +6885,7 @@ class LlamaCppBackend: continue proc.kill() + killed += 1 logger.info( f"Killed orphaned llama-server process (pid={proc.info['pid']})" ) @@ -5129,7 +6898,7 @@ class LlamaCppBackend: else: # -- Fallback: pgrep + /proc//exe (Linux only) ----------- if sys.platform != "linux": - return + return killed result = subprocess.run( ["pgrep", "-a", "-f", "llama-server"], capture_output = True, @@ -5138,7 +6907,7 @@ class LlamaCppBackend: env = child_env_without_native_path_secret(), ) if result.returncode != 0: - return + return killed for line in result.stdout.strip().splitlines(): parts = line.strip().split(None, 1) @@ -5169,6 +6938,7 @@ class LlamaCppBackend: try: os.kill(pid, signal.SIGKILL) + killed += 1 logger.info(f"Killed orphaned llama-server process (pid={pid})") except ProcessLookupError: pass @@ -5176,6 +6946,7 @@ class LlamaCppBackend: pass except Exception: logger.warning("Error during orphan server cleanup", exc_info = True) + return killed def _cleanup(self): """atexit handler to ensure llama-server is terminated.""" @@ -5198,6 +6969,139 @@ class LlamaCppBackend: return False return True + def _probe_mtp_decode(self, timeout: float = 60.0) -> bool: + """One tiny /completion to confirm MTP survives the first decode. + + MTP-draft can pass /health yet crash the flash-attn kernel only once + tokens generate (e.g. under --split-mode tensor). False on any error so + the caller can drop MTP and retry. + """ + url = f"http://127.0.0.1:{self._port}/completion" + payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False} + # Match the --api-key auth direct-stream mode uses, else a spurious 401. + headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + try: + resp = httpx.post(url, json = payload, timeout = timeout, headers = headers) + except Exception as e: + logger.debug(f"MTP decode probe failed: {e}") + return False + if resp.status_code != 200: + logger.debug(f"MTP decode probe returned HTTP {resp.status_code}") + return False + # A crash can drop the connection or kill the process right after a reply. + if self._process is not None and self._process.poll() is not None: + return False + return True + + def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool: + """Schedule one background reload without MTP after a mid-generation death. + + MTP+tensor can crash the flash-attn kernel on a later request, after + load_model returned, past the load-time fallback and decode probe. Not a + persistent ban: a fresh load re-tries MTP. Returns True if scheduled. + """ + # Cheap async-safe gate: only our live MTP+tensor launch, not cancelled, + # with a snapshot to replay. + if self._cancel_event.is_set(): + return False + if not self._mtp_runtime_fallback_active: + return False + if not self._last_load_kwargs or self._process is None: + return False + # Single-flight: the first failure claims the reload. + with self._mtp_runtime_fallback_lock: + if self._mtp_runtime_fallback_in_progress: + return False + self._mtp_runtime_fallback_in_progress = True + snapshot = dict(self._last_load_kwargs) + proc = self._process + + def _recover(): + try: + # Confirm the process really exited (the error can arrive a beat + # early) so a transient stream error can't disable MTP. + deadline = time.monotonic() + 5.0 + while proc.poll() is None and time.monotonic() < deadline: + time.sleep(0.1) + if proc.poll() is None: + logger.debug("Generation error but llama-server is alive; keeping MTP.") + return + logger.warning( + "llama-server exited mid-generation with MTP under tensor " + "parallelism (%s); reloading without speculative decoding.", + type(exc).__name__ if exc is not None else "server exited", + ) + # Re-check under the load lock (RLock allows the nested + # load_model) so a newer load isn't clobbered by this stale replay. + requested_mode = snapshot.get("speculative_type") + with self._serial_load_lock: + if self._cancel_event.is_set(): + logger.info("MTP-crash reload skipped: load was cancelled/unloaded.") + return + if self._process is not proc: + logger.info("MTP-crash reload skipped: a newer load is already active.") + return + if self._last_load_kwargs != snapshot: + logger.info("MTP-crash reload skipped: load settings changed.") + return + snapshot["speculative_type"] = "off" + # Drop user/env MTP too: append a last-wins --spec-default. + _ea = list(snapshot.get("extra_args") or []) + if _extra_args_requests_mtp(_ea, env = os.environ): + _ea.append("--spec-default") + snapshot["extra_args"] = _ea + self.load_model(**snapshot) + # Restore the requested mode + reason load_model("off") cleared, + # so /status shows the user's mode + note (like the startup fallback). + self._requested_spec_mode = _canonicalize_spec_mode(requested_mode) + self._spec_fallback_reason = "runtime_error" + logger.info("Reloaded without MTP after the tensor-parallel crash.") + except Exception as e: + logger.error(f"Reload without MTP failed: {e}") + finally: + with self._mtp_runtime_fallback_lock: + self._mtp_runtime_fallback_in_progress = False + + threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start() + return True + + def _start_mtp_crash_watchdog(self) -> None: + """Background poll that recovers on an MTP+tensor crash even when no + request observes it (direct proxy endpoints, or nothing in flight). + + Armed only for a live MTP+tensor launch; the no-MTP reload disarms it, so + it can't loop. + """ + if not self._mtp_runtime_fallback_active: + return + proc = self._process + if proc is None: + return + # Replace any prior watchdog (loads are serialised, so at most one). + self._stop_mtp_crash_watchdog() + stop = threading.Event() + self._mtp_watchdog_stop = stop + + def _watch(): + # Exit on stop or process death. _kill_process sets stop before + # terminating, so re-check it: only a real crash (stop unset) recovers. + while not stop.wait(1.0): + if proc.poll() is not None: + if not stop.is_set(): + self._maybe_recover_from_mtp_crash() + return + + t = threading.Thread(target = _watch, daemon = True, name = "mtp-crash-watchdog") + self._mtp_watchdog_thread = t + t.start() + + def _stop_mtp_crash_watchdog(self) -> None: + """Signal the crash watchdog to exit; called before any deliberate kill.""" + stop = getattr(self, "_mtp_watchdog_stop", None) + if stop is not None: + stop.set() + self._mtp_watchdog_thread = None + def _wait_for_health( self, timeout: float = 120.0, @@ -5495,6 +7399,38 @@ class LlamaCppBackend: finally: _cancel_closed.set() + def _respawn_if_dead(self) -> bool: + """Relaunch the llama-server if its process has exited. + + A loaded chat model can be SIGKILL'd mid-session (usually GPU/RAM pressure + from a training run on the same box), leaving a defunct process while + ``is_loaded`` still reads True. Replay the last ``load_model`` call to + recover, returning True once healthy. Serialised on ``_respawn_lock`` so + many generations hitting the dead server trigger at most one reload. + """ + with self._respawn_lock: + proc = self._process + if proc is None: + return False + if proc.poll() is None: + # Process is alive: either a concurrent caller already respawned + # it (healthy), or this connection error wasn't a dead server. + return self._healthy + kwargs = self._last_load_kwargs + if not kwargs: + return False + logger.warning( + f"llama-server for '{self._model_identifier}' exited " + f"(code {proc.returncode}); respawning to recover the session" + ) + with self._lock: + self._healthy = False + try: + return bool(self.load_model(**kwargs)) + except Exception as exc: + logger.error(f"Failed to respawn llama-server: {exc}") + return False + def generate_chat_completion( self, messages: list[dict], @@ -5512,7 +7448,8 @@ class LlamaCppBackend: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, seed: Optional[int] = None, - ) -> Generator[str | dict, None, None]: + _allow_respawn_retry: bool = True, + ) -> Generator[Union[str, dict], None, None]: """ Send a chat completion to llama-server and stream tokens back. @@ -5677,11 +7614,43 @@ class LlamaCppBackend: "finish_reason": _metadata_finish_reason, } - except httpx.ConnectError: + except httpx.ConnectError as e: + # Server already down. If this was an MTP+tensor crash, recover by + # reloading without MTP (scheduled in the background) and fail this + # request. Otherwise the server was likely SIGKILL'd by GPU pressure + # from a concurrent training run: respawn the same config and retry the + # generation once (bounded by the private flag, no duplicate output). + if self._maybe_recover_from_mtp_crash(e): + raise RuntimeError("Lost connection to llama-server") + if _allow_respawn_retry and not cumulative and self._respawn_if_dead(): + logger.warning( + "llama-server was unreachable; respawned it and retrying the generation" + ) + yield from self.generate_chat_completion( + messages, + image_b64 = image_b64, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + max_tokens = max_tokens, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + stop = stop, + cancel_event = cancel_event, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + seed = seed, + _allow_respawn_retry = False, + ) + return raise RuntimeError("Lost connection to llama-server") except Exception as e: if cancel_event is not None and cancel_event.is_set(): return + # Died mid-generation: recover MTP, re-raise unchanged for this request. + self._maybe_recover_from_mtp_crash(e) raise # ── Tool-calling agentic loop ────────────────────────────── @@ -5710,6 +7679,7 @@ class LlamaCppBackend: seed: Optional[int] = None, disable_parallel_tool_use: bool = False, confirm_tool_calls: bool = False, + bypass_permissions: bool = False, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -6410,7 +8380,9 @@ class LlamaCppBackend: decision.as_assistant_tool_call() ) - needs_confirm = bool(confirm_tool_calls) + # Bypass wins over the confirm gate at the loop level too, + # so a direct internal caller with both flags never prompts. + needs_confirm = bool(confirm_tool_calls) and not bypass_permissions approval_id = new_approval_id() if needs_confirm else "" decision_slot = ( begin_tool_decision(session_id, approval_id) if needs_confirm else None @@ -6471,6 +8443,7 @@ class LlamaCppBackend: timeout = _effective_timeout, session_id = session_id, rag_scope = rag_scope, + disable_sandbox = bypass_permissions, ) if decision.tool_name == "search_knowledge_base": _kb_search_count += 1 @@ -6776,7 +8749,7 @@ class LlamaCppBackend: try: # llama-server's /apply-template renders tool declarations # into the prompt when ``tools`` is supplied, so pass them - # through — otherwise tool-schema tokens go uncounted. + # through, otherwise tool-schema tokens go uncounted. template_body = {"messages": template_messages} if tools: template_body["tools"] = tools diff --git a/studio/backend/core/inference/llama_http.py b/studio/backend/core/inference/llama_http.py new file mode 100644 index 0000000000..b554949c3e --- /dev/null +++ b/studio/backend/core/inference/llama_http.py @@ -0,0 +1,56 @@ +# 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 pooled httpx.AsyncClient for NON-streaming calls to the local llama-server. + +Streaming generation must NOT use this. It relies on ``Connection: close`` and +``max_keepalive_connections=0`` so a client disconnect tears down the upstream +socket and stops GPU decode (PR #5749). This pooled client is only for short +request/response proxy calls (non-streaming completions, embeddings) where +reusing a connection removes per-request setup cost. Per-request ``timeout`` is +still passed at each call site. +""" + +from __future__ import annotations + +import asyncio +import weakref + +import httpx + +_LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32) + + +def _new_client() -> httpx.AsyncClient: + try: + return httpx.AsyncClient(limits = _LIMITS) + except Exception: + # Mirror external_provider: an unsupported env proxy scheme can raise. + return httpx.AsyncClient(limits = _LIMITS, trust_env = False) + + +# One client per running event loop: an httpx client binds its transport to the +# loop it first runs on, so a single global instance breaks across a lifespan +# restart or a second test loop. Weak keys let a finished loop drop its client. +_clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]" = ( + weakref.WeakKeyDictionary() +) + + +def nonstreaming_client() -> httpx.AsyncClient: + loop = asyncio.get_running_loop() + client = _clients.get(loop) + if client is None or client.is_closed: + client = _new_client() + _clients[loop] = client + return client + + +async def aclose() -> None: + clients = list(_clients.values()) + _clients.clear() + for client in clients: + try: + await client.aclose() + except Exception: + pass diff --git a/studio/backend/core/inference/llama_server_args.py b/studio/backend/core/inference/llama_server_args.py index 69a86fa3ba..b42be5ee0d 100644 --- a/studio/backend/core/inference/llama_server_args.py +++ b/studio/backend/core/inference/llama_server_args.py @@ -13,7 +13,8 @@ Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md from __future__ import annotations -from typing import Iterable, Optional +import os +from typing import Iterable, Mapping, Optional # Each group = every alias (short + long) of one hard-denied flag. # Extend the matching group when llama.cpp adds a new alias. @@ -124,7 +125,9 @@ def is_managed_flag(flag: str) -> bool: # from inherited extras so they can't last-wins-override an Apply that # re-sets the same field. _CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"}) -_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}) +_CACHE_TYPE_K_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k"}) +_CACHE_TYPE_V_FLAGS: frozenset[str] = frozenset({"-ctv", "--cache-type-v"}) +_CACHE_FLAGS: frozenset[str] = _CACHE_TYPE_K_FLAGS | _CACHE_TYPE_V_FLAGS _SPEC_FLAGS: frozenset[str] = frozenset( { "--spec-default", @@ -133,13 +136,22 @@ _SPEC_FLAGS: frozenset[str] = frozenset( "--spec-ngram-size", "--draft-min", "--draft-max", - # MTP path (llama.cpp #22673). --model-draft and aliases are - # Studio-managed since the separate-drafter support (Gemma 4): an - # inherited copy must not last-wins-override the auto-detected - # drafter. Explicit extras for the current load are never stripped. + # MTP path (llama.cpp #22673). The drafter selectors (local --model-draft + # and HF --spec-draft-hf aliases) are Studio-managed since the separate- + # drafter support (Gemma 4): an inherited copy must not last-wins-override + # the auto-detected drafter. Explicit extras for the current load are never + # stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld, + # --spec-draft-device) are deliberately NOT stripped: the VRAM budget reads + # them via the same parsers the child honors, so they stay consistent on + # inherit, and stripping them would silently move a CPU-offloaded drafter + # back onto the GPU. "--model-draft", "-md", "--spec-draft-model", + "--spec-draft-hf", + "-hfd", + "-hfrd", + "--hf-repo-draft", "--spec-draft-n-max", "--spec-draft-n-min", "--spec-draft-p-min", @@ -274,6 +286,20 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]: return _last_flag_value(args, _CACHE_FLAGS) +def parse_cache_override_per_axis( + args: Optional[Iterable[str]], +) -> tuple[Optional[str], Optional[str]]: + """Last-wins --cache-type-k / --cache-type-v values kept apart, as (k, v). + + parse_cache_override collapses both axes to one last-wins value; this keeps + them separate so an asymmetric K/V can be budgeted by its heavier axis. + """ + return ( + _last_flag_value(args, _CACHE_TYPE_K_FLAGS), + _last_flag_value(args, _CACHE_TYPE_V_FLAGS), + ) + + def resolve_cache_type_kv( args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str] ) -> Optional[str]: @@ -309,6 +335,60 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral return override.strip().lower() == "tensor" +def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool: + """True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio + emits --split-mode only on its tensor branch, so a tensor env on the layer + path would run the child tensor-parallel unbudgeted; this flips the budget + to tensor. Only tensor is heavier, so other modes are ignored.""" + raw = (os.environ if env is None else env).get("LLAMA_ARG_SPLIT_MODE") + return bool(raw) and raw.strip().lower() == "tensor" + + +def _effective_tensor_parallel( + extra_args: Optional[Iterable[str]], + tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Tensor-parallel decision including the inherited LLAMA_ARG_SPLIT_MODE env. + + resolve_tensor_parallel (extras + toggle), flipped on when extras set no split + mode but the child inherits a tensor split env. Shared by load_model (which + budgets and launches it) and the tensor-fallback wrapper (so an env-only + tensor crash still retries layer split).""" + resolved = resolve_tensor_parallel(extra_args, tensor_parallel) + if ( + not resolved + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + return True + return resolved + + +def _tensor_parallel_matches_loaded( + extra_args: Optional[Iterable[str]], + requested_tensor_parallel: bool, + loaded_tensor_parallel: bool, + env: Optional[Mapping[str, str]] = None, +) -> bool: + """Whether a duplicate load request matches a loaded server's tensor state. + + Env-only tensor mode is a launch hint load_model may downgrade to layer split + (capacity/buffer), scrubbing the child env. So only let an inherited tensor env + raise a match against a server that *actually* launched tensor; on a downgraded + (layer) server the env is ignored, and an identical request would downgrade the + same way -- avoiding an endless reload of a healthy server.""" + requested = resolve_tensor_parallel(extra_args, requested_tensor_parallel) + if ( + loaded_tensor_parallel + and not requested + and parse_split_mode_override(extra_args) is None + and _env_split_mode_is_tensor(env) + ): + requested = True + return requested == loaded_tensor_parallel + + _MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"}) _MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"}) diff --git a/studio/backend/core/inference/llama_stats.py b/studio/backend/core/inference/llama_stats.py new file mode 100644 index 0000000000..6047aedbc0 --- /dev/null +++ b/studio/backend/core/inference/llama_stats.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Translate llama-server's Prometheus /metrics into a periodic, vLLM-style +engine-stats log line (generation/prompt throughput, requests in flight). + +llama-server already computes these (it needs `--metrics`); this lifts them +into Studio's structured log so the terminal shows serving health, not just +per-request access lines. Emitted only while there is activity. +""" + +import os +import re +import threading +import time +import urllib.request + +# Prometheus body lines: "llamacpp:[{labels}] " (skip "#" HELP/TYPE). +_METRIC_RE = re.compile(r"^llamacpp:(\w+)(?:\{[^}]*\})?\s+([0-9.eE+-]+)", re.MULTILINE) +_OFF = {"0", "false", "no", "off"} + + +class LlamaServerStatsLogger: + """Daemon poller that logs vLLM-style engine stats from llama-server. + + Keeps retrying through transient scrape failures; the backend stops it via + stop() on unload/reload, so a brief /metrics stall does not silence stats. + """ + + def __init__( + self, + base_url, + logger, + interval_s = 10.0, + ): + self._url = f"{base_url.rstrip('/')}/metrics" + self._log = logger + self._interval = max(1.0, float(interval_s)) + self._stop = threading.Event() + self._thread = None + + def start(self): + if self._thread is None: + self._thread = threading.Thread(target = self._run, name = "llama-stats", daemon = True) + self._thread.start() + + def stop(self): + self._stop.set() + + def _scrape(self): + try: + with urllib.request.urlopen(self._url, timeout = 3) as r: + if r.status != 200: + return None + body = r.read().decode("utf-8", "replace") + except Exception: + return None + out = {} + for k, v in _METRIC_RE.findall(body): + try: # a malformed value must not kill the daemon thread + out[k] = float(v) + except ValueError: + continue + return out + + def _run(self): + misses = 0 + prev = None # (monotonic_t, tokens_predicted_total, prompt_tokens_total) + while not self._stop.wait(self._interval): + m = self._scrape() + if not m: + misses += 1 + if misses == 3: # transient stall (load/GC); keep polling. + self._log.debug("engine_stats: /metrics scrape failing, still retrying") + continue # real shutdown is driven by stop() from _kill_process + misses = 0 + # Generation tokens come from tokens_predicted_total (counter) and + # predicted_tokens_seconds (gauge); n_decode_total counts + # llama_decode() calls, not tokens, so it must not feed tok/s. + now = time.monotonic() + predicted = m.get("tokens_predicted_total", 0.0) + prompt = m.get("prompt_tokens_total", 0.0) + gen_delta = prompt_delta = 0.0 + if prev is not None and now > prev[0]: + dt = now - prev[0] + gen_delta = max(0.0, (predicted - prev[1]) / dt) + prompt_delta = max(0.0, (prompt - prev[2]) / dt) + prev = (now, predicted, prompt) + # Prefer llama.cpp's own throughput gauges; fall back to the counter + # delta for binaries that expose only the counters. + gen_tps = m.get("predicted_tokens_seconds") or gen_delta + prompt_tps = m.get("prompt_tokens_seconds") or prompt_delta + running, waiting = ( + int(m.get("requests_processing", 0)), + int(m.get("requests_deferred", 0)), + ) + # Gate on real activity this tick so a stale gauge never logs at idle. + if running or waiting or gen_delta or prompt_delta: + self._log.info( + "engine_stats", + gen_tok_s = round(float(gen_tps), 1), + prompt_tok_s = round(float(prompt_tps), 1), + running = running, + waiting = waiting, + ) + + +def maybe_start_stats_logger(base_url, logger): + """Start a stats logger unless UNSLOTH_STUDIO_ENGINE_STATS disables it.""" + if (os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS", "1") or "").strip().lower() in _OFF: + return None + try: + interval = float(os.environ.get("UNSLOTH_STUDIO_ENGINE_STATS_INTERVAL_S", "10")) + except ValueError: + interval = 10.0 + sl = LlamaServerStatsLogger(base_url, logger, interval) + sl.start() + return sl diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4ccac2912e..47e8038764 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -18,7 +18,6 @@ import atexit import base64 import os import signal -import structlog from loggers import get_logger import multiprocessing as mp import queue @@ -30,15 +29,15 @@ from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union from utils.hardware import prepare_gpu_selection +# Re-exported from the shared helper so GGUF, training, and inference share one +# type; kept importable here for backwards compatibility. +from utils.hf_xet_fallback import DownloadStallError + logger = get_logger(__name__) _CTX = mp.get_context("spawn") -class DownloadStallError(RuntimeError): - """Raised when the worker reports no download progress for too long.""" - - # Dispatcher timeout constants (seconds) _DISPATCH_READ_TIMEOUT = 30.0 _DISPATCH_POLL_INTERVAL = 0.5 @@ -61,7 +60,6 @@ class InferenceOrchestrator: self._cmd_queue: Any = None self._resp_queue: Any = None self._cancel_event: Any = None # mp.Event — set to cancel generation - self._lock = threading.Lock() self._gen_lock = threading.Lock() # Serializes generation # Dispatcher state for compare mode (adapter-controlled requests): @@ -76,7 +74,6 @@ class InferenceOrchestrator: self.active_model_name: Optional[str] = None self.models: dict = {} self.loading_models: set = set() - self.loaded_local_models: list = [] from core.inference.defaults import get_default_models self._static_models = get_default_models() @@ -84,8 +81,6 @@ class InferenceOrchestrator: self._top_hub_cache: Optional[list[str]] = None self._top_models_ready = threading.Event() - self._current_transformers_major: Optional[str] = None # "4" or "5" - atexit.register(self._cleanup) logger.info("InferenceOrchestrator initialized (subprocess mode)") @@ -177,6 +172,9 @@ class InferenceOrchestrator: daemon = True, ) self._proc.start() + from utils.process_lifetime import adopt_pid + + adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep) logger.info("Inference subprocess started (pid=%s)", self._proc.pid) def _cancel_generation(self) -> None: @@ -389,6 +387,108 @@ class InferenceOrchestrator: return logger.warning("Timed out waiting for gen_done after cancel") + # ------------------------------------------------------------------ + # Generation command + token-stream helpers (shared by all paths) + # ------------------------------------------------------------------ + + def _build_generate_cmd( + self, + request_id: str, + image_b64: Optional[str], + *, + messages: list = None, + system_prompt: str = "", + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 256, + repetition_penalty: float = 1.0, + use_adapter = None, + tools: Optional[list] = None, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + ) -> dict: + """Build the 'generate' command shared by the locked and dispatched paths.""" + cmd = { + "type": "generate", + "request_id": request_id, + "messages": messages or [], + "system_prompt": system_prompt, + "image_base64": image_b64, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + # Only forward template kwargs the caller set, for older worker compat. + if use_adapter is not None: + cmd["use_adapter"] = use_adapter + if tools is not None: + cmd["tools"] = tools + if enable_thinking is not None: + cmd["enable_thinking"] = enable_thinking + if reasoning_effort is not None: + cmd["reasoning_effort"] = reasoning_effort + if preserve_thinking is not None: + cmd["preserve_thinking"] = preserve_thinking + return cmd + + def _consume_token_stream( + self, + read_one, + drain_on_cancel, + *, + crash_context: str, + cancel_event = None, + stats_holder: Optional[dict] = None, + read_timeout: float = 30.0, + ) -> Generator[str, None, None]: + """Yield tokens from a response stream until gen_done/gen_error. + + ``read_one(timeout)`` returns the next response (or None on timeout) and + owns the queue choice — the shared resp_queue under _gen_lock, or a + per-request mailbox on the dispatcher path — so this loop stays agnostic + of which queue is read. On cancel, ``drain_on_cancel()`` consumes the + cancel ack from that same source so stale events don't leak into the + next request. + """ + while True: + resp = read_one(read_timeout) + if resp is None: + # Check subprocess health + if not self._ensure_subprocess_alive(): + yield f"Error: {self._subprocess_crash_message(crash_context)}" + return + continue + + rtype = resp.get("type", "") + if rtype == "status": + continue + # Subprocess-level error (no request_id); request-scoped failures + # arrive as gen_error below. + if rtype == "error" and not resp.get("request_id"): + yield f"Error: {resp.get('error', 'Unknown error')}" + return + + if rtype == "token": + # Cancel from route (e.g. SSE connection closed). + if cancel_event is not None and cancel_event.is_set(): + self._cancel_generation() + drain_on_cancel() + return + yield resp.get("text", "") + elif rtype == "gen_done": + if stats_holder is not None: + stats_holder["stats"] = resp.get("stats") + return + elif rtype == "gen_error": + yield f"Error: {resp.get('error', 'Unknown error')}" + return + # ------------------------------------------------------------------ # Dispatcher — per-request mailbox routing for compare mode # ------------------------------------------------------------------ @@ -451,13 +551,12 @@ class InferenceOrchestrator: continue # No matching mailbox (a _gen_lock reader or orphaned). Can't - # un-get from mp.Queue, so just log. - if rtype not in ("status",): - logger.debug( - "Dispatcher: no mailbox for request_id=%s type=%s, dropping", - rid, - rtype, - ) + # un-get from mp.Queue, so just log. (status was handled above.) + logger.debug( + "Dispatcher: no mailbox for request_id=%s type=%s, dropping", + rid, + rtype, + ) def _generate_dispatched( self, @@ -502,30 +601,23 @@ class InferenceOrchestrator: if image is not None: image_b64 = self._pil_to_base64(image) - cmd = { - "type": "generate", - "request_id": request_id, - "messages": messages or [], - "system_prompt": system_prompt, - "image_base64": image_b64, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "max_new_tokens": max_new_tokens, - "repetition_penalty": repetition_penalty, - } - - if use_adapter is not None: - cmd["use_adapter"] = use_adapter - if tools is not None: - cmd["tools"] = tools - if enable_thinking is not None: - cmd["enable_thinking"] = enable_thinking - if reasoning_effort is not None: - cmd["reasoning_effort"] = reasoning_effort - if preserve_thinking is not None: - cmd["preserve_thinking"] = preserve_thinking + cmd = self._build_generate_cmd( + request_id, + image_b64, + messages = messages, + system_prompt = system_prompt, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + max_new_tokens = max_new_tokens, + repetition_penalty = repetition_penalty, + use_adapter = use_adapter, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) # Create mailbox BEFORE sending command mailbox: queue.Queue = queue.Queue() @@ -540,36 +632,22 @@ class InferenceOrchestrator: yield f"Error: {exc}" return - # Read tokens from our private mailbox + def read_mailbox(timeout): + try: + return mailbox.get(timeout = timeout) + except queue.Empty: + return None + + # Read tokens from our private mailbox (the dispatcher owns resp_queue). try: - while True: - try: - resp = mailbox.get(timeout = _DISPATCH_READ_TIMEOUT) - except queue.Empty: - # Timeout — check subprocess health - if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message('generation')}" - return - continue - - rtype = resp.get("type", "") - - if rtype == "token": - # Cancel from route (e.g. SSE connection closed) - if cancel_event is not None and cancel_event.is_set(): - self._cancel_generation() - self._drain_mailbox(mailbox, timeout = 5.0) - return - yield resp.get("text", "") - - elif rtype == "gen_done": - if stats_holder is not None: - stats_holder["stats"] = resp.get("stats") - return - - elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" - return + yield from self._consume_token_stream( + read_mailbox, + lambda: self._drain_mailbox(mailbox, timeout = 5.0), + crash_context = "generation", + cancel_event = cancel_event, + stats_holder = stats_holder, + read_timeout = _DISPATCH_READ_TIMEOUT, + ) finally: with self._mailbox_lock: self._mailboxes.pop(request_id, None) @@ -636,6 +714,7 @@ class InferenceOrchestrator: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, ) -> bool: """Load a model for inference. @@ -659,6 +738,7 @@ class InferenceOrchestrator: "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, + "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "gpu_ids": gpu_ids, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( @@ -716,7 +796,6 @@ class InferenceOrchestrator: ) if resp.get("success"): - self._current_transformers_major = needed_major model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) self.models[self.active_model_name] = { @@ -738,7 +817,8 @@ class InferenceOrchestrator: logger.info("Model '%s' loaded successfully in subprocess", model_name) return True else: - error = resp.get("error", "Failed to load model") + # Worker reports failures (consent gate included) under "message". + error = resp.get("message") or resp.get("error") or "Failed to load model" self.loading_models.discard(model_name) self.active_model_name = None self.models.clear() @@ -862,6 +942,7 @@ class InferenceOrchestrator: session_id: Optional[str] = None, rag_scope: Optional[dict] = None, confirm_tool_calls: bool = False, + bypass_permissions: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, **_unused, @@ -924,6 +1005,7 @@ class InferenceOrchestrator: session_id = session_id, rag_scope = rag_scope, confirm_tool_calls = confirm_tool_calls, + bypass_permissions = bypass_permissions, ) def generate_with_adapter_control( @@ -982,127 +1064,42 @@ class InferenceOrchestrator: self._wait_dispatcher_idle() # Serialize generation: two concurrent readers on resp_queue would - # consume and drop each other's token events. + # consume and drop each other's token events. Hold _gen_lock across the + # cmd build + send + whole stream so we stay the sole resp_queue reader. with self._gen_lock: - yield from self._generate_locked( + request_id = str(uuid.uuid4()) + image_b64 = self._pil_to_base64(image) if image is not None else None + cmd = self._build_generate_cmd( + request_id, + image_b64, messages = messages, system_prompt = system_prompt, - image = image, temperature = temperature, top_p = top_p, top_k = top_k, min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, - cancel_event = cancel_event, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, - stats_holder = stats_holder, ) - def _generate_locked( - self, - messages: list = None, - system_prompt: str = "", - image = None, - temperature: float = 0.7, - top_p: float = 0.9, - top_k: int = 40, - min_p: float = 0.0, - max_new_tokens: int = 256, - repetition_penalty: float = 1.0, - cancel_event = None, - use_adapter = None, - tools: Optional[list] = None, - enable_thinking: Optional[bool] = None, - reasoning_effort: Optional[str] = None, - preserve_thinking: Optional[bool] = None, - stats_holder: Optional[dict] = None, - ) -> Generator[str, None, None]: - """Actual generation logic — must be called under _gen_lock.""" - request_id = str(uuid.uuid4()) - - # Convert PIL Image to base64 if needed - image_b64 = None - if image is not None: - image_b64 = self._pil_to_base64(image) - - cmd = { - "type": "generate", - "request_id": request_id, - "messages": messages or [], - "system_prompt": system_prompt, - "image_base64": image_b64, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "max_new_tokens": max_new_tokens, - "repetition_penalty": repetition_penalty, - } - - if use_adapter is not None: - cmd["use_adapter"] = use_adapter - # Only forward template kwargs the caller set, for older worker compat. - if tools is not None: - cmd["tools"] = tools - if enable_thinking is not None: - cmd["enable_thinking"] = enable_thinking - if reasoning_effort is not None: - cmd["reasoning_effort"] = reasoning_effort - if preserve_thinking is not None: - cmd["preserve_thinking"] = preserve_thinking - - try: - self._send_cmd(cmd) - except RuntimeError as exc: - yield f"Error: {exc}" - return - - # We are the only resp_queue reader (under _gen_lock). - while True: - resp = self._read_resp(timeout = 30.0) - - if resp is None: - # Check subprocess health - if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message('generation')}" - return - continue - - rtype = resp.get("type", "") - - # Status messages — skip - if rtype == "status": - continue - - # Error without request_id = subprocess-level error - resp_rid = resp.get("request_id") - if rtype == "error" and not resp_rid: - yield f"Error: {resp.get('error', 'Unknown error')}" + try: + self._send_cmd(cmd) + except RuntimeError as exc: + yield f"Error: {exc}" return - if rtype == "token": - # Cancel from route (e.g. SSE connection closed) - if cancel_event is not None and cancel_event.is_set(): - self._cancel_generation() - # Wait for the cancel ack so stale events don't leak into - # the next request. - self._drain_until_gen_done(timeout = 5.0) - return - yield resp.get("text", "") - - elif rtype == "gen_done": - if stats_holder is not None: - stats_holder["stats"] = resp.get("stats") - return - - elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" - return + yield from self._consume_token_stream( + self._read_resp, + lambda: self._drain_until_gen_done(timeout = 5.0), + crash_context = "generation", + cancel_event = cancel_event, + stats_holder = stats_holder, + ) def reset_generation_state(self): """Cancel any ongoing generation and reset state.""" @@ -1138,8 +1135,6 @@ class InferenceOrchestrator: if not self.active_model_name: raise RuntimeError("No active model") - import uuid - request_id = str(uuid.uuid4()) cmd = { @@ -1252,8 +1247,6 @@ class InferenceOrchestrator: return with self._gen_lock: - import uuid - request_id = str(uuid.uuid4()) # numpy array -> list for mp.Queue serialization @@ -1282,38 +1275,12 @@ class InferenceOrchestrator: yield f"Error: {exc}" return - # Yield tokens — same pattern as _generate_locked - while True: - resp = self._read_resp(timeout = 30.0) - - if resp is None: - if not self._ensure_subprocess_alive(): - yield ("Error: " + self._subprocess_crash_message("audio input generation")) - return - continue - - rtype = resp.get("type", "") - - if rtype == "status": - continue - - if rtype == "error" and not resp.get("request_id"): - yield f"Error: {resp.get('error', 'Unknown error')}" - return - - if rtype == "token": - if cancel_event is not None and cancel_event.is_set(): - self._cancel_generation() - self._drain_until_gen_done(timeout = 5.0) - return - yield resp.get("text", "") - - elif rtype == "gen_done": - return - - elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" - return + yield from self._consume_token_stream( + self._read_resp, + lambda: self._drain_until_gen_done(timeout = 5.0), + crash_context = "audio input generation", + cancel_event = cancel_event, + ) # ------------------------------------------------------------------ # Local helpers (no subprocess needed) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 3b6a393f3d..06b6cbe57f 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -154,6 +154,7 @@ def run_safetensors_tool_loop( session_id: Optional[str] = None, rag_scope: Optional[dict] = None, confirm_tool_calls: bool = False, + bypass_permissions: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -517,7 +518,9 @@ def run_safetensors_tool_loop( else: assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call()) - needs_confirm = bool(confirm_tool_calls) + # Bypass wins over the confirm gate at the loop level too, so a + # direct internal caller passing both flags never prompts. + needs_confirm = bool(confirm_tool_calls) and not bypass_permissions approval_id = new_approval_id() if needs_confirm else "" decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None start_event = decision.tool_start_event() @@ -575,6 +578,7 @@ def run_safetensors_tool_loop( timeout = eff_timeout, session_id = session_id, rag_scope = rag_scope, + disable_sandbox = bypass_permissions, ) except Exception as exc: logger.exception("Tool %s raised: %s", decision.tool_name, exc) diff --git a/studio/backend/core/inference/tensor_fallback.py b/studio/backend/core/inference/tensor_fallback.py index 73687165b8..3ceb1a268a 100644 --- a/studio/backend/core/inference/tensor_fallback.py +++ b/studio/backend/core/inference/tensor_fallback.py @@ -13,7 +13,7 @@ import logging from typing import Awaitable, Callable, Optional from core.inference.llama_server_args import ( - resolve_tensor_parallel, + _effective_tensor_parallel, strip_split_mode_only, ) @@ -34,18 +34,20 @@ async def load_with_tensor_fallback( True on success; it *raises* on a hard crash (llama-server aborts on some archs / older builds), which is treated the same as a False return. - Tensor mode can be requested by the toggle or by a ``--split-mode tensor`` - in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether - tensor mode is actually engaged, and it strips ``--split-mode`` from the - extras so the layer retry can't relaunch the same failing tensor load. A - non-tensor load keeps its original contract and propagates exceptions. + Tensor mode can be requested by the toggle, by a ``--split-mode tensor`` in + ``extra_args`` (an allowed shadow flag), or by an inherited + ``LLAMA_ARG_SPLIT_MODE=tensor`` env (load_model engages it the same way), so + the retry is keyed on whether tensor mode is actually engaged, and it forces + ``--split-mode layer`` on the retry so neither leftover extras nor the + inherited tensor env can relaunch the same failing tensor load. A non-tensor + load keeps its original contract and propagates exceptions. ``cancelled()`` distinguishes a real tensor-start failure from a user cancellation: ``attempt_load`` also returns False when the load was cancelled, so without this the helper would restart a load the user just cancelled. """ - tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor) + tensor_requested = _effective_tensor_parallel(extra_args, requested_tensor) try: success = await attempt_load(requested_tensor, extra_args) except Exception as exc: @@ -67,4 +69,8 @@ async def load_with_tensor_fallback( "(this model may not support tensor parallelism)", label, ) - return await attempt_load(False, strip_split_mode_only(extra_args)) + # Force --split-mode layer (CLI wins over env) so neither leftover extras nor + # an inherited LLAMA_ARG_SPLIT_MODE=tensor can re-engage tensor and re-crash + # the retry; load_model and the child both honor the explicit layer override. + layer_extras = strip_split_mode_only(extra_args) or [] + return await attempt_load(False, [*layer_extras, "--split-mode", "layer"]) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 43c9610282..6960310018 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -316,6 +316,196 @@ def _build_safe_env(workdir: str) -> dict[str, str]: return env +# Credential env vars dropped even in bypass mode so tool code cannot read the +# operator's keys. Over-strips on purpose (a benign var is harmless to lose). +_BYPASS_ENV_SECRET_NAMES = frozenset( + { + "HF_TOKEN", + "HF_HUB_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "HUGGINGFACE_TOKEN", + "HUGGINGFACEHUB_API_TOKEN", + "WANDB_API_KEY", + "GH_TOKEN", + "GITHUB_TOKEN", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GEMINI_API_KEY", + "GOOGLE_API_KEY", + "GROQ_API_KEY", + "OPENROUTER_API_KEY", + "REPLICATE_API_TOKEN", + "COHERE_API_KEY", + "MISTRAL_API_KEY", + "NGC_API_KEY", + "KAGGLE_KEY", + "MYSQL_PWD", # exact name: markers use PASSWD, not PWD (PWD is the cwd var) + "LD_PRELOAD", + # Auth brokers / capability handles: not secrets by value, but they + # hand the child the operator's live agent (ssh/gpg), kube config, or + # docker daemon. Names are listed because there is no value signal to + # key off. URL config vars (HTTP_PROXY, PIP_INDEX_URL, DATABASE_URL, + # ...) are intentionally NOT name-listed: a benign proxy/index without + # credentials must keep working in bypass mode, while a credentialed + # value is dropped by _is_secret_env_value() regardless of its name. + "SSH_AUTH_SOCK", + "SSH_AGENT_PID", + "GPG_AGENT_INFO", + "GNUPGHOME", + "KUBECONFIG", + "DOCKER_HOST", + } +) +_BYPASS_ENV_SECRET_PREFIXES = ("AWS_", "AZURE_", "GOOGLE_", "GCP_", "GCLOUD_", "DYLD_") +_BYPASS_ENV_SECRET_MARKERS = ( + "TOKEN", + "API_KEY", + "APIKEY", + "SECRET", + "PASSWORD", + "PASSWD", + "CREDENTIAL", + "PRIVATE_KEY", + "AUTH", # e.g. NPM_CONFIG__AUTH (npm _auth), REDISCLI_AUTH + # Azure App Service connection strings: SQLCONNSTR_/CUSTOMCONNSTR_/... and + # WEBSITE_CONTENTAZUREFILECONNECTIONSTRING carry DB/storage credentials. + "CONNSTR", + "CONNECTIONSTRING", +) +# Non-secret hardening flags that match a secret prefix/marker but must be KEPT +# so bypass mode does not silently undo an operator's opt-out. AWS_EC2_METADATA_ +# DISABLED tells the AWS SDK/CLI not to pull instance-role creds from IMDS; +# dropping it would re-open that path for a bypassed tool. +_BYPASS_ENV_KEEP_NAMES = frozenset( + { + "AWS_EC2_METADATA_DISABLED", + "AWS_EC2_METADATA_V1_DISABLED", + } +) +# Matches a URL that embeds userinfo before the host, covering both +# "scheme://user:pass@host" and token-only "scheme://token@host" (and +# percent-encoded variants). The userinfo must precede the first '/', so an '@' +# in a path or query does not false-positive. Used to scrub credential-bearing +# URL values regardless of the variable's name. +_URL_USERINFO_RE = re.compile(r"://[^/\s@]+@") +# Connection-string credential fields (ADO.NET / Azure storage / Service Bus): +# "...;Password=...", "...;AccountKey=...", "...;SharedAccessKey=...". Catches +# credential-bearing values whose names dodge the name classifier. "accesskey" +# also covers Shared/Secret AccessKey via substring; the Name fields (e.g. +# SharedAccessKeyName=) do not match since "=" must follow the keyword. +_SECRET_VALUE_RE = re.compile(r"(?i)(?:password|pwd|accountkey|accesskey)\s*=\s*[^\s;]") + +# Names that hold no secret value but point SDKs at the operator's real +# home/cache/config (cached tokens, cred files), defeating the HOME repoint. +# Startup always sets HF_HOME (-> $HF_HOME/token), so this is the live leak. +# Dropped in bypass mode so tools fall back to the empty repointed HOME. +_BYPASS_ENV_CRED_LOCATION_NAMES = frozenset( + { + # HF cache roots (token lives under $HF_HOME/token) + "HF_HOME", + "HF_HUB_CACHE", + "HUGGINGFACE_HUB_CACHE", + "HF_XET_CACHE", + "TRANSFORMERS_CACHE", + "HF_DATASETS_CACHE", + "HF_ASSETS_CACHE", + # XDG base dirs (resolved before $HOME) + "XDG_CONFIG_HOME", + "XDG_CACHE_HOME", + "XDG_DATA_HOME", + # explicit cred/config file pointers honoured before $HOME + "NETRC", + "PGPASSFILE", + "BOTO_CONFIG", + "PIP_CONFIG_FILE", + "CLOUDSDK_CONFIG", + "KAGGLE_CONFIG_DIR", + "DOCKER_CONFIG", + "WANDB_DIR", + "WANDB_CONFIG_DIR", + "WANDB_CACHE_DIR", + # package-manager / git / cloud config pointers to real cred files + "NPM_CONFIG_USERCONFIG", + "NPM_CONFIG_GLOBALCONFIG", + "YARN_RC_FILENAME", + "GIT_CONFIG_GLOBAL", + "GIT_CONFIG_SYSTEM", + "CARGO_HOME", + "RCLONE_CONFIG", + # auth-helper scripts that hand creds to git/ssh + "GIT_ASKPASS", + "SSH_ASKPASS", + # shell startup hook: bash -c sources $BASH_ENV (can re-export secrets) + "BASH_ENV", + # Windows: HOMEDRIVE+HOMEPATH compose a home that bypasses HOME + "HOMEDRIVE", + "HOMEPATH", + } +) +# Windows profile dirs SDKs read creds under; repointed (not dropped) since +# callers expect them present. +_BYPASS_ENV_WINDOWS_PROFILE_VARS = ("USERPROFILE", "APPDATA", "LOCALAPPDATA") + + +def _is_secret_env_name(name: str) -> bool: + """True if an env var name looks like it carries a credential.""" + upper = name.upper() + if upper in _BYPASS_ENV_KEEP_NAMES: + return False # non-secret hardening flag; keep it + if upper in _BYPASS_ENV_SECRET_NAMES: + return True + if any(upper.startswith(p) for p in _BYPASS_ENV_SECRET_PREFIXES): + return True + return any(marker in upper for marker in _BYPASS_ENV_SECRET_MARKERS) + + +def _is_cred_location_env_name(name: str) -> bool: + """True for vars that point SDKs at the real home/cache/config (cached creds).""" + return name.upper() in _BYPASS_ENV_CRED_LOCATION_NAMES + + +def _is_secret_env_value(value: str) -> bool: + """True if a value embeds credentials regardless of its name. + + Catches URL userinfo (``scheme://user:token@host`` in DATABASE_URL / + PIP_INDEX_URL / HTTP_PROXY) and connection-string credential fields + (``...;Password=...`` / ``...;AccountKey=...``) whose names dodge the name + classifier. + """ + if not value: + return False + return _URL_USERINFO_RE.search(value) is not None or _SECRET_VALUE_RE.search(value) is not None + + +def _build_bypass_env(workdir: str) -> dict[str, str]: + """Env for bypass exec: full host env (unrestricted) minus credential vars, + with HOME/TMPDIR repointed at the workdir so SDKs cannot read cached creds. + + Note: stripping the child env is necessary but not sufficient on its own - + a same-UID child can still read the parent's environment via procfs, so + callers also harden the parent (see _harden_parent_against_proc_env_leak). + """ + env = { + k: v + for k, v in os.environ.items() + if not _is_secret_env_name(k) + and not _is_secret_env_value(v) + and not _is_cred_location_env_name(k) + } + env["HOME"] = workdir + env["TMPDIR"] = workdir + # Windows tempfile / SDKs honour TEMP/TMP, not TMPDIR; repoint all three so + # the bypassed tool writes under the per-session sandbox dir on every OS. + env["TEMP"] = workdir + env["TMP"] = workdir + # Windows SDKs read creds under the profile dirs, not $HOME; repoint set + # ones to the workdir (HOMEDRIVE/HOMEPATH are dropped above). + for var in _BYPASS_ENV_WINDOWS_PROFILE_VARS: + if var in os.environ: + env[var] = workdir + return env + + def _sandbox_preexec(): """Best-effort sandbox setup for sandboxed subprocesses (modules are resolved at import time so the forked child runs no imports).""" @@ -377,6 +567,65 @@ def _sandbox_preexec(): pass +def _bypass_preexec(): + """Minimal pre-exec for bypass exec: os.setsid() only. + + Required, not a restriction: _kill_process_tree does killpg(getpgid(child)), + so without a new session a timeout/cancel would kill the Studio server too. + """ + try: + os.setsid() + except OSError: + pass + + +# Hardening the Studio parent is done once (PR_SET_DUMPABLE is process-global +# and sticky); guarded so repeated bypass calls do not re-issue the prctl. +_parent_proc_hardened = False + + +def _harden_parent_against_proc_env_leak() -> bool: + """Make the Studio process's /proc//environ unreadable to its children. + + Stripping the child env is not enough on Linux: a bypassed same-UID child + runs unsandboxed and can read /proc//environ to recover the + tool-executing process's *unfiltered* secrets (HF_TOKEN, cloud keys, ...). + Clearing the dumpable flag (PR_SET_DUMPABLE=0) reparents this process's + /proc entries to root, so a same-UID child can no longer read its environ. + + Returns True when the process is hardened or hardening is unnecessary (no + /proc leak off Linux), and False when it is needed but could not be applied + (e.g. prctl denied by a seccomp policy). Callers must fail closed - refuse + the unsandboxed exec - when this returns False, rather than running with the + parent environ still readable. + + Scope: this closes the direct parent read (the demonstrated leak). It is a + mitigation, not a full boundary - a bypassed tool is unsandboxed by design, + so it can still walk /proc to a same-UID *ancestor* (e.g. the launching + shell) or read on-disk credentials by absolute path. Complete isolation + needs a separate uid / PID+mount namespace, which is out of scope here; the + UI already warns the mode is dangerous. Applied lazily on first bypass exec + so non-bypass operation is unchanged. + """ + global _parent_proc_hardened + if _parent_proc_hardened: + return True + if sys.platform != "linux": + return True # no /proc//environ same-UID leak to close + if _libc is None: + return False # on Linux but cannot issue prctl -> cannot harden + try: + # prctl(PR_SET_DUMPABLE=4, SUID_DUMP_DISABLE=0). ctypes returns the + # syscall result (-1 on failure) and does NOT raise, so check it. + ret = _libc.prctl(4, 0, 0, 0, 0) + except (OSError, AttributeError): + return False + if ret != 0: + return False + _parent_proc_hardened = True + return True + + def _get_shell_cmd(command: str) -> list[str]: """Return the platform-appropriate shell invocation for a command string.""" if sys.platform == "win32": @@ -524,10 +773,10 @@ RENDER_HTML_TOOL = { "function": { "name": "render_html", "description": ( - "Render a self-contained HTML/CSS/JavaScript artifact for the user. " + "Render a self-contained HTML/CSS/JavaScript canvas for the user. " "Call this at most once per assistant response unless the user " "explicitly asks for changes in that response. Future user requests " - "for new artifacts may call render_html once. Put the entire document " + "for new canvases may call render_html once. Put the entire document " "in code, including any CSS in