Merge branch 'main' into pip
This commit is contained in:
Daniel Han
commit
9cc539c1b4
235 changed files with 28186 additions and 2406 deletions
45
.github/workflows/release-desktop.yml
vendored
45
.github/workflows/release-desktop.yml
vendored
|
|
@ -54,10 +54,43 @@ jobs:
|
|||
with:
|
||||
node-version: 24
|
||||
|
||||
- name: Install pinned Tauri CLI
|
||||
run: npm install --save-dev --prefix studio @tauri-apps/cli@2.10.1
|
||||
|
||||
- name: Verify pinned Tauri CLI
|
||||
shell: bash
|
||||
run: |
|
||||
out="$(npx --prefix studio tauri --version)"
|
||||
echo "$out"
|
||||
if [ "$out" != "tauri-cli 2.10.1" ]; then
|
||||
echo "Expected tauri-cli 2.10.1, got $out" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
- name: Install frontend dependencies
|
||||
working-directory: studio/frontend
|
||||
run: npm install
|
||||
|
||||
- name: Verify backend package is published
|
||||
shell: bash
|
||||
run: |
|
||||
node <<'JS'
|
||||
const { readFileSync } = require('node:fs');
|
||||
|
||||
(async () => {
|
||||
const cargo = readFileSync('studio/src-tauri/Cargo.toml', 'utf8');
|
||||
const match = cargo.match(/^version\s*=\s*"([^"]+)"/m);
|
||||
if (!match) throw new Error('Could not read desktop app version');
|
||||
|
||||
const appVersion = match[1];
|
||||
const response = await fetch(`https://pypi.org/pypi/unsloth/${appVersion}/json`);
|
||||
if (!response.ok) {
|
||||
const message = 'Publish unsloth=={app_version} to PyPI before the desktop release';
|
||||
throw new Error(`${message.replace('{app_version}', appVersion)} (HTTP ${response.status})`);
|
||||
}
|
||||
})();
|
||||
JS
|
||||
|
||||
# ── Rust ──
|
||||
- name: Install Rust stable
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
|
@ -105,13 +138,14 @@ jobs:
|
|||
# ── Linux: build + sign + upload ──
|
||||
- name: Build Linux app
|
||||
if: matrix.platform == 'ubuntu-22.04'
|
||||
uses: tauri-apps/tauri-action@fce9c6108b31ea247710505d3aaaa893ee6768d4
|
||||
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
|
||||
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
|
||||
with:
|
||||
projectPath: studio
|
||||
tauriScript: npx --prefix . tauri
|
||||
tagName: desktop-v__VERSION__
|
||||
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
|
||||
releaseBody: |
|
||||
|
|
@ -123,6 +157,7 @@ jobs:
|
|||
|
||||
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
|
||||
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
|
||||
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
|
||||
releaseDraft: ${{ inputs.draft }}
|
||||
prerelease: false
|
||||
args: -v ${{ matrix.args }}
|
||||
|
|
@ -130,7 +165,7 @@ jobs:
|
|||
# ── macOS: build + sign + notarize + upload ──
|
||||
- name: Build macOS app
|
||||
if: matrix.platform == 'macos-latest'
|
||||
uses: tauri-apps/tauri-action@fce9c6108b31ea247710505d3aaaa893ee6768d4
|
||||
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
|
||||
|
|
@ -141,6 +176,7 @@ jobs:
|
|||
APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }}
|
||||
with:
|
||||
projectPath: studio
|
||||
tauriScript: npx --prefix . tauri
|
||||
tagName: desktop-v__VERSION__
|
||||
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
|
||||
releaseBody: |
|
||||
|
|
@ -152,6 +188,7 @@ jobs:
|
|||
|
||||
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
|
||||
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
|
||||
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
|
||||
releaseDraft: ${{ inputs.draft }}
|
||||
prerelease: false
|
||||
args: -v ${{ matrix.args }}
|
||||
|
|
@ -159,7 +196,7 @@ jobs:
|
|||
# ── Windows: build + sign + upload ──
|
||||
- name: Build Windows app
|
||||
if: matrix.platform == 'windows-latest'
|
||||
uses: tauri-apps/tauri-action@fce9c6108b31ea247710505d3aaaa893ee6768d4
|
||||
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
|
||||
|
|
@ -171,6 +208,7 @@ jobs:
|
|||
AZURE_CERTIFICATE_PROFILE_NAME: ${{ secrets.AZURE_CERTIFICATE_PROFILE_NAME }}
|
||||
with:
|
||||
projectPath: studio
|
||||
tauriScript: npx --prefix . tauri
|
||||
tagName: desktop-v__VERSION__
|
||||
releaseName: 'Unsloth Studio (Desktop) v__VERSION__'
|
||||
releaseBody: |
|
||||
|
|
@ -182,6 +220,7 @@ jobs:
|
|||
|
||||
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
|
||||
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
|
||||
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
|
||||
releaseDraft: ${{ inputs.draft }}
|
||||
prerelease: false
|
||||
args: -v ${{ matrix.args }}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
repos:
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.15.11
|
||||
rev: v0.15.12
|
||||
hooks:
|
||||
- id: ruff
|
||||
args:
|
||||
|
|
|
|||
15
README.md
15
README.md
|
|
@ -79,8 +79,9 @@ irm https://unsloth.ai/install.ps1 | iex
|
|||
|
||||
#### Launch
|
||||
```bash
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
unsloth studio -p 8888
|
||||
```
|
||||
> For cloud VMs or LAN access, add `-H 0.0.0.0` to bind on all interfaces.
|
||||
|
||||
#### Update
|
||||
To update, use the same install commands as above. Or run (does not work on Windows):
|
||||
|
|
@ -167,7 +168,7 @@ The below advanced instructions are for Unsloth Studio. For Unsloth Core advance
|
|||
git clone https://github.com/unslothai/unsloth
|
||||
cd unsloth
|
||||
./install.sh --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
unsloth studio -p 8888
|
||||
```
|
||||
Then to update :
|
||||
```bash
|
||||
|
|
@ -180,7 +181,7 @@ git clone https://github.com/unslothai/unsloth.git
|
|||
cd unsloth
|
||||
Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass
|
||||
.\install.ps1 --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
unsloth studio -p 8888
|
||||
```
|
||||
Then to update :
|
||||
```bash
|
||||
|
|
@ -193,11 +194,11 @@ git clone https://github.com/unslothai/unsloth
|
|||
cd unsloth
|
||||
git checkout nightly
|
||||
./install.sh --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
unsloth studio -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
unsloth studio -p 8888
|
||||
```
|
||||
|
||||
#### Nightly: Windows:
|
||||
|
|
@ -208,11 +209,11 @@ cd unsloth
|
|||
git checkout nightly
|
||||
Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass
|
||||
.\install.ps1 --local
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
unsloth studio -p 8888
|
||||
```
|
||||
Then to launch every time:
|
||||
```bash
|
||||
unsloth studio -H 0.0.0.0 -p 8888
|
||||
unsloth studio -p 8888
|
||||
```
|
||||
|
||||
#### Uninstall
|
||||
|
|
|
|||
325
install.ps1
325
install.ps1
|
|
@ -8,6 +8,79 @@ function Install-UnslothStudio {
|
|||
$ErrorActionPreference = "Stop"
|
||||
$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq "1")
|
||||
|
||||
# ── Tauri structured output ──
|
||||
function Write-TauriLog {
|
||||
param([string]$Tag, [string]$Message)
|
||||
if ($TauriMode) {
|
||||
Write-Host "[TAURI:$Tag] $Message"
|
||||
}
|
||||
}
|
||||
|
||||
function Format-TauriDiagBool {
|
||||
param([bool]$Value)
|
||||
if ($Value) { return "true" }
|
||||
return "false"
|
||||
}
|
||||
|
||||
function Get-TauriDiagArch {
|
||||
$arch = [string]$env:PROCESSOR_ARCHITECTURE
|
||||
if ([string]::IsNullOrWhiteSpace($arch)) {
|
||||
try { $arch = [System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture.ToString() } catch { $arch = "unknown" }
|
||||
}
|
||||
$arch = $arch.ToLowerInvariant()
|
||||
switch ($arch) {
|
||||
"amd64" { return "x86_64" }
|
||||
"x64" { return "x86_64" }
|
||||
"arm64" { return "arm64" }
|
||||
"x86" { return "x86" }
|
||||
default { return ($arch -replace '[^a-z0-9_.-]', '_') }
|
||||
}
|
||||
}
|
||||
|
||||
function Get-TauriTorchIndexFamily {
|
||||
param([string]$TorchIndexUrl)
|
||||
if ($SkipTorch) { return "none" }
|
||||
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" }
|
||||
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
|
||||
if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf }
|
||||
if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf }
|
||||
return "auto"
|
||||
}
|
||||
|
||||
function Get-TauriGpuBranch {
|
||||
param([string]$TorchIndexFamily)
|
||||
if ($SkipTorch) { return "no_torch" }
|
||||
if ($TorchIndexFamily -like "cu*") { return "cuda" }
|
||||
if ($TorchIndexFamily -like "rocm*") { return "rocm" }
|
||||
if ($TorchIndexFamily -eq "cpu") { return "cpu" }
|
||||
return "unknown"
|
||||
}
|
||||
|
||||
function Write-TauriDiag {
|
||||
param(
|
||||
[string]$GpuBranch = "unknown",
|
||||
[string]$TorchIndexFamily = "none",
|
||||
[string]$PythonVersionForDiag = $PythonVersion
|
||||
)
|
||||
if ([string]::IsNullOrWhiteSpace($PythonVersionForDiag)) { $PythonVersionForDiag = "unknown" }
|
||||
Write-TauriLog "DIAG" "diag_schema=1 platform=windows arch=$(Get-TauriDiagArch) python_version=$($PythonVersionForDiag.ToLowerInvariant()) skip_torch=$(Format-TauriDiagBool $SkipTorch) mac_intel=false gpu_branch=$GpuBranch torch_index_family=$TorchIndexFamily"
|
||||
}
|
||||
|
||||
function Exit-InstallFailure {
|
||||
param(
|
||||
[Parameter(Mandatory = $true)][string]$Message,
|
||||
[int]$Code = 1
|
||||
)
|
||||
if ($Code -eq 0) { $Code = 1 }
|
||||
Write-TauriLog "ERROR" $Message
|
||||
if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) {
|
||||
Restore-StudioVenvRollback
|
||||
}
|
||||
if ($TauriMode) {
|
||||
exit $Code
|
||||
}
|
||||
}
|
||||
|
||||
# ── Parse flags ──
|
||||
$StudioLocalInstall = $false
|
||||
$PackageName = "unsloth"
|
||||
|
|
@ -26,7 +99,7 @@ function Install-UnslothStudio {
|
|||
$i++
|
||||
if ($i -ge $argList.Count) {
|
||||
Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "--package requires an argument.")
|
||||
}
|
||||
$PackageName = $argList[$i]
|
||||
}
|
||||
|
|
@ -42,22 +115,14 @@ function Install-UnslothStudio {
|
|||
$RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path
|
||||
if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) {
|
||||
Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "--local must be run from the unsloth repo root")
|
||||
}
|
||||
}
|
||||
|
||||
# Validate --package to prevent injection into shell/Python commands
|
||||
if ($PackageName -notmatch '^[a-zA-Z0-9][a-zA-Z0-9._-]*$') {
|
||||
Write-Host "[ERROR] --package name contains invalid characters (allowed: a-z A-Z 0-9 . _ -)" -ForegroundColor Red
|
||||
return
|
||||
}
|
||||
|
||||
# ── Tauri structured output ──
|
||||
function Write-TauriLog {
|
||||
param([string]$Tag, [string]$Message)
|
||||
if ($TauriMode) {
|
||||
Write-Host "[TAURI:$Tag] $Message"
|
||||
}
|
||||
return (Exit-InstallFailure "--package name contains invalid characters")
|
||||
}
|
||||
|
||||
$PythonVersion = "3.13"
|
||||
|
|
@ -487,7 +552,7 @@ try {
|
|||
} catch {}
|
||||
exit 1
|
||||
}
|
||||
`$studioCommand = '& "' + `$studioExe + '" studio -H 0.0.0.0 -p ' + `$launchPort
|
||||
`$studioCommand = '& "' + `$studioExe + '" studio -p ' + `$launchPort
|
||||
`$launchArgs = @(
|
||||
'-NoExit',
|
||||
'-NoProfile',
|
||||
|
|
@ -630,7 +695,7 @@ shell.Run cmd, 0, False
|
|||
step "winget" "not available" "Red"
|
||||
substep "Install it from https://aka.ms/getwinget" "Yellow"
|
||||
substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow"
|
||||
return
|
||||
return (Exit-InstallFailure "winget is not available")
|
||||
}
|
||||
|
||||
# ── Helper: detect a working Python 3.11-3.13 on the system ──
|
||||
|
|
@ -749,9 +814,14 @@ shell.Run cmd, 0, False
|
|||
Write-Host " Please install Python $PythonVersion manually from https://www.python.org/downloads/" -ForegroundColor Yellow
|
||||
Write-Host " Make sure to check 'Add Python to PATH' during installation." -ForegroundColor Yellow
|
||||
Write-Host " Then re-run this installer." -ForegroundColor Yellow
|
||||
return
|
||||
return (Exit-InstallFailure "Python installation failed")
|
||||
}
|
||||
}
|
||||
$DiagPythonVersion = $PythonVersion
|
||||
if ($DetectedPython) { $DiagPythonVersion = $DetectedPython.Version }
|
||||
$InitialGpuBranch = "unknown"
|
||||
if ($SkipTorch) { $InitialGpuBranch = "no_torch" }
|
||||
Write-TauriDiag -GpuBranch $InitialGpuBranch -TorchIndexFamily "none" -PythonVersionForDiag $DiagPythonVersion
|
||||
|
||||
# ── Install uv if not present ──
|
||||
Write-TauriLog "STEP" "Installing uv package manager"
|
||||
|
|
@ -773,7 +843,7 @@ shell.Run cmd, 0, False
|
|||
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
|
||||
step "uv" "could not be installed" "Red"
|
||||
substep "Install it from https://docs.astral.sh/uv/" "Yellow"
|
||||
return
|
||||
return (Exit-InstallFailure "uv could not be installed")
|
||||
}
|
||||
|
||||
# ── Create venv (migrate old layout if possible, otherwise fresh) ──
|
||||
|
|
@ -786,11 +856,68 @@ shell.Run cmd, 0, False
|
|||
|
||||
$VenvPython = Join-Path $VenvDir "Scripts\python.exe"
|
||||
$_Migrated = $false
|
||||
$script:StudioVenvRollbackDir = $null
|
||||
$script:StudioVenvRollbackTarget = $VenvDir
|
||||
$script:StudioVenvRollbackActive = $false
|
||||
|
||||
function Start-StudioVenvRollback {
|
||||
param([Parameter(Mandatory = $true)][string]$ExistingDir)
|
||||
$stamp = Get-Date -Format "yyyyMMddHHmmss"
|
||||
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID"
|
||||
$suffix = 0
|
||||
while (Test-Path $candidate) {
|
||||
$suffix++
|
||||
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix"
|
||||
}
|
||||
Move-Item -Path $ExistingDir -Destination $candidate -ErrorAction Stop
|
||||
$script:StudioVenvRollbackDir = $candidate
|
||||
$script:StudioVenvRollbackTarget = $ExistingDir
|
||||
$script:StudioVenvRollbackActive = $true
|
||||
substep "previous environment preserved for rollback"
|
||||
}
|
||||
|
||||
function Restore-StudioVenvRollback {
|
||||
if (-not $script:StudioVenvRollbackActive) { return }
|
||||
$backup = $script:StudioVenvRollbackDir
|
||||
$target = $script:StudioVenvRollbackTarget
|
||||
if (-not $backup -or -not (Test-Path $backup)) {
|
||||
$script:StudioVenvRollbackActive = $false
|
||||
return
|
||||
}
|
||||
substep "restoring previous environment after failed install..." "Yellow"
|
||||
try {
|
||||
if (Test-Path $target) {
|
||||
Remove-Item -Recurse -Force $target -ErrorAction SilentlyContinue
|
||||
}
|
||||
Move-Item -Path $backup -Destination $target -Force -ErrorAction Stop
|
||||
substep "restored previous environment"
|
||||
$script:StudioVenvRollbackActive = $false
|
||||
$script:StudioVenvRollbackDir = $null
|
||||
} catch {
|
||||
Write-Host "[WARN] Could not restore previous environment from $backup to $target" -ForegroundColor Yellow
|
||||
Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
function Complete-StudioVenvRollback {
|
||||
if (-not $script:StudioVenvRollbackActive) { return }
|
||||
$backup = $script:StudioVenvRollbackDir
|
||||
if ($backup -and (Test-Path $backup)) {
|
||||
Remove-Item -Recurse -Force $backup -ErrorAction SilentlyContinue
|
||||
}
|
||||
$script:StudioVenvRollbackActive = $false
|
||||
$script:StudioVenvRollbackDir = $null
|
||||
}
|
||||
|
||||
if (Test-Path $VenvPython) {
|
||||
# New layout already exists -- nuke for fresh install
|
||||
substep "removing existing environment for fresh install..."
|
||||
Remove-Item -Recurse -Force $VenvDir
|
||||
# New layout already exists -- replace only after preserving rollback copy.
|
||||
substep "preserving existing environment for rollback..."
|
||||
try {
|
||||
Start-StudioVenvRollback -ExistingDir $VenvDir
|
||||
} catch {
|
||||
Write-Host "[ERROR] Could not prepare existing environment for reinstall: $($_.Exception.Message)" -ForegroundColor Red
|
||||
return (Exit-InstallFailure "Could not prepare existing environment for reinstall")
|
||||
}
|
||||
} elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) {
|
||||
# Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating
|
||||
$OldVenv = Join-Path $StudioHome ".venv"
|
||||
|
|
@ -799,18 +926,23 @@ shell.Run cmd, 0, False
|
|||
$prevEAP2 = $ErrorActionPreference
|
||||
$ErrorActionPreference = "Continue"
|
||||
try {
|
||||
& $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null
|
||||
$torchOk = ($LASTEXITCODE -eq 0)
|
||||
} catch { $torchOk = $false }
|
||||
if ($SkipTorch) {
|
||||
& $OldPy -c "import sys; print(sys.executable)" 2>$null | Out-Null
|
||||
} else {
|
||||
& $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null
|
||||
}
|
||||
$legacyOk = ($LASTEXITCODE -eq 0)
|
||||
} catch { $legacyOk = $false }
|
||||
$ErrorActionPreference = $prevEAP2
|
||||
if ($torchOk) {
|
||||
if ($legacyOk) {
|
||||
substep "legacy environment is healthy -- migrating..."
|
||||
Move-Item -Path $OldVenv -Destination $VenvDir -Force
|
||||
substep "moved .venv -> unsloth_studio"
|
||||
$_Migrated = $true
|
||||
} else {
|
||||
substep "legacy environment failed validation -- creating fresh environment" "Yellow"
|
||||
Remove-Item -Recurse -Force $OldVenv -ErrorAction SilentlyContinue
|
||||
$invalidVenv = Join-Path $StudioHome (".venv.invalid.{0}.{1}" -f (Get-Date -Format "yyyyMMddHHmmss"), $PID)
|
||||
Move-Item -Path $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
} elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) {
|
||||
# CWD-relative venv from old install.ps1 -- migrate to absolute path
|
||||
|
|
@ -826,9 +958,8 @@ shell.Run cmd, 0, False
|
|||
substep "$VenvDir"
|
||||
$venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" }
|
||||
if ($venvExit -ne 0) {
|
||||
Write-TauriLog "ERROR" "Failed to create virtual environment (exit code $venvExit)"
|
||||
Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "Failed to create virtual environment (exit code $venvExit)" $venvExit)
|
||||
}
|
||||
} else {
|
||||
step "venv" "using migrated environment"
|
||||
|
|
@ -886,6 +1017,9 @@ shell.Run cmd, 0, False
|
|||
return "$baseUrl/cu126"
|
||||
}
|
||||
$TorchIndexUrl = Get-TorchIndexUrl
|
||||
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
|
||||
$GpuBranch = Get-TauriGpuBranch $TorchIndexFamily
|
||||
Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version
|
||||
|
||||
# ── Print CPU-only hint when no GPU detected ──
|
||||
if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") {
|
||||
|
|
@ -934,7 +1068,7 @@ 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.4.7" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.8" unsloth-zoo }
|
||||
if ($baseInstallExit -eq 0) {
|
||||
$NoTorchReq = Find-NoTorchRuntimeFile
|
||||
if ($NoTorchReq) {
|
||||
|
|
@ -942,18 +1076,24 @@ shell.Run cmd, 0, False
|
|||
}
|
||||
}
|
||||
} else {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.7" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.8" unsloth-zoo }
|
||||
}
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
|
||||
}
|
||||
if ($StudioLocalInstall) {
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return
|
||||
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" }
|
||||
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)
|
||||
}
|
||||
}
|
||||
} elseif ($TorchIndexUrl) {
|
||||
|
|
@ -964,9 +1104,8 @@ shell.Run cmd, 0, False
|
|||
substep "installing PyTorch ($TorchIndexUrl)..."
|
||||
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
|
||||
if ($torchInstallExit -ne 0) {
|
||||
Write-TauriLog "ERROR" "Failed to install PyTorch (exit code $torchInstallExit)"
|
||||
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -975,7 +1114,7 @@ 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.4.7" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.4.8" unsloth-zoo }
|
||||
if ($baseInstallExit -eq 0) {
|
||||
$NoTorchReq = Find-NoTorchRuntimeFile
|
||||
if ($NoTorchReq) {
|
||||
|
|
@ -983,14 +1122,13 @@ shell.Run cmd, 0, False
|
|||
}
|
||||
}
|
||||
} elseif ($StudioLocalInstall) {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.7" unsloth-zoo }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.8" unsloth-zoo }
|
||||
} else {
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
|
||||
}
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-TauriLog "ERROR" "Failed to install unsloth (exit code $baseInstallExit)"
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
|
||||
}
|
||||
|
||||
if ($StudioLocalInstall) {
|
||||
|
|
@ -998,7 +1136,13 @@ shell.Run cmd, 0, False
|
|||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return
|
||||
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" }
|
||||
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 {
|
||||
|
|
@ -1006,53 +1150,72 @@ 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.4.7" --torch-backend=auto }
|
||||
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.4.8" --torch-backend=auto }
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
|
||||
}
|
||||
substep "overlaying local repo (editable)..."
|
||||
$overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps }
|
||||
if ($overlayExit -ne 0) {
|
||||
Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red
|
||||
return
|
||||
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" }
|
||||
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" }
|
||||
if ($baseInstallExit -ne 0) {
|
||||
Write-TauriLog "ERROR" "Failed to install unsloth (exit code $baseInstallExit)"
|
||||
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Hotfix: patch install_python_stack.py for Windows GUI stdout
|
||||
# The PyPI version crashes with OSError when stdout is piped from a GUI app.
|
||||
# Copy our fixed version (bundled by Tauri) over the installed one.
|
||||
# Remove this block once PyPI ships the fix from commit 18c5aae7.
|
||||
# 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__).
|
||||
# Source paths match the Tauri bundle layout in studio/src-tauri/tauri.conf.json,
|
||||
# which bundles install_python_stack.py at the bundle root next to install.ps1.
|
||||
if ($TauriMode) {
|
||||
$rawPath = if ($PSCommandPath) { $PSCommandPath } else { $MyInvocation.ScriptName }
|
||||
$scriptDir = Split-Path -Parent ($rawPath -replace '^\\\\\?\\', '')
|
||||
$fixedPy = Join-Path $scriptDir "install_python_stack.py"
|
||||
$target = Join-Path $VenvDir "Lib\site-packages\studio\install_python_stack.py"
|
||||
$sentinel = "# UNSLOTH_DESKTOP_HOTFIX_APPLIED_v1"
|
||||
$sentinelPattern = [regex]::Escape($sentinel)
|
||||
if ((Test-Path $fixedPy) -and (Test-Path $target)) {
|
||||
$installed = Get-Content $target -Raw
|
||||
if ($installed -notmatch $sentinelPattern) {
|
||||
Copy-Item $fixedPy $target -Force
|
||||
Add-Content -Path $target -Value "`n$sentinel"
|
||||
substep "patched install_python_stack.py (stdout fix)"
|
||||
} else {
|
||||
substep "install_python_stack.py already has stdout fix"
|
||||
if ($rawPath) {
|
||||
# Strip leading \\?\ extended-length prefix if the launcher passed one.
|
||||
$scriptDir = Split-Path -Parent ($rawPath -replace '^\\\\\?\\', '')
|
||||
$overlayMap = [ordered]@{
|
||||
"install_python_stack.py" = "Lib\site-packages\studio\install_python_stack.py"
|
||||
}
|
||||
foreach ($rel in $overlayMap.Keys) {
|
||||
$src = Join-Path $scriptDir $rel
|
||||
$dst = Join-Path $VenvDir $overlayMap[$rel]
|
||||
if (-not (Test-Path $src)) { continue }
|
||||
$dstParent = Split-Path -Parent $dst
|
||||
if (-not (Test-Path $dstParent)) {
|
||||
Write-Host "[WARN] Overlay target dir missing: $dstParent; studio setup may use stale bundled file" -ForegroundColor Yellow
|
||||
continue
|
||||
}
|
||||
try {
|
||||
if (-not (Test-Path $dst)) {
|
||||
# Backfill: target file missing but parent dir exists.
|
||||
Copy-Item $src $dst -Force
|
||||
substep ("backfilled bundled " + (Split-Path -Leaf $rel))
|
||||
} else {
|
||||
# Hash-compare so re-runs are no-ops when files already match.
|
||||
$srcHash = (Get-FileHash $src -Algorithm SHA256).Hash
|
||||
$dstHash = (Get-FileHash $dst -Algorithm SHA256).Hash
|
||||
if ($srcHash -ne $dstHash) {
|
||||
Copy-Item $src $dst -Force
|
||||
substep ("applied bundled " + (Split-Path -Leaf $rel))
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
Write-Host "[WARN] Could not overlay $($rel): $($_.Exception.Message); studio setup may use stale bundled file" -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
} elseif ((Test-Path $fixedPy) -and (Test-Path (Split-Path $target))) {
|
||||
Copy-Item $fixedPy $target -Force
|
||||
Add-Content -Path $target -Value "`n$sentinel"
|
||||
substep "patched install_python_stack.py (stdout fix)"
|
||||
} else {
|
||||
Write-Host "[WARN] Could not patch install_python_stack.py (bundled file or target dir missing)" -ForegroundColor Yellow
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -1063,12 +1226,11 @@ shell.Run cmd, 0, False
|
|||
step "setup" "running unsloth studio setup..."
|
||||
$UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe"
|
||||
if (-not (Test-Path $UnslothExe)) {
|
||||
Write-TauriLog "ERROR" "unsloth CLI was not installed correctly"
|
||||
Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red
|
||||
Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow
|
||||
Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow
|
||||
Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow
|
||||
return
|
||||
return (Exit-InstallFailure "unsloth CLI was not installed correctly")
|
||||
}
|
||||
# Tell setup.ps1 to skip base package installation (install.ps1 already did it)
|
||||
$env:SKIP_STUDIO_BASE = "1"
|
||||
|
|
@ -1090,12 +1252,16 @@ shell.Run cmd, 0, False
|
|||
# and bypass the fast-path version check from PR #4667.
|
||||
$studioArgs = @('studio', 'setup')
|
||||
if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
|
||||
& $UnslothExe @studioArgs
|
||||
$setupExit = $LASTEXITCODE
|
||||
$env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1"
|
||||
try {
|
||||
& $UnslothExe @studioArgs
|
||||
$setupExit = $LASTEXITCODE
|
||||
} finally {
|
||||
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
|
||||
}
|
||||
if ($setupExit -ne 0) {
|
||||
Write-TauriLog "ERROR" "unsloth studio setup failed (exit code $setupExit)"
|
||||
Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red
|
||||
return
|
||||
return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit)
|
||||
}
|
||||
|
||||
# ── Expose `unsloth` via a shim dir containing only unsloth.exe ──
|
||||
|
|
@ -1168,6 +1334,7 @@ shell.Run cmd, 0, False
|
|||
step "path" "added unsloth launcher to PATH"
|
||||
}
|
||||
Refresh-SessionPath # sync current session with registry
|
||||
Complete-StudioVenvRollback
|
||||
|
||||
# ── Tauri mode: done, skip shortcuts and auto-launch ──
|
||||
if ($TauriMode) {
|
||||
|
|
@ -1177,15 +1344,25 @@ shell.Run cmd, 0, False
|
|||
|
||||
New-StudioShortcuts -UnslothExePath $UnslothExe
|
||||
|
||||
# Launch studio automatically in interactive terminals;
|
||||
# in non-interactive environments (CI, Docker) just print instructions.
|
||||
# In interactive terminals, ask the user before starting Studio.
|
||||
# In non-interactive environments (CI, Docker) just print instructions.
|
||||
$IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
|
||||
if ($IsInteractive) {
|
||||
& $UnslothExe studio -H 0.0.0.0 -p 8888
|
||||
Write-Host ""
|
||||
$reply = Read-Host " Start Unsloth Studio now? [Y/n]"
|
||||
if ([string]::IsNullOrWhiteSpace($reply) -or $reply -match '^[Yy]') {
|
||||
& $UnslothExe studio -p 8888
|
||||
} else {
|
||||
step "launch" "to start later, run:"
|
||||
substep "unsloth studio -p 8888"
|
||||
substep "(add -H 0.0.0.0 to allow network / cloud access)"
|
||||
Write-Host ""
|
||||
}
|
||||
} else {
|
||||
step "launch" "manual commands:"
|
||||
substep "& `"$VenvDir\Scripts\Activate.ps1`""
|
||||
substep "unsloth studio -H 0.0.0.0 -p 8888"
|
||||
substep "unsloth studio -p 8888"
|
||||
substep "(add -H 0.0.0.0 to allow network / cloud access)"
|
||||
Write-Host ""
|
||||
}
|
||||
}
|
||||
|
|
|
|||
243
install.sh
243
install.sh
|
|
@ -162,9 +162,119 @@ tauri_log() {
|
|||
fi
|
||||
}
|
||||
|
||||
tauri_diag_marker() {
|
||||
_diag_gpu_branch="${1:-unknown}"
|
||||
_diag_torch_index_family="${2:-none}"
|
||||
tauri_log "DIAG" "diag_schema=1 platform=${OS:-unknown} arch=${_ARCH:-unknown} python_version=${PYTHON_VERSION:-unknown} skip_torch=${SKIP_TORCH:-false} mac_intel=${MAC_INTEL:-false} gpu_branch=${_diag_gpu_branch} torch_index_family=${_diag_torch_index_family}"
|
||||
}
|
||||
|
||||
_tauri_torch_index_family() {
|
||||
if [ "${SKIP_TORCH:-false}" = true ]; then
|
||||
echo "none"
|
||||
return
|
||||
fi
|
||||
_diag_url="${1:-}"
|
||||
case "$_diag_url" in
|
||||
*/cu118) echo "cu118" ;;
|
||||
*/cu124) echo "cu124" ;;
|
||||
*/cu126) echo "cu126" ;;
|
||||
*/cu128) echo "cu128" ;;
|
||||
*/cu130) echo "cu130" ;;
|
||||
*/cpu) echo "cpu" ;;
|
||||
*/rocm[0-9]*.[0-9]*)
|
||||
_diag_family=${_diag_url##*/}
|
||||
case "$_diag_family" in
|
||||
rocm[0-9]*.[0-9]*) echo "$_diag_family" ;;
|
||||
*) echo "auto" ;;
|
||||
esac ;;
|
||||
"") echo "none" ;;
|
||||
*) echo "auto" ;;
|
||||
esac
|
||||
}
|
||||
|
||||
_tauri_gpu_branch() {
|
||||
_diag_family="${1:-unknown}"
|
||||
_diag_radeon="${2:-false}"
|
||||
if [ "${SKIP_TORCH:-false}" = true ]; then
|
||||
echo "no_torch"
|
||||
return
|
||||
fi
|
||||
if [ "${OS:-}" = "macos" ]; then
|
||||
echo "mac"
|
||||
return
|
||||
fi
|
||||
case "$_diag_family" in
|
||||
cu*) echo "cuda" ;;
|
||||
rocm*)
|
||||
if [ "$_diag_radeon" = true ]; then
|
||||
echo "rocm_radeon"
|
||||
else
|
||||
echo "rocm"
|
||||
fi ;;
|
||||
radeon) echo "rocm_radeon" ;;
|
||||
cpu) echo "cpu" ;;
|
||||
none) echo "no_torch" ;;
|
||||
*) echo "unknown" ;;
|
||||
esac
|
||||
}
|
||||
|
||||
PYTHON_VERSION="" # resolved after platform detection
|
||||
STUDIO_HOME="$HOME/.unsloth/studio"
|
||||
VENV_DIR="$STUDIO_HOME/unsloth_studio"
|
||||
_VENV_ROLLBACK_DIR=""
|
||||
_VENV_ROLLBACK_TARGET="$VENV_DIR"
|
||||
_VENV_ROLLBACK_ACTIVE=false
|
||||
|
||||
_start_studio_venv_replacement() {
|
||||
_existing_dir="$1"
|
||||
_stamp=$(date +%Y%m%d%H%M%S 2>/dev/null || echo "time")
|
||||
_candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$"
|
||||
_suffix=0
|
||||
while [ -e "$_candidate" ]; do
|
||||
_suffix=$((_suffix + 1))
|
||||
_candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$.$_suffix"
|
||||
done
|
||||
mv "$_existing_dir" "$_candidate"
|
||||
_VENV_ROLLBACK_DIR="$_candidate"
|
||||
_VENV_ROLLBACK_TARGET="$_existing_dir"
|
||||
_VENV_ROLLBACK_ACTIVE=true
|
||||
substep "previous environment preserved for rollback"
|
||||
}
|
||||
|
||||
_restore_studio_venv_replacement() {
|
||||
[ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0
|
||||
[ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ] || {
|
||||
_VENV_ROLLBACK_ACTIVE=false
|
||||
return 0
|
||||
}
|
||||
substep "restoring previous environment after failed install..." "$C_WARN"
|
||||
rm -rf "$_VENV_ROLLBACK_TARGET"
|
||||
if mv "$_VENV_ROLLBACK_DIR" "$_VENV_ROLLBACK_TARGET"; then
|
||||
substep "restored previous environment"
|
||||
_VENV_ROLLBACK_ACTIVE=false
|
||||
_VENV_ROLLBACK_DIR=""
|
||||
else
|
||||
echo "⚠️ Could not restore previous environment from $_VENV_ROLLBACK_DIR to $_VENV_ROLLBACK_TARGET" >&2
|
||||
fi
|
||||
}
|
||||
|
||||
_commit_studio_venv_replacement() {
|
||||
[ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0
|
||||
if [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ]; then
|
||||
rm -rf "$_VENV_ROLLBACK_DIR" || true
|
||||
fi
|
||||
_VENV_ROLLBACK_ACTIVE=false
|
||||
_VENV_ROLLBACK_DIR=""
|
||||
}
|
||||
|
||||
_on_install_exit() {
|
||||
_status=$?
|
||||
if [ "$_status" -ne 0 ]; then
|
||||
_restore_studio_venv_replacement
|
||||
fi
|
||||
exit "$_status"
|
||||
}
|
||||
trap _on_install_exit EXIT
|
||||
|
||||
# ── Helper: download a URL to a file (supports curl and wget) ──
|
||||
download() {
|
||||
|
|
@ -512,11 +622,11 @@ if [ -t 1 ]; then
|
|||
) &
|
||||
# Clear traps so exec does not trigger _release_lock (the subshell owns it)
|
||||
trap - EXIT INT TERM
|
||||
exec "$UNSLOTH_EXE" studio -H 0.0.0.0 -p "$_launch_port"
|
||||
exec "$UNSLOTH_EXE" studio -p "$_launch_port"
|
||||
else
|
||||
# ── Background mode (no TTY) ──
|
||||
# Used by macOS .app and headless invocations.
|
||||
_launch_cmd=$(printf '%q ' "$UNSLOTH_EXE" studio -H 0.0.0.0 -p "$_launch_port")
|
||||
_launch_cmd=$(printf '%q ' "$UNSLOTH_EXE" studio -p "$_launch_port")
|
||||
_launch_cmd=${_launch_cmd% }
|
||||
_spawn_terminal "$_launch_cmd"
|
||||
|
||||
|
|
@ -828,6 +938,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then
|
|||
SKIP_TORCH=true
|
||||
fi
|
||||
|
||||
_TAURI_INITIAL_GPU_BRANCH="unknown"
|
||||
if [ "$SKIP_TORCH" = true ]; then
|
||||
_TAURI_INITIAL_GPU_BRANCH="no_torch"
|
||||
elif [ "$OS" = "macos" ]; then
|
||||
_TAURI_INITIAL_GPU_BRANCH="mac"
|
||||
fi
|
||||
tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none"
|
||||
|
||||
# ── Check system dependencies ──
|
||||
# cmake and git are needed by unsloth studio setup to build the GGUF inference
|
||||
# engine (llama.cpp). build-essential and libcurl-dev are also needed on Linux.
|
||||
|
|
@ -856,9 +974,7 @@ case "$OS" in
|
|||
fi
|
||||
command -v gcc >/dev/null 2>&1 || MISSING="$MISSING build-essential"
|
||||
# libcurl dev headers for llama.cpp HTTPS support
|
||||
if command -v dpkg >/dev/null 2>&1; then
|
||||
dpkg -s libcurl4-openssl-dev >/dev/null 2>&1 || MISSING="$MISSING libcurl4-openssl-dev"
|
||||
fi
|
||||
command -v curl-config >/dev/null 2>&1 || MISSING="$MISSING libcurl4-openssl-dev"
|
||||
;;
|
||||
esac
|
||||
|
||||
|
|
@ -883,9 +999,15 @@ if [ -n "$MISSING" ]; then
|
|||
if command -v apt-get >/dev/null 2>&1; then
|
||||
_smart_apt_install $MISSING
|
||||
else
|
||||
echo " apt-get is not available. Please install with your package manager:"
|
||||
echo " Automatic system package installation is supported on apt-based"
|
||||
echo " Linux distributions (Ubuntu/Debian) only. Please install the"
|
||||
echo " missing dependencies with your package manager, then re-run setup:"
|
||||
echo " $MISSING"
|
||||
echo " Then re-run Unsloth Studio setup."
|
||||
echo ""
|
||||
echo " Examples:"
|
||||
echo " Fedora/RHEL: sudo dnf install cmake git gcc gcc-c++ make libcurl-devel"
|
||||
echo " Arch: sudo pacman -S --needed cmake git base-devel curl"
|
||||
echo " openSUSE: sudo zypper install cmake git gcc gcc-c++ make libcurl-devel"
|
||||
exit 1
|
||||
fi
|
||||
;;
|
||||
|
|
@ -957,12 +1079,19 @@ mkdir -p "$STUDIO_HOME"
|
|||
_MIGRATED=false
|
||||
|
||||
if [ -x "$VENV_DIR/bin/python" ]; then
|
||||
# New layout already exists — nuke for fresh install
|
||||
rm -rf "$VENV_DIR"
|
||||
# New layout already exists — replace only after preserving rollback copy.
|
||||
substep "preserving existing environment for rollback..."
|
||||
_start_studio_venv_replacement "$VENV_DIR"
|
||||
elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then
|
||||
# Old layout exists — validate before migrating
|
||||
# Old layout exists — validate before migrating.
|
||||
# In no-torch mode, a missing torch package is expected; validate Python only.
|
||||
substep "found legacy Studio environment, validating..."
|
||||
if "$STUDIO_HOME/.venv/bin/python" -c "
|
||||
_legacy_ok=false
|
||||
if [ "$SKIP_TORCH" = true ]; then
|
||||
if "$STUDIO_HOME/.venv/bin/python" -c "import sys; print(sys.executable)" >/dev/null 2>&1; then
|
||||
_legacy_ok=true
|
||||
fi
|
||||
elif "$STUDIO_HOME/.venv/bin/python" -c "
|
||||
import torch
|
||||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
A = torch.ones((10, 10), device=device)
|
||||
|
|
@ -972,13 +1101,17 @@ D = A + B
|
|||
E = D @ C
|
||||
torch.testing.assert_close(torch.unique(E), torch.tensor((20,), device=E.device, dtype=E.dtype))
|
||||
" >/dev/null 2>&1; then
|
||||
_legacy_ok=true
|
||||
fi
|
||||
if [ "$_legacy_ok" = true ]; then
|
||||
echo "✅ Legacy environment is healthy — migrating..."
|
||||
mv "$STUDIO_HOME/.venv" "$VENV_DIR"
|
||||
echo " Moved ~/.unsloth/studio/.venv → $VENV_DIR"
|
||||
_MIGRATED=true
|
||||
else
|
||||
echo "⚠️ Legacy environment failed validation — creating fresh environment"
|
||||
rm -rf "$STUDIO_HOME/.venv"
|
||||
_invalid_venv="$STUDIO_HOME/.venv.invalid.$(date +%Y%m%d%H%M%S 2>/dev/null || echo time).$$"
|
||||
mv "$STUDIO_HOME/.venv" "$_invalid_venv" 2>/dev/null || true
|
||||
fi
|
||||
fi
|
||||
|
||||
|
|
@ -1308,6 +1441,12 @@ case "$TORCH_INDEX_URL" in
|
|||
fi
|
||||
;;
|
||||
esac
|
||||
_TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL")
|
||||
if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then
|
||||
_TAURI_TORCH_INDEX_FAMILY="radeon"
|
||||
fi
|
||||
_TAURI_GPU_BRANCH=$(_tauri_gpu_branch "$_TAURI_TORCH_INDEX_FAMILY" "$_amd_gpu_radeon")
|
||||
tauri_diag_marker "$_TAURI_GPU_BRANCH" "$_TAURI_TORCH_INDEX_FAMILY"
|
||||
|
||||
# ── Print CPU-only hint when no GPU detected ──
|
||||
case "$TORCH_INDEX_URL" in
|
||||
|
|
@ -1347,7 +1486,7 @@ if [ "$_MIGRATED" = true ]; then
|
|||
# to prevent transitive torch resolution.
|
||||
run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
|
||||
--reinstall-package unsloth --reinstall-package unsloth-zoo \
|
||||
"unsloth>=2026.4.7" unsloth-zoo
|
||||
"unsloth>=2026.4.8" unsloth-zoo
|
||||
_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"
|
||||
|
|
@ -1355,11 +1494,15 @@ if [ "$_MIGRATED" = true ]; then
|
|||
else
|
||||
run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
|
||||
--reinstall-package unsloth --reinstall-package unsloth-zoo \
|
||||
"unsloth>=2026.4.7" unsloth-zoo
|
||||
"unsloth>=2026.4.8" unsloth-zoo
|
||||
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" \
|
||||
--no-deps --reinstall-package unsloth-zoo \
|
||||
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
|
||||
fi
|
||||
# AMD ROCm: install bitsandbytes even in migrated environments so
|
||||
# existing ROCm installs gain the AMD bitsandbytes build without a
|
||||
|
|
@ -1519,7 +1662,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; 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 \
|
||||
--upgrade-package unsloth --upgrade-package unsloth-zoo \
|
||||
"unsloth>=2026.4.7" unsloth-zoo
|
||||
"unsloth>=2026.4.8" unsloth-zoo
|
||||
_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"
|
||||
|
|
@ -1527,12 +1670,20 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
|
|||
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" \
|
||||
--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.4.7" unsloth-zoo
|
||||
--upgrade-package unsloth "unsloth>=2026.4.8" unsloth-zoo
|
||||
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" \
|
||||
--no-deps --reinstall-package unsloth-zoo \
|
||||
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
|
||||
else
|
||||
run_install_cmd "install unsloth" uv pip install --python "$_VENV_PY" \
|
||||
--upgrade-package unsloth -- "$PACKAGE_NAME"
|
||||
|
|
@ -1558,9 +1709,13 @@ 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.4.7" --torch-backend=auto
|
||||
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.4.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" \
|
||||
--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"
|
||||
fi
|
||||
|
|
@ -1679,6 +1834,8 @@ if [ "$_SETUP_EXIT" -ne 0 ]; then
|
|||
exit "$_SETUP_EXIT"
|
||||
fi
|
||||
|
||||
_commit_studio_venv_replacement
|
||||
|
||||
# ── Tauri mode: done, skip shortcuts and auto-launch ──
|
||||
if [ "$TAURI_MODE" = true ]; then
|
||||
tauri_log "DONE" ""
|
||||
|
|
@ -1690,28 +1847,46 @@ printf " ${C_TITLE}%s${C_RST}\n" "Unsloth Studio installed!"
|
|||
printf " ${C_DIM}%s${C_RST}\n" "$RULE"
|
||||
echo ""
|
||||
|
||||
# Launch studio automatically in interactive terminals;
|
||||
# in non-interactive environments (Docker, CI, cloud-init) just print instructions.
|
||||
# In interactive terminals, ask the user before starting Studio.
|
||||
# In non-interactive environments (Docker, CI, cloud-init) just print instructions.
|
||||
if [ -t 1 ]; then
|
||||
step "launch" "starting Unsloth Studio..."
|
||||
"$VENV_DIR/bin/unsloth" studio -H 0.0.0.0 -p 8888
|
||||
_LAUNCH_EXIT=$?
|
||||
if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then
|
||||
echo ""
|
||||
echo "⚠️ Unsloth Studio failed to start after migration."
|
||||
echo " Your migrated environment may be incompatible."
|
||||
echo " To fix, remove the environment and reinstall:"
|
||||
echo ""
|
||||
echo " rm -rf $VENV_DIR"
|
||||
echo " curl -fsSL https://unsloth.ai/install.sh | sh"
|
||||
echo ""
|
||||
echo ""
|
||||
printf " Start Unsloth Studio now? [Y/n] "
|
||||
if [ -r /dev/tty ]; then
|
||||
read -r _reply </dev/tty || _reply="y"
|
||||
else
|
||||
_reply="y"
|
||||
fi
|
||||
exit "$_LAUNCH_EXIT"
|
||||
case "${_reply:-y}" in
|
||||
[Yy]*|"")
|
||||
step "launch" "starting Unsloth Studio..."
|
||||
"$VENV_DIR/bin/unsloth" studio -p 8888
|
||||
_LAUNCH_EXIT=$?
|
||||
if [ "$_LAUNCH_EXIT" -ne 0 ] && [ "$_MIGRATED" = true ]; then
|
||||
echo ""
|
||||
echo "⚠️ Unsloth Studio failed to start after migration."
|
||||
echo " Your migrated environment may be incompatible."
|
||||
echo " To fix, remove the environment and reinstall:"
|
||||
echo ""
|
||||
echo " rm -rf $VENV_DIR"
|
||||
echo " curl -fsSL https://unsloth.ai/install.sh | sh"
|
||||
echo ""
|
||||
fi
|
||||
exit "$_LAUNCH_EXIT"
|
||||
;;
|
||||
*)
|
||||
step "launch" "to start later, run:"
|
||||
substep "unsloth studio -p 8888"
|
||||
substep "(add -H 0.0.0.0 to allow network / cloud access)"
|
||||
echo ""
|
||||
;;
|
||||
esac
|
||||
else
|
||||
step "launch" "manual commands:"
|
||||
substep "unsloth studio -H 0.0.0.0 -p 8888"
|
||||
substep "unsloth studio -p 8888"
|
||||
substep "or activate env first:"
|
||||
substep "source ${VENV_DIR}/bin/activate"
|
||||
substep "unsloth studio -H 0.0.0.0 -p 8888"
|
||||
substep "unsloth studio -p 8888"
|
||||
substep "(add -H 0.0.0.0 to allow network / cloud access)"
|
||||
echo ""
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -64,7 +64,7 @@
|
|||
"id": "27e68f91"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh"
|
||||
"source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local"
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
|
|
|
|||
|
|
@ -146,10 +146,18 @@ def get_connection() -> sqlite3.Connection:
|
|||
created_at TEXT NOT NULL,
|
||||
last_used_at TEXT,
|
||||
expires_at TEXT,
|
||||
is_active INTEGER NOT NULL DEFAULT 1
|
||||
is_active INTEGER NOT NULL DEFAULT 1,
|
||||
is_internal INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
"""
|
||||
)
|
||||
api_key_columns = {
|
||||
row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")
|
||||
}
|
||||
if "is_internal" not in api_key_columns:
|
||||
conn.execute(
|
||||
"ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0"
|
||||
)
|
||||
conn.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS app_secrets (
|
||||
|
|
@ -592,11 +600,15 @@ def create_api_key(
|
|||
username: str,
|
||||
name: str,
|
||||
expires_at: Optional[str] = None,
|
||||
internal: bool = False,
|
||||
) -> Tuple[str, dict]:
|
||||
"""Create a new API key for *username*.
|
||||
|
||||
Returns ``(raw_key, row_dict)`` where *raw_key* is shown to the user
|
||||
exactly once. The database only stores the SHA-256 hash.
|
||||
exactly once. The database only stores the PBKDF2 hash.
|
||||
|
||||
Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe
|
||||
runs) that should not appear in user-facing key listings.
|
||||
"""
|
||||
raw_key = API_KEY_PREFIX + secrets.token_hex(16)
|
||||
key_hash = _pbkdf2_api_key(raw_key)
|
||||
|
|
@ -607,10 +619,18 @@ def create_api_key(
|
|||
try:
|
||||
conn.execute(
|
||||
"""
|
||||
INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(username, key_prefix, key_hash, name, now, expires_at),
|
||||
(
|
||||
username,
|
||||
key_prefix,
|
||||
key_hash,
|
||||
name,
|
||||
now,
|
||||
expires_at,
|
||||
1 if internal else 0,
|
||||
),
|
||||
)
|
||||
conn.commit()
|
||||
cur = conn.execute("SELECT * FROM api_keys WHERE key_hash = ?", (key_hash,))
|
||||
|
|
@ -620,19 +640,33 @@ def create_api_key(
|
|||
conn.close()
|
||||
|
||||
|
||||
def list_api_keys(username: str) -> list:
|
||||
"""Return all API keys for *username* (never exposes ``key_hash``)."""
|
||||
def list_api_keys(username: str, include_internal: bool = False) -> list:
|
||||
"""Return API keys for *username*. Internal workflow keys are hidden
|
||||
by default so they do not clutter user-facing UIs."""
|
||||
conn = get_connection()
|
||||
try:
|
||||
cur = conn.execute(
|
||||
"""
|
||||
SELECT id, username, key_prefix, name, created_at, last_used_at, expires_at, is_active
|
||||
FROM api_keys
|
||||
WHERE username = ?
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
(username,),
|
||||
)
|
||||
if include_internal:
|
||||
cur = conn.execute(
|
||||
"""
|
||||
SELECT id, username, key_prefix, name, created_at, last_used_at,
|
||||
expires_at, is_active, is_internal
|
||||
FROM api_keys
|
||||
WHERE username = ?
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
(username,),
|
||||
)
|
||||
else:
|
||||
cur = conn.execute(
|
||||
"""
|
||||
SELECT id, username, key_prefix, name, created_at, last_used_at,
|
||||
expires_at, is_active, is_internal
|
||||
FROM api_keys
|
||||
WHERE username = ? AND is_internal = 0
|
||||
ORDER BY created_at DESC
|
||||
""",
|
||||
(username,),
|
||||
)
|
||||
return [dict(row) for row in cur.fetchall()]
|
||||
finally:
|
||||
conn.close()
|
||||
|
|
@ -652,6 +686,24 @@ def revoke_api_key(username: str, key_id: int) -> bool:
|
|||
conn.close()
|
||||
|
||||
|
||||
def revoke_internal_api_key(key_id: int) -> bool:
|
||||
"""Revoke an internal workflow-minted key without requiring a username.
|
||||
|
||||
Used by the recipe runner to retire its sk-unsloth-* key once the job
|
||||
terminates, shrinking the window a leaked key could be abused.
|
||||
"""
|
||||
conn = get_connection()
|
||||
try:
|
||||
cursor = conn.execute(
|
||||
"UPDATE api_keys SET is_active = 0 WHERE id = ? AND is_internal = 1",
|
||||
(key_id,),
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def validate_api_key(raw_key: str) -> Optional[str]:
|
||||
"""Validate *raw_key* and return the owning username, or ``None``.
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ STAGE_PREVIEW = "preview"
|
|||
STAGE_DAG = "dag"
|
||||
STAGE_HEALTHCHECK = "healthcheck"
|
||||
STAGE_SAMPLING = "sampling"
|
||||
STAGE_SOURCE = "source"
|
||||
STAGE_COLUMN_CONFIG = "column_config"
|
||||
STAGE_GENERATING = "generating"
|
||||
STAGE_BATCH = "batch"
|
||||
|
|
|
|||
|
|
@ -33,6 +33,60 @@ from .worker import run_job_process
|
|||
_CTX = mp.get_context("spawn")
|
||||
|
||||
|
||||
def _github_source_estimated_total(recipe: dict) -> int | None:
|
||||
seed_config = recipe.get("seed_config")
|
||||
if not isinstance(seed_config, dict):
|
||||
return None
|
||||
source = seed_config.get("source")
|
||||
if not isinstance(source, dict) or source.get("seed_type") != "github_repo":
|
||||
return None
|
||||
|
||||
repos_raw = source.get("repos")
|
||||
repos = (
|
||||
[repo for repo in repos_raw if isinstance(repo, str) and repo.strip()]
|
||||
if isinstance(repos_raw, list)
|
||||
else []
|
||||
)
|
||||
item_types_raw = source.get("item_types")
|
||||
item_types = (
|
||||
[
|
||||
item
|
||||
for item in item_types_raw
|
||||
if isinstance(item, str) and item in {"issues", "pulls", "commits"}
|
||||
]
|
||||
if isinstance(item_types_raw, list)
|
||||
else []
|
||||
)
|
||||
try:
|
||||
limit = int(source.get("limit") or 0)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
if not repos or not item_types or limit <= 0:
|
||||
return None
|
||||
return len(repos) * len(item_types) * limit
|
||||
|
||||
|
||||
def _source_progress_status(job: Job) -> dict[str, Any] | None:
|
||||
progress = job.source_progress
|
||||
if progress is None:
|
||||
return None
|
||||
return {
|
||||
"source": progress.source,
|
||||
"status": progress.status,
|
||||
"repo": progress.repo,
|
||||
"resource": progress.resource,
|
||||
"page": progress.page,
|
||||
"page_items": progress.page_items,
|
||||
"fetched_items": progress.fetched_items,
|
||||
"estimated_total": progress.estimated_total,
|
||||
"percent": progress.percent,
|
||||
"rate_remaining": progress.rate_remaining,
|
||||
"retry_after_sec": progress.retry_after_sec,
|
||||
"message": progress.message,
|
||||
"updated_at": progress.updated_at,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Subscription:
|
||||
replay: list[dict]
|
||||
|
|
@ -71,8 +125,20 @@ class JobManager:
|
|||
self._pump_thread: threading.Thread | None = None
|
||||
self._seq: int = 0
|
||||
|
||||
def start(self, *, recipe: dict, run: dict) -> str:
|
||||
"""Spawn the job subprocess (one at a time, no cap)."""
|
||||
def start(
|
||||
self,
|
||||
*,
|
||||
recipe: dict,
|
||||
run: dict,
|
||||
internal_api_key_id: int | None = None,
|
||||
) -> str:
|
||||
"""Spawn the job subprocess (one at a time, no cap).
|
||||
|
||||
``internal_api_key_id`` is the row id of a workflow-scoped
|
||||
sk-unsloth-* key minted by the route layer for local providers.
|
||||
JobManager revokes it when the job reaches a terminal state so the
|
||||
key's live window is no longer than the run.
|
||||
"""
|
||||
llm_columns = recipe.get("columns") or []
|
||||
llm_column_count = 0
|
||||
if isinstance(llm_columns, list):
|
||||
|
|
@ -92,18 +158,29 @@ class JobManager:
|
|||
job_id = uuid.uuid4().hex
|
||||
self._job = Job(job_id = job_id, status = "pending", started_at = time.time())
|
||||
self._job.progress_columns_total = llm_column_count
|
||||
self._job.source_progress_estimated_total = _github_source_estimated_total(
|
||||
recipe
|
||||
)
|
||||
self._job.internal_api_key_id = internal_api_key_id
|
||||
self._events.clear()
|
||||
self._seq = 0
|
||||
|
||||
run_payload = dict(run)
|
||||
run_payload["_job_id"] = job_id
|
||||
mp_q = _CTX.Queue()
|
||||
proc = _CTX.Process(
|
||||
target = run_job_process,
|
||||
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
|
||||
daemon = True,
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
proc.start()
|
||||
|
||||
with native_path_secret_removed_for_child_start():
|
||||
mp_q = _CTX.Queue()
|
||||
proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_job_process,),
|
||||
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
|
||||
daemon = True,
|
||||
)
|
||||
proc.start()
|
||||
|
||||
self._mp_q = mp_q
|
||||
self._proc = proc
|
||||
|
|
@ -163,6 +240,7 @@ class JobManager:
|
|||
"ok": job.column_progress.ok,
|
||||
"failed": job.column_progress.failed,
|
||||
},
|
||||
"source_progress": _source_progress_status(job),
|
||||
"model_usage": {
|
||||
name: {
|
||||
"model": usage.model,
|
||||
|
|
@ -405,6 +483,7 @@ class JobManager:
|
|||
for e in self._drain_queue(mp_q):
|
||||
self._handle_event(job, e)
|
||||
|
||||
retired_job: Job | None = None
|
||||
with self._lock:
|
||||
if self._job and self._job.status in {
|
||||
"pending",
|
||||
|
|
@ -429,6 +508,9 @@ class JobManager:
|
|||
"job_id": self._job.job_id,
|
||||
}
|
||||
)
|
||||
retired_job = self._job
|
||||
if retired_job is not None:
|
||||
self._retire_workflow_key(retired_job)
|
||||
return
|
||||
|
||||
def _handle_event(self, job: Job, event: dict) -> None:
|
||||
|
|
@ -436,6 +518,7 @@ class JobManager:
|
|||
et = event.get("type")
|
||||
msg = event.get("message") if et == "log" else None
|
||||
|
||||
terminal = False
|
||||
with self._lock:
|
||||
if self._job is None or self._job.job_id != job.job_id:
|
||||
return
|
||||
|
|
@ -452,18 +535,43 @@ class JobManager:
|
|||
if self._job.progress.total and self._job.progress.total > 0:
|
||||
self._job.progress.done = self._job.progress.total
|
||||
self._job.progress.percent = 100.0
|
||||
terminal = True
|
||||
if et == EVENT_JOB_ERROR:
|
||||
self._job.status = "error"
|
||||
self._job.finished_at = time.time()
|
||||
self._job.error = event.get("error") or "error"
|
||||
terminal = True
|
||||
if et == EVENT_JOB_CANCELLED:
|
||||
terminal = True
|
||||
|
||||
if msg:
|
||||
upd = parse_log_message(msg)
|
||||
if upd:
|
||||
apply_update(self._job, upd)
|
||||
|
||||
if terminal:
|
||||
self._retire_workflow_key(job)
|
||||
|
||||
self._emit(event)
|
||||
|
||||
def _retire_workflow_key(self, job: Job) -> None:
|
||||
"""Revoke the workflow-scoped sk-unsloth-* key, if one was minted.
|
||||
|
||||
Best-effort: revocation failures are swallowed. The key would
|
||||
expire on its own after 24h, so a missed revoke is a latency
|
||||
concern, not a correctness one.
|
||||
"""
|
||||
key_id = getattr(job, "internal_api_key_id", None)
|
||||
if not key_id:
|
||||
return
|
||||
try:
|
||||
from auth import storage # deferred: avoids circular import
|
||||
|
||||
storage.revoke_internal_api_key(int(key_id))
|
||||
except Exception:
|
||||
pass
|
||||
job.internal_api_key_id = None
|
||||
|
||||
|
||||
_JOB_MANAGER: JobManager | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -17,9 +18,10 @@ from .constants import (
|
|||
STAGE_PREVIEW,
|
||||
STAGE_PROFILING,
|
||||
STAGE_SAMPLING,
|
||||
STAGE_SOURCE,
|
||||
USAGE_RESET_STAGES,
|
||||
)
|
||||
from .types import Job, ModelUsage, Progress
|
||||
from .types import Job, ModelUsage, Progress, SourceProgress
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
|
|
@ -41,6 +43,7 @@ class ParsedUpdate:
|
|||
usage_requests_total: int | None = None
|
||||
usage_rpm: float | None = None
|
||||
usage_section_start: bool | None = None
|
||||
source_progress: SourceProgress | None = None
|
||||
|
||||
|
||||
# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI.
|
||||
|
|
@ -61,9 +64,165 @@ _RE_USAGE_TOKENS = re.compile(
|
|||
_RE_USAGE_REQUESTS = re.compile(
|
||||
r"requests:\s*success=(?P<success>\d+),\s*failed=(?P<failed>\d+),\s*total=(?P<total>\d+),\s*rpm=(?P<rpm>[0-9.]+)"
|
||||
)
|
||||
_RE_GITHUB_PAGE = re.compile(
|
||||
r"^\[(?P<repo>[^\]\s]+/[^\]\s]+)\]\s+"
|
||||
r"(?P<resource>issues|PRs|commits)\s+page\s+(?P<page>\d+)\s+"
|
||||
r"\(\+(?P<items>\d+)\).*?\bremaining=(?P<remaining>\d+)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RE_GITHUB_RATE_LIMIT = re.compile(
|
||||
r"Rate limit hit\. Sleeping (?P<seconds>\d+)s until reset\.",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RE_GITHUB_SECONDARY_RATE_LIMIT = re.compile(
|
||||
r"Secondary rate limit(?: on REST)?\. Sleep (?P<seconds>\d+)s\.",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RE_GITHUB_REST_RATE_LIMIT = re.compile(
|
||||
r"REST 403/429, sleep (?P<seconds>\d+)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RE_GITHUB_TRANSIENT = re.compile(
|
||||
r"^(?P<api>GraphQL|REST) (?P<code>\d{3}) transient, retrying",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RE_GITHUB_NETWORK_RETRY = re.compile(
|
||||
r"^(?P<api>GraphQL|REST) network error: .* Retry\.",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RE_GITHUB_TRIAL_LIMIT = re.compile(
|
||||
r"Trial limit reached for (?P<resource>issues|PRs|commits) \((?P<items>\d+)\)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
_RE_GITHUB_COMPLETE = re.compile(
|
||||
r"Scraper complete\. GraphQL calls=\d+ REST calls=\d+",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def parse_log_message(msg: str) -> ParsedUpdate | None:
|
||||
m = _RE_GITHUB_PAGE.search(msg)
|
||||
if m:
|
||||
resource_raw = m.group("resource")
|
||||
resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower()
|
||||
repo = m.group("repo")
|
||||
page = int(m.group("page"))
|
||||
page_items = int(m.group("items"))
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "fetching",
|
||||
repo = repo,
|
||||
resource = resource,
|
||||
page = page,
|
||||
page_items = page_items,
|
||||
rate_remaining = int(m.group("remaining")),
|
||||
message = (
|
||||
f"Scraping GitHub source: {repo} "
|
||||
f"{resource} page {page} (+{page_items})"
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
m = _RE_GITHUB_RATE_LIMIT.search(msg)
|
||||
if m:
|
||||
seconds = int(m.group("seconds"))
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "rate_limited",
|
||||
retry_after_sec = seconds,
|
||||
message = (
|
||||
"Waiting for GitHub rate limit. "
|
||||
"Studio will resume automatically."
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
m = _RE_GITHUB_SECONDARY_RATE_LIMIT.search(msg)
|
||||
if m:
|
||||
seconds = int(m.group("seconds"))
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "rate_limited",
|
||||
retry_after_sec = seconds,
|
||||
message = (
|
||||
"Waiting for GitHub secondary rate limit. "
|
||||
"Studio will resume automatically."
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
m = _RE_GITHUB_REST_RATE_LIMIT.search(msg)
|
||||
if m:
|
||||
seconds = int(m.group("seconds"))
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "rate_limited",
|
||||
retry_after_sec = seconds,
|
||||
message = (
|
||||
"Waiting for GitHub rate limit. "
|
||||
"Studio will resume automatically."
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
m = _RE_GITHUB_TRIAL_LIMIT.search(msg)
|
||||
if m:
|
||||
resource_raw = m.group("resource")
|
||||
resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower()
|
||||
items = int(m.group("items"))
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "fetching",
|
||||
resource = resource,
|
||||
message = f"GitHub {resource} trial limit reached ({items}).",
|
||||
),
|
||||
)
|
||||
|
||||
m = _RE_GITHUB_TRANSIENT.search(msg)
|
||||
if m:
|
||||
api = m.group("api")
|
||||
code = m.group("code")
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "retrying",
|
||||
message = f"GitHub {api} returned {code}; retrying automatically.",
|
||||
),
|
||||
)
|
||||
|
||||
m = _RE_GITHUB_NETWORK_RETRY.search(msg)
|
||||
if m:
|
||||
api = m.group("api")
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "retrying",
|
||||
message = f"GitHub {api} request failed; retrying automatically.",
|
||||
),
|
||||
)
|
||||
|
||||
if _RE_GITHUB_COMPLETE.search(msg):
|
||||
return ParsedUpdate(
|
||||
stage = STAGE_SOURCE,
|
||||
source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = "completed",
|
||||
message = "GitHub source scrape complete.",
|
||||
),
|
||||
)
|
||||
|
||||
m = _RE_SAMPLERS.search(msg)
|
||||
if m:
|
||||
return ParsedUpdate(
|
||||
|
|
@ -172,6 +331,8 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
|
|||
job.batch.idx = update.batch_idx
|
||||
if update.batch_total is not None:
|
||||
job.batch.total = update.batch_total
|
||||
if update.source_progress is not None:
|
||||
_apply_source_progress(job, update.source_progress)
|
||||
|
||||
if update.stage in USAGE_RESET_STAGES:
|
||||
# usage summary is a short block so we reset once we move into the next stage.
|
||||
|
|
@ -216,6 +377,67 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
|
|||
usage.rpm = update.usage_rpm
|
||||
|
||||
|
||||
def _apply_source_progress(job: Job, progress: SourceProgress) -> None:
|
||||
previous = job.source_progress
|
||||
now = time.time()
|
||||
|
||||
page_items = progress.page_items
|
||||
if progress.repo and progress.resource and progress.page is not None:
|
||||
page_key = f"{progress.repo}:{progress.resource}:{progress.page}"
|
||||
count_key = f"{progress.repo}:{progress.resource}"
|
||||
if page_key not in job._source_seen_pages:
|
||||
job._source_seen_pages.add(page_key)
|
||||
job._source_counts[count_key] = int(
|
||||
job._source_counts.get(count_key, 0)
|
||||
) + int(page_items or 0)
|
||||
|
||||
fetched_items = sum(job._source_counts.values())
|
||||
if fetched_items <= 0:
|
||||
fetched_items = progress.fetched_items or (
|
||||
previous.fetched_items if previous else None
|
||||
)
|
||||
|
||||
estimated_total = (
|
||||
progress.estimated_total
|
||||
or job.source_progress_estimated_total
|
||||
or (previous.estimated_total if previous else None)
|
||||
)
|
||||
percent: float | None = progress.percent
|
||||
if percent is None and estimated_total and fetched_items is not None:
|
||||
raw_percent = (float(fetched_items) / float(max(1, estimated_total))) * 100.0
|
||||
percent = 100.0 if progress.status == "completed" else min(99.0, raw_percent)
|
||||
if percent is None and previous is not None:
|
||||
percent = previous.percent
|
||||
|
||||
job.source_progress = SourceProgress(
|
||||
source = "github",
|
||||
status = progress.status or (previous.status if previous else None),
|
||||
repo = progress.repo or (previous.repo if previous else None),
|
||||
resource = progress.resource or (previous.resource if previous else None),
|
||||
page = (
|
||||
progress.page
|
||||
if progress.page is not None
|
||||
else (previous.page if previous else None)
|
||||
),
|
||||
page_items = (
|
||||
page_items
|
||||
if page_items is not None
|
||||
else (previous.page_items if previous else None)
|
||||
),
|
||||
fetched_items = fetched_items,
|
||||
estimated_total = estimated_total,
|
||||
percent = percent,
|
||||
rate_remaining = (
|
||||
progress.rate_remaining
|
||||
if progress.rate_remaining is not None
|
||||
else (previous.rate_remaining if previous else None)
|
||||
),
|
||||
retry_after_sec = progress.retry_after_sec,
|
||||
message = progress.message or (previous.message if previous else None),
|
||||
updated_at = now,
|
||||
)
|
||||
|
||||
|
||||
def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
|
||||
if not job.rows:
|
||||
return column_progress
|
||||
|
|
|
|||
|
|
@ -35,6 +35,23 @@ class BatchProgress:
|
|||
total: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class SourceProgress:
|
||||
source: str = "github"
|
||||
status: str | None = None
|
||||
repo: str | None = None
|
||||
resource: str | None = None
|
||||
page: int | None = None
|
||||
page_items: int | None = None
|
||||
fetched_items: int | None = None
|
||||
estimated_total: int | None = None
|
||||
percent: float | None = None
|
||||
rate_remaining: int | None = None
|
||||
retry_after_sec: int | None = None
|
||||
message: str | None = None
|
||||
updated_at: float | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelUsage:
|
||||
model: str
|
||||
|
|
@ -57,6 +74,7 @@ class Job:
|
|||
progress: Progress = field(default_factory = Progress)
|
||||
column_progress: Progress = field(default_factory = Progress)
|
||||
batch: BatchProgress = field(default_factory = BatchProgress)
|
||||
source_progress: SourceProgress | None = None
|
||||
rows: int | None = None
|
||||
cols: int | None = None
|
||||
error: str | None = None
|
||||
|
|
@ -70,8 +88,15 @@ class Job:
|
|||
processor_artifacts: dict[str, Any] | None = None
|
||||
model_usage: dict[str, ModelUsage] = field(default_factory = dict)
|
||||
progress_columns_total: int | None = None
|
||||
source_progress_estimated_total: int | None = None
|
||||
completed_columns: list[str] = field(default_factory = list)
|
||||
# Id of the internal sk-unsloth-* API key minted for a local-model
|
||||
# workflow. Revoked when the job terminates so the key's live window
|
||||
# matches the run rather than its 24h TTL.
|
||||
internal_api_key_id: int | None = None
|
||||
_current_usage_model: str | None = None
|
||||
_in_usage_summary: bool = False
|
||||
_seen_generation_columns: list[str] = field(default_factory = list)
|
||||
_column_done: dict[str, int] = field(default_factory = dict)
|
||||
_source_counts: dict[str, int] = field(default_factory = dict)
|
||||
_source_seen_pages: set[str] = field(default_factory = set)
|
||||
|
|
|
|||
|
|
@ -21,6 +21,15 @@ from ..service import build_config_builder, create_data_designer
|
|||
from utils.paths import ensure_dir, recipe_datasets_root
|
||||
|
||||
_ARTIFACT_ROOT = recipe_datasets_root()
|
||||
_RE_GITHUB_CURSOR = re.compile(r"\bcursor=[^\s,]+")
|
||||
_RE_SECRET_TOKEN = re.compile(
|
||||
r"\b(?:(?:ghp|gho|ghu|ghs|ghr|github_pat)_[A-Za-z0-9_]+|sk-unsloth-[A-Za-z0-9]+)"
|
||||
)
|
||||
|
||||
|
||||
def _sanitize_log_message(message: str) -> str:
|
||||
message = _RE_GITHUB_CURSOR.sub("cursor=<redacted>", message)
|
||||
return _RE_SECRET_TOKEN.sub("<redacted-token>", message)
|
||||
|
||||
|
||||
class _QueueLogHandler(logging.Handler):
|
||||
|
|
@ -35,7 +44,7 @@ class _QueueLogHandler(logging.Handler):
|
|||
"ts": record.created,
|
||||
"level": record.levelname,
|
||||
"logger": record.name,
|
||||
"message": record.getMessage(),
|
||||
"message": _sanitize_log_message(record.getMessage()),
|
||||
}
|
||||
self._q.put(event)
|
||||
except (OSError, RuntimeError, ValueError):
|
||||
|
|
@ -119,10 +128,16 @@ def run_job_process(
|
|||
# Attach queue logger directly to `data_designer` so parser events survive root resets.
|
||||
handler = _QueueLogHandler(event_queue)
|
||||
handler.setLevel(logging.INFO)
|
||||
data_designer_logger = logging.getLogger("data_designer")
|
||||
data_designer_logger.addHandler(handler)
|
||||
data_designer_logger.setLevel(logging.INFO)
|
||||
data_designer_logger.propagate = True
|
||||
for logger_name in (
|
||||
"data_designer",
|
||||
"scraper",
|
||||
"gh_client",
|
||||
"data_designer_github_repo_seed",
|
||||
):
|
||||
logger = logging.getLogger(logger_name)
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(logging.INFO)
|
||||
logger.propagate = True
|
||||
|
||||
if run_config_raw:
|
||||
designer.set_run_config(RunConfig.model_validate(run_config_raw))
|
||||
|
|
@ -180,8 +195,8 @@ def run_job_process(
|
|||
{
|
||||
"type": EVENT_JOB_ERROR,
|
||||
"ts": time.time(),
|
||||
"error": str(exc),
|
||||
"stack": traceback.format_exc(limit = 20),
|
||||
"error": _sanitize_log_message(str(exc)),
|
||||
"stack": _sanitize_log_message(traceback.format_exc(limit = 20)),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ _OXC_TOOL_DIR = Path(__file__).resolve().parent / "oxc-validator"
|
|||
_OXC_RUNNER_PATH = _OXC_TOOL_DIR / "validate.mjs"
|
||||
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -248,7 +249,7 @@ def _run_oxc_batch(
|
|||
}
|
||||
try:
|
||||
tmp_dir = ensure_dir(oxc_validator_tmp_root())
|
||||
env = dict(os.environ)
|
||||
env = child_env_without_native_path_secret()
|
||||
tmp_dir_str = str(tmp_dir)
|
||||
env["TMPDIR"] = tmp_dir_str
|
||||
env["TMP"] = tmp_dir_str
|
||||
|
|
|
|||
|
|
@ -28,6 +28,8 @@ from core.inference import get_inference_backend
|
|||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
|
||||
|
||||
|
||||
def _is_wsl():
|
||||
"""Detect if running under Windows Subsystem for Linux."""
|
||||
|
|
@ -529,6 +531,31 @@ class ExportBackend:
|
|||
# Convert quantization method to lowercase for unsloth
|
||||
quant_method = quantization_method.lower()
|
||||
|
||||
# Pin convert_hf_to_gguf.py to the same llama.cpp ref as the
|
||||
# llama-quantize binary (Studio installs at a tagged ref via
|
||||
# setup.sh) so it can't drift past the pinned binary's gguf API.
|
||||
# Set before both branches; hub-only export has save_directory == "".
|
||||
global _LLAMA_CPP_SCRIPTS_WARNING_EMITTED
|
||||
try:
|
||||
from unsloth_zoo.llama_cpp import (
|
||||
LLAMA_CPP_DEFAULT_DIR,
|
||||
_resolve_local_convert_script, # noqa: F401
|
||||
)
|
||||
|
||||
os.environ.setdefault(
|
||||
"UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR
|
||||
)
|
||||
except ImportError:
|
||||
if not _LLAMA_CPP_SCRIPTS_WARNING_EMITTED:
|
||||
logger.warning(
|
||||
"Unsloth: installed unsloth_zoo does not honor "
|
||||
"UNSLOTH_LLAMA_CPP_SCRIPTS_DIR; convert_hf_to_gguf.py will "
|
||||
"still be downloaded from llama.cpp master and may drift "
|
||||
"past the pinned llama-quantize binary. Upgrade unsloth_zoo "
|
||||
"to activate the local script pin."
|
||||
)
|
||||
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True
|
||||
|
||||
# Save locally if requested
|
||||
if save_directory:
|
||||
save_directory = str(resolve_export_dir(save_directory))
|
||||
|
|
|
|||
|
|
@ -163,21 +163,28 @@ class ExportOrchestrator:
|
|||
|
||||
def _spawn_subprocess(self, config: dict) -> None:
|
||||
"""Spawn a new export subprocess."""
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
|
||||
from .worker import run_export_process
|
||||
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
with native_path_secret_removed_for_child_start():
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
|
||||
self._proc = _CTX.Process(
|
||||
target = run_export_process,
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
self._proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_export_process,),
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
|
||||
|
||||
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from typing import Optional, Tuple
|
|||
import numpy as np
|
||||
import torch
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -105,6 +106,7 @@ class AudioCodecManager:
|
|||
spark_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
|
||||
|
|
@ -143,6 +145,7 @@ class AudioCodecManager:
|
|||
outetts_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
# Remove files that pull in heavy / incompatible dependencies
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
120
studio/backend/core/inference/llama_server_args.py
Normal file
120
studio/backend/core/inference/llama_server_args.py
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Validator for user-supplied llama-server pass-through args.
|
||||
|
||||
Studio runs llama-server as a managed subprocess and lets callers pass
|
||||
extra flags directly (CLI: ``unsloth run ... --top-k 20``; HTTP:
|
||||
``LoadRequest.llama_extra_args``). This module is the boundary that
|
||||
rejects only flags Studio fundamentally cannot share with the user --
|
||||
model identity, the auth key, and the network endpoint Studio's HTTP
|
||||
proxy targets. Anything else passes through.
|
||||
|
||||
User-supplied args are appended to ``cmd`` after Studio's auto-set
|
||||
flags, so llama.cpp's last-wins CLI parsing makes the user's value
|
||||
override the auto-set one. That covers tunable knobs the user might
|
||||
reasonably want to override -- ``-c``/``--ctx-size``,
|
||||
``-np``/``--parallel``, ``-fa``/``--flash-attn``,
|
||||
``-ngl``/``--gpu-layers``, ``-t``/``--threads``, ``-fit``/``--fit*``,
|
||||
``--cache-type-k/v``, ``--chat-template-file/-kwargs``,
|
||||
``--spec-*``, ``--jinja``/``--no-jinja``,
|
||||
``--no-context-shift``/``--context-shift``, sampling params, etc.
|
||||
|
||||
Reference: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable, Optional
|
||||
|
||||
# Each group is the full set of aliases (short + long) for one
|
||||
# hard-denied flag, taken from the llama-server README. If llama.cpp
|
||||
# adds a new alias for an existing denied flag, extend the relevant
|
||||
# group.
|
||||
#
|
||||
# Flags NOT in this list (e.g. -c, --parallel, --flash-attn, -ngl,
|
||||
# -t/--threads, --jinja, --no-context-shift, --fit*, --cache-type-*,
|
||||
# --chat-template-*, --spec-*) pass through and override Studio's
|
||||
# auto-set version via llama.cpp's last-wins CLI parsing.
|
||||
_DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
|
||||
# Model identity -- Studio resolves the model from LoadRequest and
|
||||
# passes -m / mmproj after downloading from HF if needed. A second
|
||||
# -m would point at a different model than the one Studio thinks
|
||||
# is loaded.
|
||||
frozenset({"-m", "--model"}),
|
||||
frozenset({"-mu", "--model-url"}),
|
||||
frozenset({"-dr", "--docker-repo"}),
|
||||
frozenset({"-hf", "-hfr", "--hf-repo"}),
|
||||
frozenset({"-hff", "--hf-file"}),
|
||||
frozenset({"-hfv", "-hfrv", "--hf-repo-v"}),
|
||||
frozenset({"-hffv", "--hf-file-v"}),
|
||||
frozenset({"-hft", "--hf-token"}),
|
||||
frozenset({"-mm", "--mmproj"}),
|
||||
frozenset({"-mmu", "--mmproj-url"}),
|
||||
# Networking -- Studio binds llama-server's port and reverse-proxies
|
||||
# HTTP traffic to it. Retargeting host/port/path/prefix would
|
||||
# orphan Studio's proxy and the UI would lose the server.
|
||||
frozenset({"--host"}),
|
||||
frozenset({"--port"}),
|
||||
frozenset({"--path"}),
|
||||
frozenset({"--api-prefix"}),
|
||||
frozenset({"--reuse-port"}),
|
||||
# Auth / TLS -- Studio terminates auth at its own layer; an
|
||||
# upstream --api-key would shadow Studio's UNSLOTH_DIRECT_STREAM
|
||||
# key, and TLS on llama-server would break the local proxy hop.
|
||||
frozenset({"--api-key"}),
|
||||
frozenset({"--api-key-file"}),
|
||||
frozenset({"--ssl-key-file"}),
|
||||
frozenset({"--ssl-cert-file"}),
|
||||
# Single-model server -- Studio runs one model per llama-server
|
||||
# process and serves its own UI. Enabling multi-model loading or
|
||||
# llama-server's built-in web UI changes the surface clients see.
|
||||
frozenset({"--webui", "--no-webui"}),
|
||||
frozenset({"--models-dir"}),
|
||||
frozenset({"--models-preset"}),
|
||||
frozenset({"--models-max"}),
|
||||
frozenset({"--models-autoload", "--no-models-autoload"}),
|
||||
)
|
||||
|
||||
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
|
||||
|
||||
|
||||
def _flag_name(token: str) -> Optional[str]:
|
||||
"""Return the flag name for a token, or None if it isn't a flag.
|
||||
|
||||
Peels ``--key=value`` to the bare ``--key``. Plain numeric values
|
||||
like ``-1`` or ``-0.5`` (e.g. ``--seed -1``) are values, not flags;
|
||||
llama-server short-form flags always start with a letter.
|
||||
"""
|
||||
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 validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
|
||||
"""Validate user-supplied llama-server args.
|
||||
|
||||
Returns the args as a flat list ready to extend the llama-server
|
||||
command. Raises ``ValueError`` (with the offending flag in the
|
||||
message) the moment a token resolves to a Studio-managed flag.
|
||||
"""
|
||||
if not args:
|
||||
return []
|
||||
out: list[str] = []
|
||||
for raw in args:
|
||||
token = str(raw)
|
||||
flag = _flag_name(token)
|
||||
if flag is not None and flag in _DENYLIST:
|
||||
raise ValueError(
|
||||
f"llama-server flag '{flag}' is managed by Unsloth Studio "
|
||||
f"and cannot be passed as an extra arg"
|
||||
)
|
||||
out.append(token)
|
||||
return out
|
||||
|
||||
|
||||
def is_managed_flag(flag: str) -> bool:
|
||||
"""True if ``flag`` is a Studio-managed llama-server flag."""
|
||||
return flag in _DENYLIST
|
||||
|
|
@ -166,23 +166,30 @@ class InferenceOrchestrator:
|
|||
|
||||
def _spawn_subprocess(self, config: dict) -> None:
|
||||
"""Spawn a new inference subprocess."""
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
|
||||
from .worker import run_inference_process
|
||||
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
self._cancel_event = _CTX.Event()
|
||||
with native_path_secret_removed_for_child_start():
|
||||
self._cmd_queue = _CTX.Queue()
|
||||
self._resp_queue = _CTX.Queue()
|
||||
self._cancel_event = _CTX.Event()
|
||||
|
||||
self._proc = _CTX.Process(
|
||||
target = run_inference_process,
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"cancel_event": self._cancel_event,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
self._proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_inference_process,),
|
||||
kwargs = {
|
||||
"cmd_queue": self._cmd_queue,
|
||||
"resp_queue": self._resp_queue,
|
||||
"cancel_event": self._cancel_event,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
self._proc.start()
|
||||
logger.info("Inference subprocess started (pid=%s)", self._proc.pid)
|
||||
|
||||
def _cancel_generation(self) -> None:
|
||||
|
|
@ -708,6 +715,17 @@ class InferenceOrchestrator:
|
|||
|
||||
def unload_model(self, model_name: str) -> bool:
|
||||
"""Unload a model from the subprocess."""
|
||||
if model_name in self.loading_models:
|
||||
logger.info(
|
||||
"Cancelling in-flight load for model '%s' by terminating subprocess",
|
||||
model_name,
|
||||
)
|
||||
self._shutdown_subprocess(timeout = 0.5)
|
||||
self.loading_models.discard(model_name)
|
||||
self.active_model_name = None
|
||||
self.models.clear()
|
||||
return True
|
||||
|
||||
if not self._ensure_subprocess_alive():
|
||||
# No subprocess — just clear local state
|
||||
self.models.pop(model_name, None)
|
||||
|
|
|
|||
75
studio/backend/core/training/resume.py
Normal file
75
studio/backend/core/training/resume.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Helpers for validating resumable training outputs."""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from utils.paths import outputs_root, resolve_output_dir
|
||||
|
||||
|
||||
def _is_under_outputs(path: Path) -> bool:
|
||||
resolved = path.resolve(strict = False)
|
||||
root = outputs_root().resolve(strict = False)
|
||||
try:
|
||||
resolved.relative_to(root)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def has_resume_state(path_value: Optional[str]) -> bool:
|
||||
if not path_value:
|
||||
return False
|
||||
return get_resume_checkpoint_path(path_value) is not None
|
||||
|
||||
|
||||
def _checkpoint_step(path: Path) -> int:
|
||||
try:
|
||||
return int(path.name.removeprefix("checkpoint-"))
|
||||
except ValueError:
|
||||
return -1
|
||||
|
||||
|
||||
def get_resume_checkpoint_path(path_value: str) -> Optional[str]:
|
||||
path = resolve_output_dir(path_value)
|
||||
if not _is_under_outputs(path) or not path.is_dir():
|
||||
return None
|
||||
if (path / "trainer_state.json").is_file():
|
||||
return str(path)
|
||||
|
||||
checkpoints = [
|
||||
child
|
||||
for child in path.glob("checkpoint-*")
|
||||
if child.is_dir() and (child / "trainer_state.json").is_file()
|
||||
]
|
||||
if not checkpoints:
|
||||
return None
|
||||
return str(max(checkpoints, key = _checkpoint_step))
|
||||
|
||||
|
||||
def normalize_resume_output_dir(path_value: str) -> str:
|
||||
path = resolve_output_dir(path_value)
|
||||
if not _is_under_outputs(path):
|
||||
raise ValueError("Resume checkpoint must be inside Studio outputs.")
|
||||
return str(path)
|
||||
|
||||
|
||||
def can_resume_run(run: dict) -> bool:
|
||||
if run.get("resumed_later"):
|
||||
return False
|
||||
|
||||
final_step = run.get("final_step")
|
||||
total_steps = run.get("total_steps")
|
||||
has_remaining_steps = (
|
||||
not isinstance(final_step, int)
|
||||
or not isinstance(total_steps, int)
|
||||
or total_steps <= 0
|
||||
or final_step < total_steps
|
||||
)
|
||||
return (
|
||||
run.get("status") == "stopped"
|
||||
and has_remaining_steps
|
||||
and has_resume_state(run.get("output_dir"))
|
||||
)
|
||||
|
|
@ -70,6 +70,7 @@ from utils.paths import (
|
|||
)
|
||||
from trl import SFTTrainer, SFTConfig
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -376,6 +377,7 @@ class UnslothTrainer:
|
|||
def _finalize_training(self, output_dir, label = ""):
|
||||
"""Save model after training and update progress. Used by all training branches."""
|
||||
if self.should_stop and self.save_on_stop:
|
||||
self.trainer._save_checkpoint(self.trainer.model, trial = None)
|
||||
self.trainer.save_model()
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
self._patch_adapter_config(output_dir)
|
||||
|
|
@ -1770,6 +1772,7 @@ class UnslothTrainer:
|
|||
spark_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
|
||||
|
|
@ -2004,6 +2007,7 @@ class UnslothTrainer:
|
|||
outetts_code_dir,
|
||||
],
|
||||
check = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
for fpath in [
|
||||
|
|
@ -2828,7 +2832,9 @@ class UnslothTrainer:
|
|||
total_steps = total, status_message = "Starting CSM training..."
|
||||
)
|
||||
logger.info(f"CSM training config: {config}\n")
|
||||
self.trainer.train()
|
||||
self.trainer.train(
|
||||
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
|
||||
)
|
||||
self._finalize_training(output_dir, "CSM")
|
||||
return
|
||||
|
||||
|
|
@ -2867,7 +2873,9 @@ class UnslothTrainer:
|
|||
total_steps = total, status_message = "Starting SNAC training..."
|
||||
)
|
||||
logger.info(f"SNAC training config: {config}\n")
|
||||
self.trainer.train()
|
||||
self.trainer.train(
|
||||
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
|
||||
)
|
||||
self._finalize_training(output_dir, "SNAC")
|
||||
return
|
||||
|
||||
|
|
@ -2913,7 +2921,9 @@ class UnslothTrainer:
|
|||
total_steps = total, status_message = "Starting Whisper training..."
|
||||
)
|
||||
logger.info(f"Whisper training config: {config}\n")
|
||||
self.trainer.train()
|
||||
self.trainer.train(
|
||||
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
|
||||
)
|
||||
self._finalize_training(output_dir, "Whisper")
|
||||
return
|
||||
|
||||
|
|
@ -3408,7 +3418,9 @@ class UnslothTrainer:
|
|||
# ========== START TRAINING ==========
|
||||
self._update_progress(status_message = "Starting training...")
|
||||
logger.info("Starting training...\n")
|
||||
self.trainer.train()
|
||||
self.trainer.train(
|
||||
resume_from_checkpoint = training_args.get("resume_from_checkpoint")
|
||||
)
|
||||
|
||||
# ========== SAVE MODEL ==========
|
||||
self._finalize_training(output_dir)
|
||||
|
|
|
|||
|
|
@ -29,6 +29,10 @@ from typing import Optional, Tuple, Any
|
|||
|
||||
import matplotlib.pyplot as plt
|
||||
from utils.hardware import prepare_gpu_selection
|
||||
from utils.native_path_leases import (
|
||||
native_path_secret_removed_for_child_start,
|
||||
run_without_native_path_secret,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -185,6 +189,7 @@ class TrainingBackend:
|
|||
"wandb_project": kwargs.get("wandb_project", "unsloth-training"),
|
||||
"enable_tensorboard": kwargs.get("enable_tensorboard", False),
|
||||
"tensorboard_dir": kwargs.get("tensorboard_dir", "runs"),
|
||||
"resume_from_checkpoint": kwargs.get("resume_from_checkpoint"),
|
||||
"trust_remote_code": kwargs.get("trust_remote_code", False),
|
||||
"gpu_ids": kwargs.get("gpu_ids"),
|
||||
}
|
||||
|
|
@ -212,20 +217,22 @@ class TrainingBackend:
|
|||
|
||||
from .worker import run_training_process
|
||||
|
||||
event_queue = _CTX.Queue()
|
||||
stop_queue = _CTX.Queue()
|
||||
|
||||
proc = _CTX.Process(
|
||||
target = run_training_process,
|
||||
kwargs = {
|
||||
"event_queue": event_queue,
|
||||
"stop_queue": stop_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
try:
|
||||
proc.start()
|
||||
with native_path_secret_removed_for_child_start():
|
||||
event_queue = _CTX.Queue()
|
||||
stop_queue = _CTX.Queue()
|
||||
|
||||
proc = _CTX.Process(
|
||||
target = run_without_native_path_secret,
|
||||
args = (run_training_process,),
|
||||
kwargs = {
|
||||
"event_queue": event_queue,
|
||||
"stop_queue": stop_queue,
|
||||
"config": config,
|
||||
},
|
||||
daemon = True,
|
||||
)
|
||||
proc.start()
|
||||
except Exception:
|
||||
logger.error("Failed to start training subprocess", exc_info = True)
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -35,6 +35,15 @@ from utils.wheel_utils import (
|
|||
)
|
||||
|
||||
|
||||
def _output_dir_from_resume_checkpoint(
|
||||
resume_from_checkpoint: str | None,
|
||||
) -> str | None:
|
||||
if not resume_from_checkpoint:
|
||||
return None
|
||||
path = Path(resume_from_checkpoint)
|
||||
return str(path.parent if path.name.startswith("checkpoint-") else path)
|
||||
|
||||
|
||||
_CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4"
|
||||
_CAUSAL_CONV1D_PACKAGE_VERSION = "1.6.1"
|
||||
_MAMBA_SSM_RELEASE_TAG = "v2.3.1"
|
||||
|
|
@ -50,6 +59,8 @@ def _model_wants_causal_conv1d(model_name: str) -> bool:
|
|||
for key in (
|
||||
"qwen3.5",
|
||||
"qwen3_5",
|
||||
"qwen3.6",
|
||||
"qwen3_6",
|
||||
"qwen3-next",
|
||||
"qwen3_next",
|
||||
"nemotron_h",
|
||||
|
|
@ -755,7 +766,10 @@ def run_training_process(
|
|||
return
|
||||
|
||||
# Generate output dir
|
||||
output_dir = config.get("output_dir")
|
||||
resume_from_checkpoint = config.get("resume_from_checkpoint")
|
||||
output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint(
|
||||
resume_from_checkpoint
|
||||
)
|
||||
if not output_dir:
|
||||
output_dir = f"{model_name.replace('/', '_')}_{int(time.time())}"
|
||||
output_dir = str(resolve_output_dir(output_dir))
|
||||
|
|
@ -803,6 +817,7 @@ def run_training_process(
|
|||
max_seq_length = config.get("max_seq_length", 2048),
|
||||
optim = config.get("optim", "adamw_8bit"),
|
||||
lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
|
||||
resume_from_checkpoint = resume_from_checkpoint,
|
||||
)
|
||||
|
||||
_tqdm_stop.set()
|
||||
|
|
@ -819,10 +834,13 @@ def run_training_process(
|
|||
}
|
||||
)
|
||||
else:
|
||||
saved_output_dir = (
|
||||
None if trainer.should_stop and not trainer.save_on_stop else output_dir
|
||||
)
|
||||
event_queue.put(
|
||||
{
|
||||
"type": "complete",
|
||||
"output_dir": output_dir,
|
||||
"output_dir": saved_output_dir,
|
||||
"status_message": progress.status_message or "Training completed",
|
||||
"ts": time.time(),
|
||||
}
|
||||
|
|
@ -1107,11 +1125,15 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
|
|||
)
|
||||
return
|
||||
|
||||
output_dir = config.get("output_dir")
|
||||
resume_from_checkpoint = config.get("resume_from_checkpoint")
|
||||
output_dir = config.get("output_dir") or _output_dir_from_resume_checkpoint(
|
||||
resume_from_checkpoint
|
||||
)
|
||||
if not output_dir:
|
||||
output_dir = str(
|
||||
resolve_output_dir(f"{model_name.replace('/', '_')}_{int(time.time())}")
|
||||
)
|
||||
output_dir = str(resolve_output_dir(output_dir))
|
||||
|
||||
num_epochs = config.get("num_epochs", 2)
|
||||
batch_size = config.get("batch_size", 256)
|
||||
|
|
@ -1219,7 +1241,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
|
|||
callbacks = [_EmbeddingProgressCallback()],
|
||||
)
|
||||
|
||||
trainer.train()
|
||||
trainer.train(resume_from_checkpoint = resume_from_checkpoint)
|
||||
except Exception as e:
|
||||
event_queue.put(
|
||||
{
|
||||
|
|
@ -1245,6 +1267,8 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
|
|||
|
||||
_send_status(event_queue, "Saving model...")
|
||||
try:
|
||||
if _should_stop and _save_on_stop:
|
||||
trainer._save_checkpoint(trainer.model, trial = None)
|
||||
model.save_pretrained(output_dir)
|
||||
model.tokenizer.save_pretrained(output_dir)
|
||||
logger.info("Embedding model saved to %s", output_dir)
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ from typing import Optional
|
|||
|
||||
import structlog
|
||||
|
||||
from loggers.handlers import filter_sensitive_data
|
||||
|
||||
|
||||
class LogConfig:
|
||||
"""Structured logging configuration for the application.
|
||||
|
|
@ -58,6 +60,8 @@ class LogConfig:
|
|||
structlog.processors.TimeStamper(fmt = "iso"), # timestamp first
|
||||
structlog.processors.add_log_level, # level second
|
||||
structlog.contextvars.merge_contextvars,
|
||||
structlog.processors.format_exc_info,
|
||||
filter_sensitive_data,
|
||||
# Custom processor to flatten the extra field
|
||||
lambda logger, method_name, event_dict: {
|
||||
"timestamp": event_dict.get("timestamp"),
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ Key Components:
|
|||
- get_logger: Factory function for structured loggers
|
||||
"""
|
||||
|
||||
import re
|
||||
import time
|
||||
from typing import Callable
|
||||
|
||||
|
|
@ -22,7 +23,12 @@ import structlog
|
|||
from fastapi import Request, Response
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
|
||||
from utils.native_path_leases import redact_native_paths
|
||||
|
||||
logger = structlog.get_logger(__name__)
|
||||
_NATIVE_PATH_LEASE_RE = re.compile(
|
||||
r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"
|
||||
)
|
||||
|
||||
|
||||
class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
|
|
@ -75,6 +81,12 @@ def filter_sensitive_data(logger, method_name, event_dict):
|
|||
"""Structlog processor to filter out base64 data from logs."""
|
||||
|
||||
def filter_value(value):
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = redact_native_paths(value)
|
||||
except Exception:
|
||||
pass
|
||||
value = _NATIVE_PATH_LEASE_RE.sub(r"\1<redacted native path lease>", value)
|
||||
if (
|
||||
isinstance(value, str)
|
||||
and len(value) > 100
|
||||
|
|
@ -83,12 +95,22 @@ def filter_sensitive_data(logger, method_name, event_dict):
|
|||
# Likely base64 data, truncate it
|
||||
return value[:20] + "..."
|
||||
elif isinstance(value, dict):
|
||||
return {k: filter_value(v) for k, v in value.items()}
|
||||
return {
|
||||
k: "<redacted native path lease>"
|
||||
if str(k).replace("_", "").lower() == "nativepathlease"
|
||||
else filter_value(v)
|
||||
for k, v in value.items()
|
||||
}
|
||||
elif isinstance(value, list):
|
||||
return [filter_value(item) for item in value]
|
||||
return value
|
||||
|
||||
return {k: filter_value(v) for k, v in event_dict.items()}
|
||||
return {
|
||||
k: "<redacted native path lease>"
|
||||
if str(k).replace("_", "").lower() == "nativepathlease"
|
||||
else filter_value(v)
|
||||
for k, v in event_dict.items()
|
||||
}
|
||||
|
||||
|
||||
def get_logger(name: str) -> structlog.BoundLogger:
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ from routes import (
|
|||
datasets_router,
|
||||
export_router,
|
||||
inference_router,
|
||||
inference_studio_router,
|
||||
models_router,
|
||||
training_history_router,
|
||||
training_router,
|
||||
|
|
@ -77,6 +78,7 @@ from utils.hardware import (
|
|||
import utils.hardware.hardware as _hw_module
|
||||
|
||||
from utils.cache_cleanup import clear_unsloth_compiled_cache
|
||||
from utils.native_path_leases import native_path_leases_supported
|
||||
|
||||
|
||||
def get_unsloth_version() -> str:
|
||||
|
|
@ -186,6 +188,8 @@ if _api_only:
|
|||
"tauri://localhost", # Linux/macOS Tauri webview
|
||||
"http://tauri.localhost", # Windows Tauri webview
|
||||
"http://localhost", # dev fallback
|
||||
"http://localhost:5173", # Tauri dev/Vite
|
||||
"http://127.0.0.1:5173", # Tauri dev/Vite fallback
|
||||
]
|
||||
_cors_origin_regex = None
|
||||
else:
|
||||
|
|
@ -207,6 +211,9 @@ app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"])
|
|||
app.include_router(training_router, prefix = "/api/train", tags = ["training"])
|
||||
app.include_router(models_router, prefix = "/api/models", tags = ["models"])
|
||||
app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
|
||||
# Studio-only inference endpoints (cancel, etc.) are intentionally NOT
|
||||
# exposed on the /v1 OpenAI-compat prefix below.
|
||||
app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"])
|
||||
|
||||
# OpenAI-compatible endpoints: mount the same inference router at /v1
|
||||
# so external tools (Open WebUI, SillyTavern, etc.) can use the
|
||||
|
|
@ -238,6 +245,7 @@ async def health_check():
|
|||
"chat_only": _hw_module.CHAT_ONLY,
|
||||
"desktop_protocol_version": 1,
|
||||
"supports_desktop_auth": True,
|
||||
"native_path_leases_supported": native_path_leases_supported(),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -18,6 +18,9 @@ class LoadRequest(BaseModel):
|
|||
"""Request to load a model for inference"""
|
||||
|
||||
model_path: str = Field(..., description = "Model identifier or local path")
|
||||
native_path_lease: Optional[str] = Field(
|
||||
None, description = "Frontend-visible signed native path grant"
|
||||
)
|
||||
hf_token: Optional[str] = Field(
|
||||
None, description = "HuggingFace token for gated models"
|
||||
)
|
||||
|
|
@ -52,6 +55,16 @@ class LoadRequest(BaseModel):
|
|||
None,
|
||||
description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.",
|
||||
)
|
||||
llama_extra_args: Optional[List[str]] = Field(
|
||||
None,
|
||||
description = (
|
||||
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
|
||||
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
|
||||
"Studio-managed flags (model identity, port, context length, GPU placement, "
|
||||
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
|
||||
"non-GGUF models."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class UnloadRequest(BaseModel):
|
||||
|
|
@ -69,6 +82,9 @@ class ValidateModelRequest(BaseModel):
|
|||
"""
|
||||
|
||||
model_path: str = Field(..., description = "Model identifier or local path")
|
||||
native_path_lease: Optional[str] = Field(
|
||||
None, description = "Frontend-visible signed native path grant"
|
||||
)
|
||||
hf_token: Optional[str] = Field(
|
||||
None, description = "HuggingFace token for gated models"
|
||||
)
|
||||
|
|
@ -283,6 +299,10 @@ class InferenceStatusResponse(BaseModel):
|
|||
supports_tools: bool = Field(
|
||||
False, description = "Whether the active model supports tool calling"
|
||||
)
|
||||
chat_template: Optional[str] = Field(
|
||||
None,
|
||||
description = "Jinja2 chat template string for the active model",
|
||||
)
|
||||
context_length: Optional[int] = Field(
|
||||
None, description = "Context length of the active model"
|
||||
)
|
||||
|
|
@ -396,7 +416,10 @@ class ChatMessage(BaseModel):
|
|||
if self.name is not None and self.role != "tool":
|
||||
raise ValueError('"name" is only valid on role="tool" messages.')
|
||||
|
||||
# Per-role content requirements.
|
||||
# Per-role content requirements. OpenAI-compatible clients may send
|
||||
# ``content=""`` for image-only turns when the image travels in a
|
||||
# companion field such as Studio's ``image_base64`` extension, so treat
|
||||
# empty strings as present content for user/system messages.
|
||||
if self.role == "tool":
|
||||
if not self.tool_call_id:
|
||||
raise ValueError(
|
||||
|
|
@ -411,10 +434,8 @@ class ChatMessage(BaseModel):
|
|||
'role="assistant" messages require either "content" or "tool_calls".'
|
||||
)
|
||||
else: # "user" | "system"
|
||||
if not self.content:
|
||||
raise ValueError(
|
||||
f'role="{self.role}" messages require non-empty "content".'
|
||||
)
|
||||
if self.content is None or self.content == []:
|
||||
raise ValueError(f'role="{self.role}" messages require "content".')
|
||||
return self
|
||||
|
||||
|
||||
|
|
@ -531,6 +552,10 @@ class ChatCompletionRequest(BaseModel):
|
|||
None,
|
||||
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
|
||||
)
|
||||
cancel_id: Optional[str] = Field(
|
||||
None,
|
||||
description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.",
|
||||
)
|
||||
|
||||
|
||||
# ── Streaming response chunks ────────────────────────────────────
|
||||
|
|
@ -992,6 +1017,7 @@ class AnthropicMessagesRequest(BaseModel):
|
|||
enable_tools: Optional[bool] = None
|
||||
enabled_tools: Optional[list[str]] = None
|
||||
session_id: Optional[str] = None
|
||||
cancel_id: Optional[str] = None
|
||||
model_config = {"extra": "allow"}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -127,6 +127,9 @@ class TrainingStartRequest(BaseModel):
|
|||
wandb_project: Optional[str] = Field(None, description = "W&B project name")
|
||||
enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging")
|
||||
tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory")
|
||||
resume_from_checkpoint: Optional[str] = Field(
|
||||
None, description = "Saved training output directory to resume from"
|
||||
)
|
||||
|
||||
# GPU selection
|
||||
gpu_ids: Optional[List[int]] = Field(
|
||||
|
|
@ -220,6 +223,8 @@ class TrainingRunSummary(BaseModel):
|
|||
duration_seconds: Optional[float] = None
|
||||
error_message: Optional[str] = None
|
||||
loss_sparkline: Optional[List[float]] = None
|
||||
can_resume: bool = False
|
||||
resumed_later: bool = False
|
||||
|
||||
|
||||
class TrainingRunListResponse(BaseModel):
|
||||
|
|
|
|||
|
|
@ -0,0 +1,73 @@
|
|||
# data-designer-github-repo-seed
|
||||
|
||||
A Data Designer seed-reader plugin for **Unsloth Studio** that scrapes real
|
||||
GitHub data (issues, pull requests, commits) from one or more repositories
|
||||
and hands it to the recipe pipeline as a seed dataset.
|
||||
|
||||
Designed to ship with Studio as a default seed source so any user with a
|
||||
GitHub token can build training datasets straight from live repos.
|
||||
|
||||
## What it does
|
||||
|
||||
Given a list of `owner/name` repos, a GitHub token, and a per-resource
|
||||
`limit`, the plugin uses GitHub's GraphQL API to fetch issues, pull
|
||||
requests, and/or commits, with labels, state, authors, and the first N
|
||||
comments of each item, and materialises a single JSONL with uniform
|
||||
columns so the rest of the recipe (LLM text / LLM structured / processors)
|
||||
can treat it like any other seed table.
|
||||
|
||||
| Column | Description |
|
||||
|---------------|------------------------------------------------|
|
||||
| `item_type` | `issue` / `pull` / `commit` |
|
||||
| `repo` | `owner/name` |
|
||||
| `number` | Issue/PR number, or commit SHA |
|
||||
| `title` | Title (or commit message headline) |
|
||||
| `body` | Issue/PR body (or full commit message) |
|
||||
| `state` | `OPEN` / `CLOSED` / `MERGED` (empty for commit)|
|
||||
| `author` | GitHub login of the author |
|
||||
| `created_at` | ISO8601 |
|
||||
| `closed_at` | ISO8601 (empty for commits) |
|
||||
| `url` | Permalink |
|
||||
| `labels` | List of label names |
|
||||
| `comments` | First N comments concatenated |
|
||||
|
||||
## Usage in a recipe
|
||||
|
||||
```json
|
||||
{
|
||||
"seed_config": {
|
||||
"source": {
|
||||
"seed_type": "github_repo",
|
||||
"repos": ["unslothai/unsloth", "unslothai/unsloth-zoo"],
|
||||
"token": "",
|
||||
"item_types": ["issues", "pulls"],
|
||||
"limit": 100,
|
||||
"include_comments": true,
|
||||
"max_comments_per_item": 30
|
||||
},
|
||||
"sampling_strategy": "shuffle",
|
||||
"selection_strategy": null
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Leave `token` empty to fall back to the server's `GH_TOKEN` / `GITHUB_TOKEN`
|
||||
environment variable, useful when the recipe is published and shouldn't
|
||||
carry a secret.
|
||||
|
||||
## Auth
|
||||
|
||||
A GitHub personal access token with `public_repo` scope is enough for public
|
||||
repositories; `repo` scope is required for private ones. GraphQL requests
|
||||
are rate-limit aware: the client inspects `x-ratelimit-*` headers and
|
||||
sleeps until reset when the budget drops below a safety threshold.
|
||||
|
||||
## Install
|
||||
|
||||
Shipped as a default Studio plugin. For development:
|
||||
|
||||
```bash
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
Registered automatically via the `data_designer.plugins` entry point.
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=68", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "data-designer-github-repo-seed"
|
||||
version = "0.1.0"
|
||||
description = "Unsloth Studio seed plugin that scrapes GitHub issues, PRs, and commits."
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"data-designer-engine>=0.5.4,<0.6",
|
||||
"requests>=2.31",
|
||||
]
|
||||
|
||||
[project.entry-points."data_designer.plugins"]
|
||||
github_repo_seed = "data_designer_github_repo_seed.plugin:github_repo_seed_plugin"
|
||||
|
||||
[tool.setuptools]
|
||||
package-dir = {"" = "src"}
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
# Intentionally empty. Data-designer loads submodules lazily via qualified names
|
||||
# (impl_qualified_name / config_qualified_name in plugin.py), so importing this
|
||||
# package must NOT touch modules that depend on data_designer.engine.* during
|
||||
# Studio's bootstrap (circular import).
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field, field_validator, model_validator
|
||||
|
||||
from data_designer.config.seed_source import SeedSource
|
||||
|
||||
|
||||
class GitHubRepoSeedSource(SeedSource):
|
||||
seed_type: Literal["github_repo"] = "github_repo"
|
||||
|
||||
repos: list[str] = Field(
|
||||
default_factory = list,
|
||||
description = "List of GitHub repositories to scrape, each in `owner/name` form.",
|
||||
)
|
||||
token: str = Field(
|
||||
default = "",
|
||||
description = "Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.",
|
||||
)
|
||||
item_types: list[Literal["issues", "pulls", "commits"]] = Field(
|
||||
default = ["issues", "pulls"],
|
||||
description = "Which GitHub item types to fetch per repo.",
|
||||
)
|
||||
limit: int = Field(
|
||||
default = 100,
|
||||
ge = 1,
|
||||
le = 5000,
|
||||
description = "Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).",
|
||||
)
|
||||
include_comments: bool = Field(
|
||||
default = True,
|
||||
description = "Fetch the first N comments of each issue/PR and include them in the `comments` column.",
|
||||
)
|
||||
max_comments_per_item: int = Field(default = 30, ge = 0, le = 200)
|
||||
|
||||
@field_validator("repos")
|
||||
@classmethod
|
||||
def _validate_repos(cls, v: list[str]) -> list[str]:
|
||||
out: list[str] = []
|
||||
for r in v or []:
|
||||
r = r.strip()
|
||||
if not r:
|
||||
continue
|
||||
if r.count("/") != 1 or not all(r.split("/")):
|
||||
raise ValueError(f"Each repo must be `owner/name`; got {r!r}")
|
||||
out.append(r)
|
||||
return out
|
||||
|
||||
@field_validator("item_types")
|
||||
@classmethod
|
||||
def _validate_item_types(cls, v: list[str]) -> list[str]:
|
||||
if not v:
|
||||
raise ValueError("item_types must not be empty")
|
||||
return list(dict.fromkeys(v))
|
||||
|
||||
@model_validator(mode = "after")
|
||||
def _ensure_repos(self) -> "GitHubRepoSeedSource":
|
||||
if not self.repos:
|
||||
raise ValueError("At least one repo is required")
|
||||
return self
|
||||
|
|
@ -0,0 +1,83 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import tempfile
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import data_designer.lazy_heavy_imports as lazy
|
||||
from data_designer.engine.resources.seed_reader import SeedReader
|
||||
|
||||
from .config import GitHubRepoSeedSource
|
||||
from .scraper import ScrapeConfig, materialize_to_jsonl
|
||||
|
||||
|
||||
# In-process cache mapping a stable config signature to the JSONL materialization
|
||||
# path. A single recipe job invokes the seed reader multiple times (validation,
|
||||
# preview, per-column sampling), and the default flow re-scrapes the repo on
|
||||
# every call: for a 2-repo preview that is ~15s of redundant GitHub GraphQL
|
||||
# traffic before any generation fires. Memoize the materialization so the second
|
||||
# and third passes reuse the file the first pass wrote. Cache key excludes the
|
||||
# raw token and uses a short SHA-256 digest so token values never hit memory
|
||||
# twice and token rotation invalidates cleanly.
|
||||
_SCRAPE_CACHE: dict[tuple, str] = {}
|
||||
_SCRAPE_CACHE_LOCK = threading.Lock()
|
||||
|
||||
|
||||
def _scrape_cache_key(cfg: ScrapeConfig) -> tuple:
|
||||
token_digest = hashlib.sha256(
|
||||
(cfg.token or "").encode("utf-8"),
|
||||
).hexdigest()[:16]
|
||||
return (
|
||||
tuple(cfg.repos),
|
||||
tuple(cfg.item_types),
|
||||
cfg.limit,
|
||||
bool(cfg.include_comments),
|
||||
cfg.max_comments_per_item,
|
||||
token_digest,
|
||||
)
|
||||
|
||||
|
||||
def _lookup_cached_scrape(key: tuple) -> Optional[str]:
|
||||
with _SCRAPE_CACHE_LOCK:
|
||||
path = _SCRAPE_CACHE.get(key)
|
||||
if path and Path(path).exists():
|
||||
return path
|
||||
# Stale entry (tmp cleanup, user restarted, ...); drop it so the caller
|
||||
# materializes a fresh file rather than returning a dangling path.
|
||||
if path:
|
||||
with _SCRAPE_CACHE_LOCK:
|
||||
_SCRAPE_CACHE.pop(key, None)
|
||||
return None
|
||||
|
||||
|
||||
def _store_cached_scrape(key: tuple, path: str) -> None:
|
||||
with _SCRAPE_CACHE_LOCK:
|
||||
_SCRAPE_CACHE[key] = path
|
||||
|
||||
|
||||
class GitHubRepoSeedReader(SeedReader[GitHubRepoSeedSource]):
|
||||
def create_duckdb_connection(self):
|
||||
return lazy.duckdb.connect()
|
||||
|
||||
def get_dataset_uri(self) -> str:
|
||||
out_dir = Path(tempfile.gettempdir()) / "studio-github-repo-seed"
|
||||
cfg = ScrapeConfig(
|
||||
repos = list(self.source.repos),
|
||||
token = self.source.token,
|
||||
item_types = list(self.source.item_types),
|
||||
limit = self.source.limit,
|
||||
include_comments = self.source.include_comments,
|
||||
max_comments_per_item = self.source.max_comments_per_item,
|
||||
)
|
||||
cache_key = _scrape_cache_key(cfg)
|
||||
cached_path = _lookup_cached_scrape(cache_key)
|
||||
if cached_path is not None:
|
||||
return cached_path
|
||||
path = materialize_to_jsonl(cfg, out_dir)
|
||||
_store_cached_scrape(cache_key, str(path))
|
||||
return str(path)
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from data_designer.plugins.plugin import Plugin, PluginType
|
||||
|
||||
github_repo_seed_plugin = Plugin(
|
||||
impl_qualified_name = "data_designer_github_repo_seed.impl.GitHubRepoSeedReader",
|
||||
config_qualified_name = "data_designer_github_repo_seed.config.GitHubRepoSeedSource",
|
||||
plugin_type = PluginType.SEED_READER,
|
||||
)
|
||||
|
|
@ -0,0 +1,236 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Multi-repo GitHub scraper for the Studio seed plugin.
|
||||
|
||||
Drives the GraphQL-based scraper in `scraper_impl/` per repo. Each repo is
|
||||
scraped with a trial_limits cap so we stop at `limit` items per resource.
|
||||
After scraping, we read the per-resource JSONL shards and flatten them into
|
||||
a single unified JSONL with stable columns (`item_type`, `repo`, `number`,
|
||||
`title`, `body`, ...).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
# Defer scraper_impl imports until `scrape()` runs with a resolved token.
|
||||
_IMPL_DIR = Path(__file__).parent / "scraper_impl"
|
||||
|
||||
|
||||
def _ensure_impl_on_path() -> None:
|
||||
if str(_IMPL_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_IMPL_DIR))
|
||||
|
||||
|
||||
def _load_impl():
|
||||
_ensure_impl_on_path()
|
||||
import importlib
|
||||
|
||||
gh_client = importlib.import_module("gh_client") # type: ignore
|
||||
scraper_mod = importlib.import_module("scraper") # type: ignore
|
||||
return gh_client.GitHubClient, scraper_mod.RepoScraper
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScrapeConfig:
|
||||
repos: list[str]
|
||||
token: str
|
||||
item_types: list[str]
|
||||
limit: int
|
||||
include_comments: bool
|
||||
max_comments_per_item: int
|
||||
|
||||
|
||||
def _resolve_token(token: str) -> str:
|
||||
tok = token or os.environ.get("GH_TOKEN", "") or os.environ.get("GITHUB_TOKEN", "")
|
||||
if not tok:
|
||||
raise ValueError(
|
||||
"GitHub token is required. Set it in the recipe config or the GH_TOKEN / GITHUB_TOKEN env var."
|
||||
)
|
||||
return tok
|
||||
|
||||
|
||||
def _read_jsonl(path: Path, max_rows: int | None = None):
|
||||
if not path.exists():
|
||||
return
|
||||
with path.open(encoding = "utf-8") as f:
|
||||
for i, line in enumerate(f):
|
||||
if not line.strip():
|
||||
continue
|
||||
if max_rows is not None and i >= max_rows:
|
||||
return
|
||||
try:
|
||||
yield json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
|
||||
def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
|
||||
labels = [
|
||||
l.get("name")
|
||||
for l in (r.get("labels", {}) or {}).get("nodes", [])
|
||||
if l.get("name")
|
||||
]
|
||||
comments_nodes = (r.get("comments") or {}).get("nodes") or []
|
||||
comments_text = ""
|
||||
if include_comments and comments_nodes:
|
||||
kept = comments_nodes[:max_c]
|
||||
comments_text = "\n\n".join(
|
||||
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}"
|
||||
for c in kept
|
||||
)
|
||||
return {
|
||||
"item_type": "issue",
|
||||
"repo": repo,
|
||||
"number": r.get("number"),
|
||||
"title": r.get("title") or "",
|
||||
"body": r.get("body") or "",
|
||||
"state": r.get("state") or "",
|
||||
"author": (r.get("author") or {}).get("login", ""),
|
||||
"created_at": r.get("createdAt") or "",
|
||||
"closed_at": r.get("closedAt") or "",
|
||||
"url": r.get("url") or r.get("permalink") or "",
|
||||
"labels": labels,
|
||||
"comments": comments_text,
|
||||
}
|
||||
|
||||
|
||||
def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict:
|
||||
labels = [
|
||||
l.get("name")
|
||||
for l in (r.get("labels", {}) or {}).get("nodes", [])
|
||||
if l.get("name")
|
||||
]
|
||||
comments_nodes = (r.get("comments") or {}).get("nodes") or []
|
||||
comments_text = ""
|
||||
if include_comments and comments_nodes:
|
||||
kept = comments_nodes[:max_c]
|
||||
comments_text = "\n\n".join(
|
||||
f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}"
|
||||
for c in kept
|
||||
)
|
||||
return {
|
||||
"item_type": "pull",
|
||||
"repo": repo,
|
||||
"number": r.get("number"),
|
||||
"title": r.get("title") or "",
|
||||
"body": r.get("body") or "",
|
||||
"state": r.get("state") or "",
|
||||
"author": (r.get("author") or {}).get("login", ""),
|
||||
"created_at": r.get("createdAt") or "",
|
||||
"closed_at": r.get("closedAt") or "",
|
||||
"url": r.get("url") or r.get("permalink") or "",
|
||||
"labels": labels,
|
||||
"comments": comments_text,
|
||||
}
|
||||
|
||||
|
||||
def _flatten_commit_row(r: dict, repo: str) -> dict:
|
||||
msg = r.get("messageHeadline") or r.get("message") or ""
|
||||
body = r.get("messageBody") or r.get("message") or msg
|
||||
author = r.get("author") or {}
|
||||
return {
|
||||
"item_type": "commit",
|
||||
"repo": repo,
|
||||
"number": r.get("oid") or r.get("sha") or "",
|
||||
"title": msg,
|
||||
"body": body,
|
||||
"state": "",
|
||||
"author": (author.get("user") or {}).get("login") or author.get("name", ""),
|
||||
"created_at": (author.get("date") or r.get("committedDate") or ""),
|
||||
"closed_at": "",
|
||||
"url": r.get("url") or "",
|
||||
"labels": [],
|
||||
"comments": "",
|
||||
}
|
||||
|
||||
|
||||
def scrape(cfg: ScrapeConfig, base_dir: Path):
|
||||
token = _resolve_token(cfg.token)
|
||||
GitHubClient, RepoScraper = _load_impl()
|
||||
client = GitHubClient(token = token)
|
||||
base_dir.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
# Per-resource trial limits. limit <= 0 means "all": use a very large cap.
|
||||
effective_limit = cfg.limit if cfg.limit and cfg.limit > 0 else 1_000_000
|
||||
trial_limits: dict[str, int] = {}
|
||||
if "issues" in cfg.item_types:
|
||||
trial_limits["issues"] = effective_limit
|
||||
if "pulls" in cfg.item_types:
|
||||
trial_limits["pull_requests"] = effective_limit
|
||||
if "commits" in cfg.item_types:
|
||||
trial_limits["commits"] = effective_limit
|
||||
|
||||
all_rows: list[dict] = []
|
||||
for repo in cfg.repos:
|
||||
owner, name = repo.split("/", 1)
|
||||
scraper = RepoScraper(
|
||||
owner = owner,
|
||||
name = name,
|
||||
base_dir = base_dir,
|
||||
client = client,
|
||||
trial_limits = trial_limits,
|
||||
light = True,
|
||||
)
|
||||
try:
|
||||
repo_meta = scraper.scrape_repo_meta()
|
||||
if "issues" in cfg.item_types:
|
||||
scraper.scrape_issues()
|
||||
if "pulls" in cfg.item_types:
|
||||
scraper.scrape_prs()
|
||||
if "commits" in cfg.item_types:
|
||||
default_ref = repo_meta.get("defaultBranchRef") or {}
|
||||
default_branch = (
|
||||
default_ref.get("name") if isinstance(default_ref, dict) else None
|
||||
)
|
||||
branch = (
|
||||
f"refs/heads/{default_branch}"
|
||||
if default_branch
|
||||
else "refs/heads/main"
|
||||
)
|
||||
scraper.scrape_commits(branch = branch)
|
||||
finally:
|
||||
scraper.close()
|
||||
|
||||
read_cap = cfg.limit if cfg.limit and cfg.limit > 0 else None
|
||||
repo_dir = base_dir / f"{owner}__{name}"
|
||||
if "issues" in cfg.item_types:
|
||||
for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap):
|
||||
all_rows.append(
|
||||
_flatten_issue_row(
|
||||
row, repo, cfg.include_comments, cfg.max_comments_per_item
|
||||
)
|
||||
)
|
||||
if "pulls" in cfg.item_types:
|
||||
for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap):
|
||||
all_rows.append(
|
||||
_flatten_pr_row(
|
||||
row, repo, cfg.include_comments, cfg.max_comments_per_item
|
||||
)
|
||||
)
|
||||
if "commits" in cfg.item_types:
|
||||
for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap):
|
||||
all_rows.append(_flatten_commit_row(row, repo))
|
||||
|
||||
return all_rows
|
||||
|
||||
|
||||
def materialize_to_jsonl(cfg: ScrapeConfig, out_dir: Path) -> Path:
|
||||
out_dir.mkdir(parents = True, exist_ok = True)
|
||||
tag = "-".join(r.replace("/", "__") for r in cfg.repos)[:120]
|
||||
kinds = "-".join(cfg.item_types)
|
||||
run_id = f"{int(time.time())}-{uuid.uuid4().hex[:12]}"
|
||||
fname = f"github_{tag}__{kinds}__{cfg.limit}_{run_id}.jsonl"
|
||||
out = out_dir / fname
|
||||
rows = scrape(cfg, out_dir / "raw-runs" / run_id)
|
||||
with out.open("w", encoding = "utf-8") as f:
|
||||
for r in rows:
|
||||
f.write(json.dumps(r, ensure_ascii = False) + "\n")
|
||||
return out
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
|
@ -0,0 +1,248 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""GitHub API client with rate-limit awareness, retry, and dual REST/GraphQL support."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import logging
|
||||
from typing import Any, Dict, Iterable, Iterator, List, Optional
|
||||
|
||||
import requests
|
||||
|
||||
log = logging.getLogger("gh_client")
|
||||
|
||||
GRAPHQL_URL = "https://api.github.com/graphql"
|
||||
REST_BASE = "https://api.github.com"
|
||||
|
||||
BASE_HEADERS = {
|
||||
"Accept": "application/vnd.github+json",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
"User-Agent": "github-data-gatherer/1.0",
|
||||
}
|
||||
|
||||
|
||||
class RateLimitError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class GitHubClient:
|
||||
def __init__(
|
||||
self,
|
||||
min_remaining_graphql: int = 100,
|
||||
min_remaining_rest: int = 100,
|
||||
token: str | None = None,
|
||||
):
|
||||
token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
|
||||
if not token:
|
||||
raise RuntimeError("GH_TOKEN not set in environment")
|
||||
self.session = requests.Session()
|
||||
self.session.headers.update(
|
||||
{**BASE_HEADERS, "Authorization": f"Bearer {token}"}
|
||||
)
|
||||
self.min_remaining_graphql = min_remaining_graphql
|
||||
self.min_remaining_rest = min_remaining_rest
|
||||
self.graphql_remaining: Optional[int] = None
|
||||
self.graphql_reset: Optional[int] = None
|
||||
self.rest_remaining: Optional[int] = None
|
||||
self.rest_reset: Optional[int] = None
|
||||
self.calls_graphql = 0
|
||||
self.calls_rest = 0
|
||||
self.retry_count = 0
|
||||
|
||||
def _sleep_until(self, reset_ts: int, buffer_s: int = 10) -> None:
|
||||
now = int(time.time())
|
||||
wait = max(0, reset_ts - now) + buffer_s
|
||||
log.warning("Rate limit hit. Sleeping %ds until reset.", wait)
|
||||
time.sleep(wait)
|
||||
|
||||
def _check_rate_and_wait(self, kind: str) -> None:
|
||||
if kind == "graphql":
|
||||
remaining = self.graphql_remaining
|
||||
reset = self.graphql_reset
|
||||
min_remaining = self.min_remaining_graphql
|
||||
else:
|
||||
remaining = self.rest_remaining
|
||||
reset = self.rest_reset
|
||||
min_remaining = self.min_remaining_rest
|
||||
if remaining is not None and remaining < min_remaining:
|
||||
if reset:
|
||||
self._sleep_until(reset)
|
||||
# Reset remaining so we don't spin
|
||||
if kind == "graphql":
|
||||
self.graphql_remaining = None
|
||||
else:
|
||||
self.rest_remaining = None
|
||||
|
||||
def graphql(
|
||||
self,
|
||||
query: str,
|
||||
variables: Optional[Dict[str, Any]] = None,
|
||||
max_retries: int = 20,
|
||||
) -> Dict[str, Any]:
|
||||
self._check_rate_and_wait("graphql")
|
||||
backoff = 2
|
||||
last_err = None
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
r = self.session.post(
|
||||
GRAPHQL_URL,
|
||||
json = {"query": query, "variables": variables or {}},
|
||||
timeout = 120,
|
||||
)
|
||||
self.calls_graphql += 1
|
||||
# Update rate info from response headers
|
||||
rem = r.headers.get("X-RateLimit-Remaining")
|
||||
rst = r.headers.get("X-RateLimit-Reset")
|
||||
if rem is not None:
|
||||
try:
|
||||
self.graphql_remaining = int(rem)
|
||||
except ValueError:
|
||||
pass
|
||||
if rst is not None:
|
||||
try:
|
||||
self.graphql_reset = int(rst)
|
||||
except ValueError:
|
||||
pass
|
||||
if r.status_code in (502, 503, 504):
|
||||
log.warning("GraphQL %s transient, retrying", r.status_code)
|
||||
time.sleep(backoff)
|
||||
backoff = min(backoff * 2, 60)
|
||||
continue
|
||||
if r.status_code == 403 or r.status_code == 429:
|
||||
# Check for secondary/abuse
|
||||
retry_after = r.headers.get("Retry-After")
|
||||
if retry_after:
|
||||
t = int(retry_after)
|
||||
log.warning("Secondary rate limit. Sleep %ds.", t)
|
||||
time.sleep(t + 2)
|
||||
continue
|
||||
if self.graphql_reset:
|
||||
self._sleep_until(self.graphql_reset)
|
||||
continue
|
||||
time.sleep(60)
|
||||
continue
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
if "errors" in data and data["errors"]:
|
||||
# Surface errors but allow partial data
|
||||
errs = data["errors"]
|
||||
# Retry on RATE_LIMITED
|
||||
for e in errs:
|
||||
if e.get("type") == "RATE_LIMITED":
|
||||
self._sleep_until(
|
||||
(self.graphql_reset or int(time.time()) + 60)
|
||||
)
|
||||
break
|
||||
else:
|
||||
# No rate-limit error, log and return partial
|
||||
log.warning("GraphQL errors: %s", json.dumps(errs)[:400])
|
||||
return data
|
||||
continue
|
||||
return data
|
||||
except requests.RequestException as e:
|
||||
last_err = e
|
||||
log.warning("GraphQL network error: %s. Retry.", e)
|
||||
time.sleep(backoff)
|
||||
backoff = min(backoff * 2, 60)
|
||||
raise RuntimeError(f"GraphQL failed after {max_retries} retries: {last_err}")
|
||||
|
||||
def rest(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
params: Optional[Dict[str, Any]] = None,
|
||||
json_body: Optional[Dict[str, Any]] = None,
|
||||
max_retries: int = 6,
|
||||
) -> requests.Response:
|
||||
self._check_rate_and_wait("rest")
|
||||
if path.startswith("http"):
|
||||
url = path
|
||||
else:
|
||||
url = REST_BASE + path
|
||||
backoff = 2
|
||||
last_err = None
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
r = self.session.request(
|
||||
method, url, params = params, json = json_body, timeout = 120
|
||||
)
|
||||
self.calls_rest += 1
|
||||
rem = r.headers.get("X-RateLimit-Remaining")
|
||||
rst = r.headers.get("X-RateLimit-Reset")
|
||||
if rem is not None:
|
||||
try:
|
||||
self.rest_remaining = int(rem)
|
||||
except ValueError:
|
||||
pass
|
||||
if rst is not None:
|
||||
try:
|
||||
self.rest_reset = int(rst)
|
||||
except ValueError:
|
||||
pass
|
||||
if r.status_code in (502, 503, 504):
|
||||
log.warning("REST %s transient, retrying", r.status_code)
|
||||
time.sleep(backoff)
|
||||
backoff = min(backoff * 2, 60)
|
||||
continue
|
||||
if r.status_code in (403, 429):
|
||||
retry_after = r.headers.get("Retry-After")
|
||||
if retry_after:
|
||||
t = int(retry_after)
|
||||
log.warning("Secondary rate limit on REST. Sleep %ds.", t)
|
||||
time.sleep(t + 2)
|
||||
continue
|
||||
# Check if primary rate
|
||||
if self.rest_remaining == 0 and self.rest_reset:
|
||||
self._sleep_until(self.rest_reset)
|
||||
continue
|
||||
log.warning("REST 403/429, sleep 60")
|
||||
time.sleep(60)
|
||||
continue
|
||||
return r
|
||||
except requests.RequestException as e:
|
||||
last_err = e
|
||||
log.warning("REST network error: %s. Retry.", e)
|
||||
time.sleep(backoff)
|
||||
backoff = min(backoff * 2, 60)
|
||||
raise RuntimeError(f"REST failed after {max_retries} retries: {last_err}")
|
||||
|
||||
def rest_paginate(
|
||||
self, path: str, params: Optional[Dict[str, Any]] = None, per_page: int = 100
|
||||
) -> Iterator[dict]:
|
||||
params = dict(params or {})
|
||||
params.setdefault("per_page", per_page)
|
||||
url = path
|
||||
while True:
|
||||
r = self.rest("GET", url, params = params if url == path else None)
|
||||
if r.status_code != 200:
|
||||
log.error(
|
||||
"REST paginate got %s at %s: %s", r.status_code, url, r.text[:200]
|
||||
)
|
||||
return
|
||||
items = r.json()
|
||||
if isinstance(items, dict):
|
||||
# Some endpoints return dict with list field
|
||||
items = items.get("items", [])
|
||||
for it in items:
|
||||
yield it
|
||||
# Follow link header
|
||||
link = r.headers.get("Link", "")
|
||||
nxt = None
|
||||
for part in link.split(","):
|
||||
if 'rel="next"' in part:
|
||||
nxt = part.split(";")[0].strip().strip("<>")
|
||||
break
|
||||
if not nxt:
|
||||
return
|
||||
url = nxt
|
||||
params = None
|
||||
|
||||
def rate_snapshot(self) -> Dict[str, Any]:
|
||||
r = self.rest("GET", "/rate_limit")
|
||||
if r.status_code == 200:
|
||||
return r.json()
|
||||
return {}
|
||||
|
|
@ -0,0 +1,685 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""GraphQL queries for GitHub data scraping.
|
||||
|
||||
GitHub's GraphQL rejects queries that define unused fragments, so each query
|
||||
only includes the fragments it actually references.
|
||||
"""
|
||||
|
||||
# ---- Fragments (kept as raw strings, composed per query) ----
|
||||
F_ACTOR = """
|
||||
fragment ActorFields on Actor {
|
||||
__typename
|
||||
login
|
||||
url
|
||||
avatarUrl
|
||||
... on User { id databaseId name }
|
||||
... on Bot { id databaseId }
|
||||
... on Organization { id databaseId name }
|
||||
}
|
||||
"""
|
||||
|
||||
F_LABEL = """
|
||||
fragment LabelFields on Label {
|
||||
id
|
||||
name
|
||||
color
|
||||
description
|
||||
createdAt
|
||||
}
|
||||
"""
|
||||
|
||||
F_TIMELINE = """
|
||||
fragment TimelineItem on IssueTimelineItems {
|
||||
__typename
|
||||
... on Node { id }
|
||||
... on AddedToProjectEvent { createdAt actor { ...ActorFields } }
|
||||
... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
|
||||
... on ClosedEvent { createdAt actor { ...ActorFields } stateReason closer { __typename ... on Commit { oid url } ... on PullRequest { number url } } }
|
||||
... on CommentDeletedEvent { createdAt actor { ...ActorFields } }
|
||||
... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url repository { nameWithOwner } } ... on PullRequest { number url repository { nameWithOwner } } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } }
|
||||
... on ConvertedNoteToIssueEvent { createdAt actor { ...ActorFields } }
|
||||
... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } }
|
||||
... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
|
||||
... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } }
|
||||
... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } }
|
||||
... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } }
|
||||
... on LockedEvent { createdAt actor { ...ActorFields } lockReason }
|
||||
... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } }
|
||||
... on MentionedEvent { createdAt actor { ...ActorFields } }
|
||||
... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
|
||||
... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } }
|
||||
... on PinnedEvent { createdAt actor { ...ActorFields } }
|
||||
... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } }
|
||||
... on RemovedFromProjectEvent { createdAt actor { ...ActorFields } }
|
||||
... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle }
|
||||
... on ReopenedEvent { createdAt actor { ...ActorFields } }
|
||||
... on SubscribedEvent { createdAt actor { ...ActorFields } }
|
||||
... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } }
|
||||
... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
|
||||
... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } }
|
||||
... on UnlockedEvent { createdAt actor { ...ActorFields } }
|
||||
... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } }
|
||||
... on UnpinnedEvent { createdAt actor { ...ActorFields } }
|
||||
... on UnsubscribedEvent { createdAt actor { ...ActorFields } }
|
||||
... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration }
|
||||
}
|
||||
"""
|
||||
|
||||
F_PR_TIMELINE = """
|
||||
fragment PRTimelineItem on PullRequestTimelineItems {
|
||||
__typename
|
||||
... on Node { id }
|
||||
... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
|
||||
... on AutoMergeDisabledEvent { createdAt actor { ...ActorFields } reason }
|
||||
... on AutoMergeEnabledEvent { createdAt actor { ...ActorFields } }
|
||||
... on AutoRebaseEnabledEvent { createdAt actor { ...ActorFields } }
|
||||
... on AutoSquashEnabledEvent { createdAt actor { ...ActorFields } }
|
||||
... on AutomaticBaseChangeFailedEvent { createdAt actor { ...ActorFields } oldBase newBase }
|
||||
... on AutomaticBaseChangeSucceededEvent { createdAt actor { ...ActorFields } oldBase newBase }
|
||||
... on BaseRefChangedEvent { createdAt actor { ...ActorFields } previousRefName currentRefName }
|
||||
... on BaseRefDeletedEvent { createdAt actor { ...ActorFields } baseRefName }
|
||||
... on BaseRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } }
|
||||
... on ClosedEvent { createdAt actor { ...ActorFields } stateReason }
|
||||
... on CommentDeletedEvent { createdAt actor { ...ActorFields } }
|
||||
... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url } ... on PullRequest { number url } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } }
|
||||
... on ConvertToDraftEvent { createdAt actor { ...ActorFields } }
|
||||
... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } }
|
||||
... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
|
||||
... on DeployedEvent { createdAt actor { ...ActorFields } }
|
||||
... on DeploymentEnvironmentChangedEvent { createdAt actor { ...ActorFields } }
|
||||
... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } }
|
||||
... on HeadRefDeletedEvent { createdAt actor { ...ActorFields } headRefName }
|
||||
... on HeadRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } }
|
||||
... on HeadRefRestoredEvent { createdAt actor { ...ActorFields } }
|
||||
... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } }
|
||||
... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } }
|
||||
... on LockedEvent { createdAt actor { ...ActorFields } lockReason }
|
||||
... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } }
|
||||
... on MentionedEvent { createdAt actor { ...ActorFields } }
|
||||
... on MergedEvent { createdAt actor { ...ActorFields } commit { oid url } mergeRefName }
|
||||
... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle }
|
||||
... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } }
|
||||
... on PinnedEvent { createdAt actor { ...ActorFields } }
|
||||
... on PullRequestCommit { commit { oid url message author { user { login } date } committedDate } }
|
||||
... on PullRequestCommitCommentThread { commit { oid } }
|
||||
... on PullRequestReview { id databaseId createdAt submittedAt author { ...ActorFields } body state url reactionGroups { content reactors { totalCount } } }
|
||||
... on PullRequestReviewThread { id isResolved isOutdated path line diffSide }
|
||||
... on PullRequestRevisionMarker { createdAt lastSeenCommit { oid } }
|
||||
... on ReadyForReviewEvent { createdAt actor { ...ActorFields } }
|
||||
... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } }
|
||||
... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle }
|
||||
... on ReopenedEvent { createdAt actor { ...ActorFields } }
|
||||
... on ReviewDismissedEvent { createdAt actor { ...ActorFields } dismissalMessage previousReviewState }
|
||||
... on ReviewRequestRemovedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } }
|
||||
... on ReviewRequestedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } }
|
||||
... on SubscribedEvent { createdAt actor { ...ActorFields } }
|
||||
... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } }
|
||||
... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } }
|
||||
... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } }
|
||||
... on UnlockedEvent { createdAt actor { ...ActorFields } }
|
||||
... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } }
|
||||
... on UnpinnedEvent { createdAt actor { ...ActorFields } }
|
||||
... on UnsubscribedEvent { createdAt actor { ...ActorFields } }
|
||||
... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration }
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
def _q(parts: list[str], body: str) -> str:
|
||||
return "\n".join(parts + [body])
|
||||
|
||||
|
||||
ISSUES_PAGE_QUERY = _q(
|
||||
[F_ACTOR, F_LABEL, F_TIMELINE],
|
||||
"""
|
||||
query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
totalCount
|
||||
nodes {
|
||||
id databaseId number title body state stateReason
|
||||
createdAt updatedAt closedAt
|
||||
url
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
labels(first: 50) { nodes { ...LabelFields } }
|
||||
assignees(first: 20) { nodes { login id } }
|
||||
milestone { title number state dueOn }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
comments(first: 100) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
timelineItems(first: 100) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes { ...TimelineItem }
|
||||
}
|
||||
trackedInIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } }
|
||||
trackedIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } }
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
PRS_PAGE_QUERY = _q(
|
||||
[F_ACTOR, F_LABEL, F_PR_TIMELINE],
|
||||
"""
|
||||
query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
totalCount
|
||||
nodes {
|
||||
id databaseId number title body state isDraft
|
||||
createdAt updatedAt closedAt mergedAt
|
||||
url
|
||||
headRefName headRefOid
|
||||
baseRefName baseRefOid
|
||||
additions deletions changedFiles
|
||||
mergeable merged mergeStateStatus
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
mergedBy { ...ActorFields }
|
||||
labels(first: 50) { nodes { ...LabelFields } }
|
||||
assignees(first: 20) { nodes { login id } }
|
||||
milestone { title number state dueOn }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
closingIssuesReferences(first: 20) { totalCount nodes { number url repository { nameWithOwner } title } }
|
||||
comments(first: 100) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
reviewThreads(first: 50) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id isResolved isOutdated path line diffSide
|
||||
comments(first: 50) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body path diffHunk
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
position originalPosition line originalLine
|
||||
commit { oid }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
reviews(first: 50) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId state createdAt submittedAt body url
|
||||
author { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
commits(first: 100) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
commit {
|
||||
oid
|
||||
message
|
||||
messageHeadline
|
||||
committedDate
|
||||
authoredDate
|
||||
author { name email user { login } date }
|
||||
committer { name email user { login } date }
|
||||
additions deletions changedFilesIfAvailable
|
||||
parents(first: 3) { nodes { oid } }
|
||||
}
|
||||
}
|
||||
}
|
||||
files(first: 100) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
path additions deletions changeType
|
||||
}
|
||||
}
|
||||
timelineItems(first: 100) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes { ...PRTimelineItem }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
PRS_PAGE_QUERY_LIGHT = _q(
|
||||
[F_ACTOR, F_LABEL],
|
||||
"""
|
||||
query PRsPageLight($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
totalCount
|
||||
nodes {
|
||||
id databaseId number title body state isDraft
|
||||
createdAt updatedAt closedAt mergedAt
|
||||
url
|
||||
author { ...ActorFields }
|
||||
labels(first: 50) { nodes { ...LabelFields } }
|
||||
comments(first: 30) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body
|
||||
author { ...ActorFields }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
ISSUES_PAGE_QUERY_LIGHT = _q(
|
||||
[F_ACTOR, F_LABEL],
|
||||
"""
|
||||
query IssuesPageLight($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
totalCount
|
||||
nodes {
|
||||
id databaseId number title body state
|
||||
createdAt updatedAt closedAt
|
||||
url
|
||||
author { ...ActorFields }
|
||||
labels(first: 50) { nodes { ...LabelFields } }
|
||||
comments(first: 30) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body
|
||||
author { ...ActorFields }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
ISSUE_COMMENTS_QUERY = _q(
|
||||
[F_ACTOR],
|
||||
"""
|
||||
query IssueComments($owner: String!, $name: String!, $number: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
issueOrPullRequest(number: $number) {
|
||||
__typename
|
||||
... on Issue {
|
||||
comments(first: 100, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
}
|
||||
... on PullRequest {
|
||||
comments(first: 100, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
ISSUE_TIMELINE_QUERY = _q(
|
||||
[F_ACTOR, F_TIMELINE],
|
||||
"""
|
||||
query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
issue(number: $number) {
|
||||
timelineItems(first: 100, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes { ...TimelineItem }
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
PR_TIMELINE_QUERY = _q(
|
||||
[F_ACTOR, F_PR_TIMELINE],
|
||||
"""
|
||||
query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
pullRequest(number: $number) {
|
||||
timelineItems(first: 100, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes { ...PRTimelineItem }
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
PR_COMMITS_QUERY = """
|
||||
query PRCommits($owner: String!, $name: String!, $number: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
pullRequest(number: $number) {
|
||||
commits(first: 100, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
commit {
|
||||
oid message messageHeadline committedDate authoredDate
|
||||
author { name email user { login } date }
|
||||
committer { name email user { login } date }
|
||||
additions deletions changedFilesIfAvailable
|
||||
parents(first: 3) { nodes { oid } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
"""
|
||||
|
||||
PR_FILES_QUERY = """
|
||||
query PRFiles($owner: String!, $name: String!, $number: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
pullRequest(number: $number) {
|
||||
files(first: 100, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes { path additions deletions changeType }
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
"""
|
||||
|
||||
PR_REVIEW_THREADS_QUERY = _q(
|
||||
[F_ACTOR],
|
||||
"""
|
||||
query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
pullRequest(number: $number) {
|
||||
reviewThreads(first: 50, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id isResolved isOutdated path line diffSide
|
||||
comments(first: 50) {
|
||||
totalCount
|
||||
nodes {
|
||||
id databaseId createdAt updatedAt url body path diffHunk
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
position originalPosition line originalLine
|
||||
commit { oid }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
DISCUSSIONS_PAGE_QUERY = _q(
|
||||
[F_ACTOR, F_LABEL],
|
||||
"""
|
||||
query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
discussions(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
totalCount
|
||||
nodes {
|
||||
id databaseId number title body
|
||||
createdAt updatedAt url
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
locked
|
||||
answerChosenAt
|
||||
closed closedAt
|
||||
category { id name emoji description isAnswerable }
|
||||
labels(first: 30) { nodes { ...LabelFields } }
|
||||
upvoteCount
|
||||
answer { id databaseId body author { ...ActorFields } createdAt url }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
comments(first: 50) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId body createdAt updatedAt url
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
upvoteCount
|
||||
isAnswer
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
replies(first: 50) {
|
||||
totalCount
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId body createdAt updatedAt url
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
DISCUSSION_COMMENTS_QUERY = _q(
|
||||
[F_ACTOR],
|
||||
"""
|
||||
query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
discussion(number: $number) {
|
||||
comments(first: 50, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId body createdAt updatedAt url
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
upvoteCount
|
||||
isAnswer
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
replies(first: 50) {
|
||||
totalCount
|
||||
nodes {
|
||||
id databaseId body createdAt updatedAt url
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
DISCUSSION_REPLIES_QUERY = _q(
|
||||
[F_ACTOR],
|
||||
"""
|
||||
query DiscussionReplies($commentId: ID!, $after: String) {
|
||||
node(id: $commentId) {
|
||||
... on DiscussionComment {
|
||||
replies(first: 50, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId body createdAt updatedAt url
|
||||
author { ...ActorFields }
|
||||
editor { ...ActorFields }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
COMMITS_PAGE_QUERY = """
|
||||
query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String, $branch: String!) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
ref(qualifiedName: $branch) {
|
||||
target {
|
||||
... on Commit {
|
||||
history(first: $first, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
totalCount
|
||||
nodes {
|
||||
oid
|
||||
message
|
||||
messageHeadline
|
||||
committedDate
|
||||
authoredDate
|
||||
url
|
||||
additions deletions changedFilesIfAvailable
|
||||
author { name email date user { login id } }
|
||||
committer { name email date user { login id } }
|
||||
parents(first: 3) { nodes { oid } }
|
||||
associatedPullRequests(first: 5) { nodes { number url state } }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
"""
|
||||
|
||||
RELEASES_QUERY = _q(
|
||||
[F_ACTOR],
|
||||
"""
|
||||
query Releases($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
releases(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id databaseId name tagName description
|
||||
createdAt publishedAt updatedAt
|
||||
isDraft isPrerelease isLatest
|
||||
url
|
||||
author { ...ActorFields }
|
||||
tagCommit { oid url }
|
||||
reactionGroups { content reactors { totalCount } }
|
||||
releaseAssets(first: 50) {
|
||||
nodes { name contentType size downloadUrl createdAt updatedAt }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
LABELS_QUERY = _q(
|
||||
[F_LABEL],
|
||||
"""
|
||||
query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
labels(first: $first, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes { ...LabelFields }
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
""",
|
||||
)
|
||||
|
||||
MILESTONES_QUERY = """
|
||||
query Milestones($owner: String!, $name: String!, $first: Int!, $after: String) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
milestones(first: $first, after: $after) {
|
||||
pageInfo { hasNextPage endCursor }
|
||||
nodes {
|
||||
id number title description state
|
||||
createdAt updatedAt closedAt dueOn
|
||||
creator { login }
|
||||
}
|
||||
}
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
"""
|
||||
|
||||
REPO_META_QUERY = """
|
||||
query RepoMeta($owner: String!, $name: String!) {
|
||||
repository(owner: $owner, name: $name) {
|
||||
id databaseId name nameWithOwner description url
|
||||
createdAt updatedAt pushedAt
|
||||
isArchived isDisabled isFork isPrivate
|
||||
primaryLanguage { name }
|
||||
languages(first: 20, orderBy: {field: SIZE, direction: DESC}) {
|
||||
edges { size node { name } }
|
||||
totalSize
|
||||
}
|
||||
stargazerCount forkCount watchers { totalCount }
|
||||
diskUsage
|
||||
licenseInfo { key name }
|
||||
homepageUrl
|
||||
defaultBranchRef { name }
|
||||
}
|
||||
rateLimit { cost remaining resetAt }
|
||||
}
|
||||
"""
|
||||
|
|
@ -0,0 +1,756 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Main scraper orchestration. Collects issues, PRs, discussions, commits, releases, etc.
|
||||
|
||||
Resumable via state file. Writes JSONL shards under data/{repo}/{resource}.jsonl.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
||||
|
||||
# Allow running as a module or script
|
||||
THIS_DIR = Path(__file__).resolve().parent
|
||||
if str(THIS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(THIS_DIR))
|
||||
|
||||
from gh_client import GitHubClient
|
||||
from state_store import JsonlWriter, StateStore
|
||||
import queries as Q
|
||||
|
||||
log = logging.getLogger("scraper")
|
||||
|
||||
|
||||
def ts() -> str:
|
||||
return time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
class RepoScraper:
|
||||
def __init__(
|
||||
self,
|
||||
owner: str,
|
||||
name: str,
|
||||
base_dir: Path,
|
||||
client: GitHubClient,
|
||||
trial_limits: Optional[Dict[str, int]] = None,
|
||||
light: bool = False,
|
||||
):
|
||||
self.owner = owner
|
||||
self.name = name
|
||||
self.base_dir = base_dir
|
||||
self.client = client
|
||||
self.trial_limits = trial_limits or {}
|
||||
# When light=True, use trimmed GraphQL queries (no reviewThreads,
|
||||
# reviews, commits, timelineItems, files) so PR pages can be much
|
||||
# larger without blowing GitHub's node-count ceiling.
|
||||
self.light = light
|
||||
self.repo_dir = base_dir / f"{owner}__{name}"
|
||||
self.repo_dir.mkdir(parents = True, exist_ok = True)
|
||||
self.state = StateStore(base_dir / "state" / f"{owner}__{name}.json")
|
||||
|
||||
# Writers
|
||||
self.writers: Dict[str, JsonlWriter] = {}
|
||||
for key in (
|
||||
"issues",
|
||||
"pull_requests",
|
||||
"discussions",
|
||||
"commits",
|
||||
"releases",
|
||||
"labels",
|
||||
"milestones",
|
||||
"pr_extra_comments",
|
||||
"pr_extra_timeline",
|
||||
"pr_extra_reviews",
|
||||
"issue_extra_comments",
|
||||
"issue_extra_timeline",
|
||||
"discussion_extra_comments",
|
||||
"discussion_extra_replies",
|
||||
"repo_meta",
|
||||
):
|
||||
self.writers[key] = JsonlWriter(self.repo_dir / f"{key}.jsonl")
|
||||
|
||||
# ----- helpers -----
|
||||
def _trial_stop(self, key: str, counter: int) -> bool:
|
||||
lim = self.trial_limits.get(key)
|
||||
if lim is None:
|
||||
return False
|
||||
return counter >= lim
|
||||
|
||||
def _log_rate(self, where: str, data: Dict[str, Any]) -> None:
|
||||
rl = (
|
||||
data.get("data", {}).get("rateLimit")
|
||||
if isinstance(data.get("data"), dict)
|
||||
else None
|
||||
)
|
||||
if rl:
|
||||
log.debug(
|
||||
"[%s] rate cost=%s remaining=%s resetAt=%s",
|
||||
where,
|
||||
rl.get("cost"),
|
||||
rl.get("remaining"),
|
||||
rl.get("resetAt"),
|
||||
)
|
||||
|
||||
# ----- repo meta -----
|
||||
def scrape_repo_meta(self) -> Dict[str, Any]:
|
||||
data = self.client.graphql(
|
||||
Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name}
|
||||
)
|
||||
self._log_rate("repo_meta", data)
|
||||
repo = data.get("data", {}).get("repository") or {}
|
||||
repo["_fetchedAt"] = ts()
|
||||
self.writers["repo_meta"].write(repo)
|
||||
return repo
|
||||
|
||||
# ----- issues -----
|
||||
def scrape_issues(self) -> int:
|
||||
key = "issues"
|
||||
cursor = self.state.get(f"{key}_cursor")
|
||||
done = self.state.get(f"{key}_done", False)
|
||||
if done:
|
||||
log.info("%s/%s issues already complete", self.owner, self.name)
|
||||
return 0
|
||||
total_new = 0
|
||||
page = 0
|
||||
# Light query skips heavy nested fields; safe at 50 per page.
|
||||
# Clamp by trial_limit so e.g. limit=1 asks GitHub for first:1
|
||||
# instead of fetching a full 50-item page and discarding 49.
|
||||
page_cap = 50 if self.light else 15
|
||||
trial_cap = self.trial_limits.get(key)
|
||||
per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap
|
||||
while True:
|
||||
page += 1
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"first": per_page,
|
||||
"after": cursor,
|
||||
}
|
||||
query = Q.ISSUES_PAGE_QUERY_LIGHT if self.light else Q.ISSUES_PAGE_QUERY
|
||||
data = self.client.graphql(query, vars_)
|
||||
self._log_rate("issues", data)
|
||||
repo = (data.get("data") or {}).get("repository") or {}
|
||||
issues = repo.get("issues") or {}
|
||||
nodes = issues.get("nodes") or []
|
||||
for it in nodes:
|
||||
it["_owner"] = self.owner
|
||||
it["_repo"] = self.name
|
||||
it["_fetchedAt"] = ts()
|
||||
if not self.light:
|
||||
if it.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
|
||||
self._paginate_issue_comments(
|
||||
it["number"], it["comments"]["pageInfo"]["endCursor"]
|
||||
)
|
||||
if (
|
||||
it.get("timelineItems", {})
|
||||
.get("pageInfo", {})
|
||||
.get("hasNextPage")
|
||||
):
|
||||
self._paginate_issue_timeline(
|
||||
it["number"],
|
||||
it["timelineItems"]["pageInfo"]["endCursor"],
|
||||
)
|
||||
if self.writers[key].write(it):
|
||||
total_new += 1
|
||||
info = issues.get("pageInfo") or {}
|
||||
cursor = info.get("endCursor")
|
||||
self.state.set(f"{key}_cursor", cursor)
|
||||
log.info(
|
||||
"[%s/%s] issues page %d (+%d) cursor=%s remaining=%s",
|
||||
self.owner,
|
||||
self.name,
|
||||
page,
|
||||
len(nodes),
|
||||
str(cursor)[:20],
|
||||
self.client.graphql_remaining,
|
||||
)
|
||||
if self._trial_stop(key, total_new):
|
||||
log.info("Trial limit reached for issues (%d)", total_new)
|
||||
return total_new
|
||||
if not info.get("hasNextPage"):
|
||||
self.state.set(f"{key}_done", True)
|
||||
break
|
||||
return total_new
|
||||
|
||||
def _paginate_issue_comments(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_)
|
||||
item = ((data.get("data") or {}).get("repository") or {}).get(
|
||||
"issueOrPullRequest"
|
||||
) or {}
|
||||
comments = item.get("comments") or {}
|
||||
for c in comments.get("nodes") or []:
|
||||
c["_owner"] = self.owner
|
||||
c["_repo"] = self.name
|
||||
c["_issueNumber"] = number
|
||||
self.writers["issue_extra_comments"].write(c)
|
||||
info = comments.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
def _paginate_issue_timeline(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.ISSUE_TIMELINE_QUERY, vars_)
|
||||
item = ((data.get("data") or {}).get("repository") or {}).get("issue") or {}
|
||||
tl = item.get("timelineItems") or {}
|
||||
for ev in tl.get("nodes") or []:
|
||||
ev["_owner"] = self.owner
|
||||
ev["_repo"] = self.name
|
||||
ev["_issueNumber"] = number
|
||||
self.writers["issue_extra_timeline"].write(ev)
|
||||
info = tl.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
# ----- PRs -----
|
||||
def scrape_prs(self) -> int:
|
||||
key = "pull_requests"
|
||||
cursor = self.state.get(f"{key}_cursor")
|
||||
done = self.state.get(f"{key}_done", False)
|
||||
if done:
|
||||
log.info("%s/%s PRs already complete", self.owner, self.name)
|
||||
return 0
|
||||
total_new = 0
|
||||
page = 0
|
||||
# Heavy nested PR query is capped at 3 per page (GitHub node-count
|
||||
# ceiling); light query skips reviewThreads/reviews/commits/etc and
|
||||
# can safely go to 25 per page. Clamp by trial_limit for small
|
||||
# previews so limit=1 does not fetch a whole 25-item page.
|
||||
page_cap = 25 if self.light else 3
|
||||
trial_cap = self.trial_limits.get(key)
|
||||
per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap
|
||||
while True:
|
||||
page += 1
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"first": per_page,
|
||||
"after": cursor,
|
||||
}
|
||||
query = Q.PRS_PAGE_QUERY_LIGHT if self.light else Q.PRS_PAGE_QUERY
|
||||
data = self.client.graphql(query, vars_)
|
||||
self._log_rate("prs", data)
|
||||
repo = (data.get("data") or {}).get("repository") or {}
|
||||
prs = repo.get("pullRequests") or {}
|
||||
nodes = prs.get("nodes") or []
|
||||
for pr in nodes:
|
||||
pr["_owner"] = self.owner
|
||||
pr["_repo"] = self.name
|
||||
pr["_fetchedAt"] = ts()
|
||||
num = pr["number"]
|
||||
if not self.light:
|
||||
if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
|
||||
self._paginate_pr_comments(
|
||||
num, pr["comments"]["pageInfo"]["endCursor"]
|
||||
)
|
||||
if (
|
||||
pr.get("timelineItems", {})
|
||||
.get("pageInfo", {})
|
||||
.get("hasNextPage")
|
||||
):
|
||||
self._paginate_pr_timeline(
|
||||
num, pr["timelineItems"]["pageInfo"]["endCursor"]
|
||||
)
|
||||
if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"):
|
||||
self._paginate_pr_commits(
|
||||
num, pr["commits"]["pageInfo"]["endCursor"]
|
||||
)
|
||||
if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"):
|
||||
self._paginate_pr_files(
|
||||
num, pr["files"]["pageInfo"]["endCursor"]
|
||||
)
|
||||
if (
|
||||
pr.get("reviewThreads", {})
|
||||
.get("pageInfo", {})
|
||||
.get("hasNextPage")
|
||||
):
|
||||
self._paginate_pr_review_threads(
|
||||
num, pr["reviewThreads"]["pageInfo"]["endCursor"]
|
||||
)
|
||||
if self.writers[key].write(pr):
|
||||
total_new += 1
|
||||
info = prs.get("pageInfo") or {}
|
||||
cursor = info.get("endCursor")
|
||||
self.state.set(f"{key}_cursor", cursor)
|
||||
log.info(
|
||||
"[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s",
|
||||
self.owner,
|
||||
self.name,
|
||||
page,
|
||||
len(nodes),
|
||||
str(cursor)[:20],
|
||||
self.client.graphql_remaining,
|
||||
)
|
||||
if self._trial_stop(key, total_new):
|
||||
log.info("Trial limit reached for PRs (%d)", total_new)
|
||||
return total_new
|
||||
if not info.get("hasNextPage"):
|
||||
self.state.set(f"{key}_done", True)
|
||||
break
|
||||
return total_new
|
||||
|
||||
def _paginate_pr_comments(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_)
|
||||
item = ((data.get("data") or {}).get("repository") or {}).get(
|
||||
"issueOrPullRequest"
|
||||
) or {}
|
||||
comments = item.get("comments") or {}
|
||||
for c in comments.get("nodes") or []:
|
||||
c["_owner"] = self.owner
|
||||
c["_repo"] = self.name
|
||||
c["_prNumber"] = number
|
||||
self.writers["pr_extra_comments"].write(c)
|
||||
info = comments.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
def _paginate_pr_timeline(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_)
|
||||
item = ((data.get("data") or {}).get("repository") or {}).get(
|
||||
"pullRequest"
|
||||
) or {}
|
||||
tl = item.get("timelineItems") or {}
|
||||
for ev in tl.get("nodes") or []:
|
||||
ev["_owner"] = self.owner
|
||||
ev["_repo"] = self.name
|
||||
ev["_prNumber"] = number
|
||||
self.writers["pr_extra_timeline"].write(ev)
|
||||
info = tl.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
def _paginate_pr_commits(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
out_key = "pr_extra_commits"
|
||||
if out_key not in self.writers:
|
||||
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_)
|
||||
item = ((data.get("data") or {}).get("repository") or {}).get(
|
||||
"pullRequest"
|
||||
) or {}
|
||||
cc = item.get("commits") or {}
|
||||
for c in cc.get("nodes") or []:
|
||||
c["_owner"] = self.owner
|
||||
c["_repo"] = self.name
|
||||
c["_prNumber"] = number
|
||||
self.writers[out_key].write(c)
|
||||
info = cc.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
def _paginate_pr_files(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
out_key = "pr_extra_files"
|
||||
if out_key not in self.writers:
|
||||
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.PR_FILES_QUERY, vars_)
|
||||
item = ((data.get("data") or {}).get("repository") or {}).get(
|
||||
"pullRequest"
|
||||
) or {}
|
||||
ff = item.get("files") or {}
|
||||
for f in ff.get("nodes") or []:
|
||||
f["_owner"] = self.owner
|
||||
f["_repo"] = self.name
|
||||
f["_prNumber"] = number
|
||||
# files don't have id, synthesize one
|
||||
f["_syntheticId"] = f"{self.owner}/{self.name}#{number}:{f.get('path')}"
|
||||
self.writers[out_key].write(f)
|
||||
info = ff.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
def _paginate_pr_review_threads(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
out_key = "pr_extra_review_threads"
|
||||
if out_key not in self.writers:
|
||||
self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl")
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_)
|
||||
item = ((data.get("data") or {}).get("repository") or {}).get(
|
||||
"pullRequest"
|
||||
) or {}
|
||||
rt = item.get("reviewThreads") or {}
|
||||
for th in rt.get("nodes") or []:
|
||||
th["_owner"] = self.owner
|
||||
th["_repo"] = self.name
|
||||
th["_prNumber"] = number
|
||||
self.writers[out_key].write(th)
|
||||
info = rt.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
# ----- Discussions -----
|
||||
def scrape_discussions(self) -> int:
|
||||
key = "discussions"
|
||||
cursor = self.state.get(f"{key}_cursor")
|
||||
done = self.state.get(f"{key}_done", False)
|
||||
if done:
|
||||
log.info("%s/%s discussions already complete", self.owner, self.name)
|
||||
return 0
|
||||
total_new = 0
|
||||
page = 0
|
||||
per_page = 15
|
||||
while True:
|
||||
page += 1
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"first": per_page,
|
||||
"after": cursor,
|
||||
}
|
||||
data = self.client.graphql(Q.DISCUSSIONS_PAGE_QUERY, vars_)
|
||||
self._log_rate("discussions", data)
|
||||
repo = (data.get("data") or {}).get("repository") or {}
|
||||
dd = repo.get("discussions") or {}
|
||||
nodes = dd.get("nodes") or []
|
||||
for d in nodes:
|
||||
d["_owner"] = self.owner
|
||||
d["_repo"] = self.name
|
||||
d["_fetchedAt"] = ts()
|
||||
num = d["number"]
|
||||
if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"):
|
||||
self._paginate_discussion_comments(
|
||||
num, d["comments"]["pageInfo"]["endCursor"]
|
||||
)
|
||||
# paginate replies per comment if needed
|
||||
for c in d.get("comments", {}).get("nodes", []) or []:
|
||||
if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"):
|
||||
self._paginate_discussion_replies(
|
||||
c["id"], c["replies"]["pageInfo"]["endCursor"], num
|
||||
)
|
||||
if self.writers[key].write(d):
|
||||
total_new += 1
|
||||
info = dd.get("pageInfo") or {}
|
||||
cursor = info.get("endCursor")
|
||||
self.state.set(f"{key}_cursor", cursor)
|
||||
log.info(
|
||||
"[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s",
|
||||
self.owner,
|
||||
self.name,
|
||||
page,
|
||||
len(nodes),
|
||||
str(cursor)[:20],
|
||||
self.client.graphql_remaining,
|
||||
)
|
||||
if self._trial_stop(key, total_new):
|
||||
return total_new
|
||||
if not info.get("hasNextPage"):
|
||||
self.state.set(f"{key}_done", True)
|
||||
break
|
||||
return total_new
|
||||
|
||||
def _paginate_discussion_comments(self, number: int, after: str) -> None:
|
||||
cur = after
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"number": number,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_)
|
||||
disc = ((data.get("data") or {}).get("repository") or {}).get(
|
||||
"discussion"
|
||||
) or {}
|
||||
cc = disc.get("comments") or {}
|
||||
for c in cc.get("nodes") or []:
|
||||
c["_owner"] = self.owner
|
||||
c["_repo"] = self.name
|
||||
c["_discussionNumber"] = number
|
||||
self.writers["discussion_extra_comments"].write(c)
|
||||
info = cc.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
def _paginate_discussion_replies(
|
||||
self, comment_id: str, after: str, disc_number: int
|
||||
) -> None:
|
||||
cur = after
|
||||
while cur:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"commentId": comment_id,
|
||||
"after": cur,
|
||||
}
|
||||
data = self.client.graphql(Q.DISCUSSION_REPLIES_QUERY, vars_)
|
||||
node = (data.get("data") or {}).get("node") or {}
|
||||
replies = node.get("replies") or {}
|
||||
for r in replies.get("nodes") or []:
|
||||
r["_owner"] = self.owner
|
||||
r["_repo"] = self.name
|
||||
r["_discussionNumber"] = disc_number
|
||||
r["_commentId"] = comment_id
|
||||
self.writers["discussion_extra_replies"].write(r)
|
||||
info = replies.get("pageInfo") or {}
|
||||
cur = info.get("endCursor") if info.get("hasNextPage") else None
|
||||
|
||||
# ----- Commits -----
|
||||
def scrape_commits(self, branch: str = "refs/heads/main") -> int:
|
||||
key = "commits"
|
||||
cursor = self.state.get(f"{key}_cursor")
|
||||
done = self.state.get(f"{key}_done", False)
|
||||
if done:
|
||||
return 0
|
||||
total_new = 0
|
||||
page = 0
|
||||
page_cap = 100
|
||||
trial_cap = self.trial_limits.get(key)
|
||||
per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap
|
||||
while True:
|
||||
page += 1
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"first": per_page,
|
||||
"after": cursor,
|
||||
"branch": branch,
|
||||
}
|
||||
data = self.client.graphql(Q.COMMITS_PAGE_QUERY, vars_)
|
||||
self._log_rate("commits", data)
|
||||
ref = ((data.get("data") or {}).get("repository") or {}).get("ref") or {}
|
||||
tgt = ref.get("target") or {}
|
||||
hist = tgt.get("history") or {}
|
||||
nodes = hist.get("nodes") or []
|
||||
for c in nodes:
|
||||
c["_owner"] = self.owner
|
||||
c["_repo"] = self.name
|
||||
c["_fetchedAt"] = ts()
|
||||
if self.writers[key].write(c):
|
||||
total_new += 1
|
||||
info = hist.get("pageInfo") or {}
|
||||
cursor = info.get("endCursor")
|
||||
self.state.set(f"{key}_cursor", cursor)
|
||||
log.info(
|
||||
"[%s/%s] commits page %d (+%d) remaining=%s",
|
||||
self.owner,
|
||||
self.name,
|
||||
page,
|
||||
len(nodes),
|
||||
self.client.graphql_remaining,
|
||||
)
|
||||
if self._trial_stop(key, total_new):
|
||||
return total_new
|
||||
if not info.get("hasNextPage"):
|
||||
self.state.set(f"{key}_done", True)
|
||||
break
|
||||
return total_new
|
||||
|
||||
# ----- Releases/Labels/Milestones -----
|
||||
def scrape_releases(self) -> int:
|
||||
return self._scrape_simple("releases", Q.RELEASES_QUERY, "releases")
|
||||
|
||||
def scrape_labels(self) -> int:
|
||||
return self._scrape_simple("labels", Q.LABELS_QUERY, "labels")
|
||||
|
||||
def scrape_milestones(self) -> int:
|
||||
return self._scrape_simple("milestones", Q.MILESTONES_QUERY, "milestones")
|
||||
|
||||
def _scrape_simple(self, key: str, query: str, field: str) -> int:
|
||||
cursor = self.state.get(f"{key}_cursor")
|
||||
done = self.state.get(f"{key}_done", False)
|
||||
if done:
|
||||
return 0
|
||||
total_new = 0
|
||||
while True:
|
||||
vars_ = {
|
||||
"owner": self.owner,
|
||||
"name": self.name,
|
||||
"first": 50,
|
||||
"after": cursor,
|
||||
}
|
||||
data = self.client.graphql(query, vars_)
|
||||
repo = (data.get("data") or {}).get("repository") or {}
|
||||
col = repo.get(field) or {}
|
||||
for it in col.get("nodes") or []:
|
||||
it["_owner"] = self.owner
|
||||
it["_repo"] = self.name
|
||||
it["_fetchedAt"] = ts()
|
||||
if self.writers[key].write(it):
|
||||
total_new += 1
|
||||
info = col.get("pageInfo") or {}
|
||||
cursor = info.get("endCursor")
|
||||
self.state.set(f"{key}_cursor", cursor)
|
||||
if self._trial_stop(key, total_new):
|
||||
return total_new
|
||||
if not info.get("hasNextPage"):
|
||||
self.state.set(f"{key}_done", True)
|
||||
break
|
||||
log.info("[%s/%s] %s done +%d", self.owner, self.name, key, total_new)
|
||||
return total_new
|
||||
|
||||
def close(self) -> None:
|
||||
for w in self.writers.values():
|
||||
try:
|
||||
w.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def setup_logging(log_file: Path) -> None:
|
||||
log_file.parent.mkdir(parents = True, exist_ok = True)
|
||||
fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
|
||||
handlers = [
|
||||
logging.StreamHandler(sys.stdout),
|
||||
logging.FileHandler(log_file, mode = "a", encoding = "utf-8"),
|
||||
]
|
||||
logging.basicConfig(level = logging.INFO, format = fmt, handlers = handlers, force = True)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument(
|
||||
"--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper"
|
||||
)
|
||||
ap.add_argument(
|
||||
"--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"]
|
||||
)
|
||||
ap.add_argument("--trial", action = "store_true", help = "Small trial run")
|
||||
ap.add_argument(
|
||||
"--only",
|
||||
nargs = "+",
|
||||
default = None,
|
||||
help = "Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--hf-upload-interval",
|
||||
type = int,
|
||||
default = 900,
|
||||
help = "Seconds between HF uploads (0 to disable)",
|
||||
)
|
||||
args = ap.parse_args()
|
||||
|
||||
base = Path(args.base_dir)
|
||||
data_dir = base / "data"
|
||||
data_dir.mkdir(parents = True, exist_ok = True)
|
||||
setup_logging(base / "logs" / f"scraper_{time.strftime('%Y%m%d_%H%M%S')}.log")
|
||||
log.info("Scraper starting: repos=%s trial=%s", args.repos, args.trial)
|
||||
|
||||
client = GitHubClient(min_remaining_graphql = 80, min_remaining_rest = 80)
|
||||
rl = client.rate_snapshot()
|
||||
log.info(
|
||||
"Rate limit snapshot: %s",
|
||||
json.dumps(rl.get("resources", {}), default = str)[:400],
|
||||
)
|
||||
|
||||
# Start HF uploader in background if requested
|
||||
uploader = None
|
||||
if args.hf_upload_interval > 0:
|
||||
from hf_uploader import HFUploader
|
||||
|
||||
uploader = HFUploader(data_dir, interval_s = args.hf_upload_interval)
|
||||
uploader.start()
|
||||
|
||||
trial_limits = None
|
||||
if args.trial:
|
||||
trial_limits = {
|
||||
"issues": 5,
|
||||
"pull_requests": 5,
|
||||
"discussions": 3,
|
||||
"commits": 20,
|
||||
"releases": 3,
|
||||
"labels": 20,
|
||||
"milestones": 20,
|
||||
}
|
||||
|
||||
only = set(args.only or [])
|
||||
|
||||
try:
|
||||
for repo_spec in args.repos:
|
||||
owner, name = repo_spec.split("/")
|
||||
scraper = RepoScraper(owner, name, data_dir, client, trial_limits)
|
||||
try:
|
||||
repo_meta: Dict[str, Any] = {}
|
||||
if not only or "meta" in only or "commits" in only:
|
||||
repo_meta = scraper.scrape_repo_meta()
|
||||
if not only or "labels" in only:
|
||||
scraper.scrape_labels()
|
||||
if not only or "milestones" in only:
|
||||
scraper.scrape_milestones()
|
||||
if not only or "releases" in only:
|
||||
scraper.scrape_releases()
|
||||
if not only or "discussions" in only:
|
||||
scraper.scrape_discussions()
|
||||
if not only or "issues" in only:
|
||||
scraper.scrape_issues()
|
||||
if not only or "pulls" in only:
|
||||
scraper.scrape_prs()
|
||||
if not only or "commits" in only:
|
||||
default_ref = repo_meta.get("defaultBranchRef") or {}
|
||||
default_branch = (
|
||||
default_ref.get("name")
|
||||
if isinstance(default_ref, dict)
|
||||
else None
|
||||
)
|
||||
branch = (
|
||||
f"refs/heads/{default_branch}"
|
||||
if default_branch
|
||||
else "refs/heads/main"
|
||||
)
|
||||
scraper.scrape_commits(branch = branch)
|
||||
finally:
|
||||
scraper.close()
|
||||
finally:
|
||||
if uploader:
|
||||
log.info("Stopping uploader and final sync...")
|
||||
uploader.stop(final_upload = True)
|
||||
log.info(
|
||||
"Scraper complete. GraphQL calls=%d REST calls=%d",
|
||||
client.calls_graphql,
|
||||
client.calls_rest,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,105 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Checkpoint state management for resumable scraping."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
|
||||
|
||||
class StateStore:
|
||||
def __init__(self, path: str | Path):
|
||||
self.path = Path(path)
|
||||
self.path.parent.mkdir(parents = True, exist_ok = True)
|
||||
self._lock = threading.Lock()
|
||||
self._data: Dict[str, Any] = {}
|
||||
if self.path.exists():
|
||||
try:
|
||||
with self.path.open() as f:
|
||||
self._data = json.load(f)
|
||||
except Exception:
|
||||
self._data = {}
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
with self._lock:
|
||||
return self._data.get(key, default)
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
with self._lock:
|
||||
self._data[key] = value
|
||||
self._flush()
|
||||
|
||||
def update(self, key: str, **kwargs) -> None:
|
||||
with self._lock:
|
||||
sub = dict(self._data.get(key, {}))
|
||||
sub.update(kwargs)
|
||||
self._data[key] = sub
|
||||
self._flush()
|
||||
|
||||
def all(self) -> Dict[str, Any]:
|
||||
with self._lock:
|
||||
return dict(self._data)
|
||||
|
||||
def _flush(self) -> None:
|
||||
tmp = self.path.with_suffix(self.path.suffix + ".tmp")
|
||||
with tmp.open("w") as f:
|
||||
json.dump(self._data, f, indent = 2, default = str)
|
||||
os.replace(tmp, self.path)
|
||||
|
||||
|
||||
class JsonlWriter:
|
||||
"""Append-only JSONL writer, thread-safe, with line buffering."""
|
||||
|
||||
def __init__(self, path: str | Path):
|
||||
self.path = Path(path)
|
||||
self.path.parent.mkdir(parents = True, exist_ok = True)
|
||||
self._lock = threading.Lock()
|
||||
self._fh = self.path.open("a", buffering = 1)
|
||||
self._count_seen_keys: set[str] = set()
|
||||
# Preload seen keys if file exists (for dedup across resumes)
|
||||
if self.path.exists() and self.path.stat().st_size > 0:
|
||||
try:
|
||||
with self.path.open() as f:
|
||||
for line in f:
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
k = self._key(obj)
|
||||
if k is not None:
|
||||
self._count_seen_keys.add(k)
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _key(self, obj: dict) -> str | None:
|
||||
for k in ("id", "node_id", "number", "sha", "url"):
|
||||
if k in obj:
|
||||
return f"{k}:{obj[k]}"
|
||||
return None
|
||||
|
||||
def has(self, key: str) -> bool:
|
||||
return key in self._count_seen_keys
|
||||
|
||||
def write(self, obj: dict) -> bool:
|
||||
"""Return True if newly written, False if already present."""
|
||||
k = self._key(obj)
|
||||
with self._lock:
|
||||
if k is not None and k in self._count_seen_keys:
|
||||
return False
|
||||
if k is not None:
|
||||
self._count_seen_keys.add(k)
|
||||
self._fh.write(json.dumps(obj, default = str, ensure_ascii = False))
|
||||
self._fh.write("\n")
|
||||
self._fh.flush()
|
||||
return True
|
||||
|
||||
def close(self) -> None:
|
||||
try:
|
||||
self._fh.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
|
@ -19,7 +19,8 @@ ruff<1,>=0.14.10
|
|||
scipy<2,>=1.11.0
|
||||
sqlfluff<4,>=3.2.0
|
||||
tiktoken<1,>=0.8.0
|
||||
# Unstructured-seed plugin deps (plugin installed with --no-deps)
|
||||
# Local seed plugin deps (plugins installed with --no-deps)
|
||||
requests>=2.31
|
||||
pymupdf>=1.24.0
|
||||
pymupdf4llm>=0.0.17
|
||||
mammoth>=1.8.0
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ API Routes
|
|||
from routes.training import router as training_router
|
||||
from routes.models import router as models_router
|
||||
from routes.inference import router as inference_router
|
||||
from routes.inference import studio_router as inference_studio_router
|
||||
from routes.datasets import router as datasets_router
|
||||
from routes.auth import router as auth_router
|
||||
from routes.data_recipe import router as data_recipe_router
|
||||
|
|
@ -18,6 +19,7 @@ __all__ = [
|
|||
"training_router",
|
||||
"models_router",
|
||||
"inference_router",
|
||||
"inference_studio_router",
|
||||
"datasets_router",
|
||||
"auth_router",
|
||||
"data_recipe_router",
|
||||
|
|
|
|||
|
|
@ -5,8 +5,9 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
import copy
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query, Request
|
||||
|
|
@ -94,14 +95,111 @@ def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]:
|
|||
return aliases
|
||||
|
||||
|
||||
def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None:
|
||||
def _inject_local_structured_response_format(
|
||||
recipe: dict[str, Any], local_provider_names: set[str]
|
||||
) -> None:
|
||||
"""For each llm-structured column that targets a local-provider model_config,
|
||||
clone the model_config and inject an OpenAI ``response_format`` with the
|
||||
column's ``output_format`` JSON schema. The column is rewritten to point at
|
||||
the clone so llm-text / llm-judge columns that share the same alias keep
|
||||
free-form sampling.
|
||||
|
||||
Without this, data_designer only injects a prompt-level "return JSON in a
|
||||
```json fence" instruction. Small GGUF models frequently break format,
|
||||
wasting the full ``max_tokens`` budget per row and then failing to parse.
|
||||
Forwarding ``response_format`` lets llama-server apply grammar-constrained
|
||||
sampling from the JSON schema, which guarantees a parseable response and
|
||||
terminates early.
|
||||
"""
|
||||
columns = recipe.get("columns")
|
||||
model_configs = recipe.get("model_configs")
|
||||
if not isinstance(columns, list) or not isinstance(model_configs, list):
|
||||
return
|
||||
|
||||
# alias -> model_config (only configs referencing a local provider qualify).
|
||||
alias_to_local_mc: dict[str, dict[str, Any]] = {}
|
||||
for mc in model_configs:
|
||||
if not isinstance(mc, dict):
|
||||
continue
|
||||
if mc.get("provider") in local_provider_names and isinstance(
|
||||
mc.get("alias"), str
|
||||
):
|
||||
alias_to_local_mc[mc["alias"]] = mc
|
||||
|
||||
if not alias_to_local_mc:
|
||||
return
|
||||
|
||||
# Clone per (alias, column) so each llm-structured column gets its own
|
||||
# schema without leaking response_format onto other columns that share the
|
||||
# same base alias.
|
||||
seen_clone_aliases: set[str] = {
|
||||
mc.get("alias") for mc in model_configs if isinstance(mc.get("alias"), str)
|
||||
}
|
||||
new_configs: list[dict[str, Any]] = []
|
||||
for column in columns:
|
||||
if not isinstance(column, dict):
|
||||
continue
|
||||
if column.get("column_type") != "llm-structured":
|
||||
continue
|
||||
alias = column.get("model_alias")
|
||||
if not isinstance(alias, str) or alias not in alias_to_local_mc:
|
||||
continue
|
||||
output_format = column.get("output_format")
|
||||
if not isinstance(output_format, dict) or not output_format:
|
||||
continue
|
||||
base_mc = alias_to_local_mc[alias]
|
||||
column_name = column.get("name") or "structured"
|
||||
clone_alias_base = f"{alias}__{column_name}_structured"
|
||||
clone_alias = clone_alias_base
|
||||
counter = 1
|
||||
while clone_alias in seen_clone_aliases:
|
||||
counter += 1
|
||||
clone_alias = f"{clone_alias_base}_{counter}"
|
||||
seen_clone_aliases.add(clone_alias)
|
||||
|
||||
clone = copy.deepcopy(base_mc)
|
||||
clone["alias"] = clone_alias
|
||||
params = clone.get("inference_parameters")
|
||||
if not isinstance(params, dict):
|
||||
params = {}
|
||||
clone["inference_parameters"] = params
|
||||
# data_designer's BaseInferenceParams is a pydantic model with
|
||||
# extra="forbid", so response_format cannot sit at the top level of
|
||||
# inference_parameters. It does expose an `extra_body: dict` pass-
|
||||
# through that the OpenAI client spreads into the request body at the
|
||||
# top level, which is where llama-server reads response_format from.
|
||||
# llama.cpp server shape (tools/server/README.md): the schema sits
|
||||
# directly under response_format, not nested in a json_schema object
|
||||
# the way OpenAI's Chat Completions API expects. llama-server converts
|
||||
# the schema to a GBNF grammar and applies it during sampling.
|
||||
extra_body = params.get("extra_body")
|
||||
if not isinstance(extra_body, dict):
|
||||
extra_body = {}
|
||||
extra_body["response_format"] = {
|
||||
"type": "json_schema",
|
||||
"schema": output_format,
|
||||
}
|
||||
params["extra_body"] = extra_body
|
||||
new_configs.append(clone)
|
||||
column["model_alias"] = clone_alias
|
||||
|
||||
if new_configs:
|
||||
model_configs.extend(new_configs)
|
||||
|
||||
|
||||
def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optional[int]:
|
||||
"""
|
||||
Mutate recipe dict in-place: for any provider with is_local=True,
|
||||
generate a JWT and fill in the endpoint pointing at this server.
|
||||
fill in the endpoint pointing at this server and inject a short-lived
|
||||
internal sk-unsloth-* API key for workflow auth.
|
||||
|
||||
Returns the row id of the minted internal key (so the caller can
|
||||
revoke it on job completion) or ``None`` when no local provider is
|
||||
actually reachable from an LLM column.
|
||||
"""
|
||||
providers = recipe.get("model_providers")
|
||||
if not providers:
|
||||
return
|
||||
return None
|
||||
|
||||
# Collect local providers and pop is_local from ALL dicts unconditionally.
|
||||
# Strict `is True` guard so malformed payloads (is_local: 1,
|
||||
|
|
@ -115,7 +213,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None:
|
|||
local_indices.append(i)
|
||||
|
||||
if not local_indices:
|
||||
return
|
||||
return None
|
||||
|
||||
endpoint = _resolve_local_v1_endpoint(request)
|
||||
|
||||
|
|
@ -138,6 +236,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None:
|
|||
}
|
||||
|
||||
token = ""
|
||||
internal_key_id: Optional[int] = None
|
||||
if local_names & referenced_providers:
|
||||
# Verify a model is loaded.
|
||||
# NOTE: This is a point-in-time check (TOCTOU). The model could be unloaded
|
||||
|
|
@ -158,18 +257,21 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None:
|
|||
"No model loaded in Chat. Load a model first, then run the recipe."
|
||||
)
|
||||
|
||||
from auth.authentication import (
|
||||
create_access_token,
|
||||
) # deferred: avoids circular import
|
||||
from auth import storage # deferred: avoids circular import
|
||||
|
||||
# Uses the "unsloth" admin subject. If the user changes their password,
|
||||
# the JWT secret rotates and this token becomes invalid mid-run.
|
||||
# Acceptable for v1 - recipes typically finish well within one session.
|
||||
token = create_access_token(
|
||||
subject = "unsloth",
|
||||
expires_delta = timedelta(hours = 24),
|
||||
desktop = _request_has_desktop_access_token(request),
|
||||
# Mint an internal sk-unsloth-* key scoped to this workflow run.
|
||||
# Uses the unified API-key issuance path (one mint/revoke/verify
|
||||
# surface instead of a second JWT code path). The key is marked
|
||||
# internal so it is hidden from the user's API-key list, and the
|
||||
# caller revokes it when the job terminates.
|
||||
expires_at = (datetime.now(timezone.utc) + timedelta(hours = 24)).isoformat()
|
||||
token, row = storage.create_api_key(
|
||||
username = "unsloth",
|
||||
name = "data-recipe workflow",
|
||||
expires_at = expires_at,
|
||||
internal = True,
|
||||
)
|
||||
internal_key_id = int(row["id"])
|
||||
|
||||
# Defensively strip any stale "external"-only fields the frontend may
|
||||
# have left on the dict (extra_headers/extra_body/api_key_env). The UI
|
||||
|
|
@ -196,6 +298,37 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None:
|
|||
continue
|
||||
if mc.get("provider") in local_names:
|
||||
mc["skip_health_check"] = True
|
||||
# Disable thinking for data-recipe inference on local providers.
|
||||
# Reasoning models emit a <think>...</think> preamble before the
|
||||
# answer, which roughly doubles generated token count per row and
|
||||
# pushes the visible answer past data_designer's json-fence
|
||||
# regex. Forward chat_template_kwargs={enable_thinking: False}
|
||||
# through the OpenAI SDK's extra_body passthrough so llama-server
|
||||
# renders the template without the reasoning preamble. Free-form
|
||||
# llm-text columns benefit from the latency cut, and structured
|
||||
# columns also stop leaking think tags into the grammar-
|
||||
# constrained JSON (llama-server's GBNF path still enforces the
|
||||
# schema either way).
|
||||
params = mc.get("inference_parameters")
|
||||
if not isinstance(params, dict):
|
||||
params = {}
|
||||
mc["inference_parameters"] = params
|
||||
extra_body = params.get("extra_body")
|
||||
if not isinstance(extra_body, dict):
|
||||
extra_body = {}
|
||||
tpl_kwargs = extra_body.get("chat_template_kwargs")
|
||||
if not isinstance(tpl_kwargs, dict):
|
||||
tpl_kwargs = {}
|
||||
tpl_kwargs.setdefault("enable_thinking", False)
|
||||
extra_body["chat_template_kwargs"] = tpl_kwargs
|
||||
params["extra_body"] = extra_body
|
||||
|
||||
# Forward each llm-structured column's output_format as an OpenAI
|
||||
# response_format so llama-server uses grammar-constrained sampling and
|
||||
# small GGUFs stop wasting the full max_tokens budget on broken JSON.
|
||||
_inject_local_structured_response_format(recipe, local_names)
|
||||
|
||||
return internal_key_id
|
||||
|
||||
|
||||
def _normalize_run_name(value: Any) -> str | None:
|
||||
|
|
@ -240,21 +373,49 @@ def create_job(payload: RecipePayload, request: Request):
|
|||
) from exc
|
||||
|
||||
try:
|
||||
_inject_local_providers(recipe, request)
|
||||
internal_api_key_id = _inject_local_providers(recipe, request)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code = 400, detail = str(exc)) from exc
|
||||
|
||||
mgr = get_job_manager()
|
||||
# Single try block covers get_job_manager() AND mgr.start() so a workflow
|
||||
# key minted above never outlives the request even when an unexpected
|
||||
# exception type (TypeError from a stale kwarg, OSError from a queue
|
||||
# write, etc.) bubbles up. Without the bare except, such exceptions let
|
||||
# the sk-unsloth-* key live until its 24h TTL.
|
||||
try:
|
||||
job_id = mgr.start(recipe = recipe, run = run)
|
||||
mgr = get_job_manager()
|
||||
job_id = mgr.start(
|
||||
recipe = recipe,
|
||||
run = run,
|
||||
internal_api_key_id = internal_api_key_id,
|
||||
)
|
||||
except RuntimeError as exc:
|
||||
if internal_api_key_id is not None:
|
||||
_revoke_internal_api_key_safe(internal_api_key_id)
|
||||
raise HTTPException(status_code = 409, detail = str(exc)) from exc
|
||||
except ValueError as exc:
|
||||
if internal_api_key_id is not None:
|
||||
_revoke_internal_api_key_safe(internal_api_key_id)
|
||||
raise HTTPException(status_code = 400, detail = str(exc)) from exc
|
||||
except Exception:
|
||||
if internal_api_key_id is not None:
|
||||
_revoke_internal_api_key_safe(internal_api_key_id)
|
||||
raise
|
||||
|
||||
return {"job_id": job_id}
|
||||
|
||||
|
||||
def _revoke_internal_api_key_safe(key_id: int) -> None:
|
||||
"""Best-effort revoke of a workflow-minted key; swallow any error so
|
||||
that revocation failures never mask the caller's own error path."""
|
||||
try:
|
||||
from auth import storage # deferred: avoids circular import
|
||||
|
||||
storage.revoke_internal_api_key(key_id)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@router.get("/jobs/{job_id}/status")
|
||||
def job_status(job_id: str):
|
||||
mgr = get_job_manager()
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from __future__ import annotations
|
|||
import base64
|
||||
import binascii
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from itertools import islice
|
||||
from pathlib import Path
|
||||
|
|
@ -627,3 +628,14 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons
|
|||
split = None,
|
||||
subset = None,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/seed/github/env-token")
|
||||
def get_github_env_token_status() -> dict:
|
||||
"""Report whether the server has a GH_TOKEN / GITHUB_TOKEN env var.
|
||||
|
||||
The value is never returned; the UI uses this to tell the user they
|
||||
can leave the token field blank.
|
||||
"""
|
||||
has_token = bool(os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN"))
|
||||
return {"has_token": has_token}
|
||||
|
|
|
|||
|
|
@ -14,10 +14,63 @@ from core.data_recipe.service import (
|
|||
create_data_designer,
|
||||
validate_recipe,
|
||||
)
|
||||
from loggers import get_logger
|
||||
from models.data_recipe import RecipePayload, ValidateError, ValidateResponse
|
||||
|
||||
logger = get_logger(__name__)
|
||||
router = APIRouter()
|
||||
|
||||
_GITHUB_VALIDATE_NOTE = "Recipe shape is valid. GitHub access and rate limits are checked when the run starts."
|
||||
_GITHUB_ITEM_TYPES = {"issues", "pulls", "commits"}
|
||||
|
||||
|
||||
def _github_seed_source(recipe: dict[str, Any]) -> dict[str, Any] | None:
|
||||
seed_config = recipe.get("seed_config")
|
||||
if not isinstance(seed_config, dict):
|
||||
return None
|
||||
source = seed_config.get("source")
|
||||
if not isinstance(source, dict) or source.get("seed_type") != "github_repo":
|
||||
return None
|
||||
return source
|
||||
|
||||
|
||||
def _validate_github_seed_static(source: dict[str, Any]) -> list[ValidateError]:
|
||||
errors: list[ValidateError] = []
|
||||
|
||||
repos = source.get("repos")
|
||||
if not isinstance(repos, list) or not repos:
|
||||
errors.append(ValidateError(message = "GitHub seed requires at least one repo."))
|
||||
else:
|
||||
for repo in repos:
|
||||
if not isinstance(repo, str) or not repo.strip() or "/" not in repo:
|
||||
errors.append(
|
||||
ValidateError(message = "GitHub repos must be owner/name strings.")
|
||||
)
|
||||
break
|
||||
|
||||
item_types = source.get("item_types")
|
||||
if not isinstance(item_types, list) or not item_types:
|
||||
errors.append(
|
||||
ValidateError(message = "GitHub seed requires at least one item type.")
|
||||
)
|
||||
else:
|
||||
invalid_items = [item for item in item_types if item not in _GITHUB_ITEM_TYPES]
|
||||
if invalid_items:
|
||||
errors.append(
|
||||
ValidateError(
|
||||
message = "GitHub item types must be issues, pulls, or commits."
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
limit = int(source.get("limit"))
|
||||
except (TypeError, ValueError):
|
||||
limit = 0
|
||||
if limit < 1 or limit > 5000:
|
||||
errors.append(ValidateError(message = "GitHub limit must be from 1 to 5000."))
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]:
|
||||
try:
|
||||
|
|
@ -93,6 +146,38 @@ def validate(payload: RecipePayload) -> ValidateResponse:
|
|||
|
||||
_patch_local_providers(recipe)
|
||||
|
||||
github_source = _github_seed_source(recipe)
|
||||
if github_source is not None:
|
||||
static_errors = _validate_github_seed_static(github_source)
|
||||
if static_errors:
|
||||
return ValidateResponse(valid = False, errors = static_errors)
|
||||
try:
|
||||
build_config_builder(recipe)
|
||||
except ModuleNotFoundError as exc:
|
||||
# data_designer is an optional runtime dep. Static validation
|
||||
# already passed; live access + full config validation are
|
||||
# deferred to run start (per _GITHUB_VALIDATE_NOTE), so a missing
|
||||
# optional import at validate time should not block the recipe.
|
||||
# Restrict the bypass to the data_designer module specifically so
|
||||
# other ImportErrors (e.g. broken internal imports or missing
|
||||
# transitive deps after a package upgrade) still surface as
|
||||
# validation failures instead of being silently swallowed.
|
||||
if not (exc.name or "").startswith("data_designer"):
|
||||
raise
|
||||
logger.debug(
|
||||
"data_designer not installed; deferring full config "
|
||||
"validation to run start",
|
||||
missing_module = exc.name,
|
||||
)
|
||||
except Exception as exc:
|
||||
detail = str(exc).strip() or "Validation failed."
|
||||
return ValidateResponse(
|
||||
valid = False,
|
||||
errors = [ValidateError(message = detail)],
|
||||
raw_detail = detail,
|
||||
)
|
||||
return ValidateResponse(valid = True, raw_detail = _GITHUB_VALIDATE_NOTE)
|
||||
|
||||
try:
|
||||
validate_recipe(recipe)
|
||||
except RuntimeError as exc:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -8,6 +8,7 @@ Model Management API routes
|
|||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
|
@ -623,7 +624,7 @@ def _scan_ollama_dir(
|
|||
gguf_link_path: Optional[str] = None
|
||||
quant = f"-{file_type}" if file_type else ""
|
||||
safe_name = repo_name.replace("/", "-")
|
||||
for layer in manifest.get("layers", []):
|
||||
for layer in manifest.get("layers") or []:
|
||||
media = layer.get("mediaType", "")
|
||||
digest = layer.get("digest", "")
|
||||
if not digest:
|
||||
|
|
@ -1684,6 +1685,338 @@ async def scan_loras(
|
|||
)
|
||||
|
||||
|
||||
def _is_path_under(path: Path, root: Path) -> bool:
|
||||
try:
|
||||
path.resolve().relative_to(root.resolve())
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _is_path_under_lexically(path: Path, root: Path) -> bool:
|
||||
"""Check containment without resolving the final path's symlink target."""
|
||||
try:
|
||||
absolute_path = Path(os.path.abspath(str(path)))
|
||||
absolute_root = Path(os.path.abspath(str(root)))
|
||||
absolute_path.relative_to(absolute_root)
|
||||
return True
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) -> bool:
|
||||
try:
|
||||
active = Path(active_model).expanduser().resolve()
|
||||
target = deleted_path.resolve()
|
||||
return active == target or (target.is_dir() and active.is_relative_to(target))
|
||||
except (OSError, RuntimeError, ValueError) as e:
|
||||
logger.debug(
|
||||
"Could not resolve loaded/deleted model paths; falling back to string comparison: %s",
|
||||
e,
|
||||
)
|
||||
active_lower = active_model.lower()
|
||||
target_lower = str(deleted_path).lower()
|
||||
return active_lower == target_lower or active_lower.startswith(
|
||||
f"{target_lower}{os.sep}"
|
||||
)
|
||||
|
||||
|
||||
def _loading_model_matches_deleted_path(
|
||||
loading_model: object,
|
||||
deleted_path: Path,
|
||||
) -> bool:
|
||||
if not loading_model:
|
||||
return False
|
||||
return _loaded_model_matches_deleted_path(str(loading_model), deleted_path)
|
||||
|
||||
|
||||
def _prune_empty_parents(start: Path, stop_at: Path) -> None:
|
||||
"""Remove empty ancestor directories of ``start`` up to (but not including) ``stop_at``.
|
||||
|
||||
Used after deleting a model checkpoint so the enclosing run directory does
|
||||
not linger as an empty entry in scan results.
|
||||
"""
|
||||
try:
|
||||
stop_resolved = stop_at.resolve()
|
||||
except OSError:
|
||||
return
|
||||
parent = start.parent
|
||||
while True:
|
||||
try:
|
||||
parent_resolved = parent.resolve()
|
||||
except OSError:
|
||||
return
|
||||
if parent_resolved == stop_resolved:
|
||||
return
|
||||
try:
|
||||
parent_resolved.relative_to(stop_resolved)
|
||||
except ValueError:
|
||||
return
|
||||
try:
|
||||
parent.rmdir()
|
||||
except OSError:
|
||||
return
|
||||
parent = parent.parent
|
||||
|
||||
|
||||
def _delete_gguf_variant_files(root: Path, variant: str) -> tuple[int, int]:
|
||||
deleted_count = 0
|
||||
deleted_bytes = 0
|
||||
for path in root.rglob("*"):
|
||||
if not path.is_file() or not _is_main_gguf_filename(path.name):
|
||||
continue
|
||||
if _extract_quant_label(path.name).lower() != variant.lower():
|
||||
continue
|
||||
try:
|
||||
deleted_bytes += path.stat().st_size
|
||||
except OSError:
|
||||
pass
|
||||
path.unlink()
|
||||
deleted_count += 1
|
||||
return deleted_count, deleted_bytes
|
||||
|
||||
|
||||
@router.delete("/delete-finetuned")
|
||||
async def delete_finetuned_model(
|
||||
model_path: str = Body(...),
|
||||
source: str = Body(...),
|
||||
export_type: Optional[str] = Body(None),
|
||||
gguf_variant: Optional[str] = Body(None),
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Delete a Studio-trained or exported model from disk.
|
||||
|
||||
Only paths under Studio's outputs/exports roots are accepted. Exported
|
||||
GGUF entries can delete one quantization variant at a time.
|
||||
"""
|
||||
if source not in {"training", "exported"}:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Only trained or exported Studio models can be deleted",
|
||||
)
|
||||
|
||||
if not model_path or not model_path.strip():
|
||||
raise HTTPException(status_code = 400, detail = "model_path is required")
|
||||
|
||||
if export_type == "gguf" and not gguf_variant:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "gguf_variant is required when export_type is 'gguf'",
|
||||
)
|
||||
|
||||
raw_path = Path(model_path).expanduser()
|
||||
if source == "training":
|
||||
target_path = raw_path
|
||||
allowed_root = outputs_root()
|
||||
else:
|
||||
allowed_root = exports_root()
|
||||
target_path = (
|
||||
raw_path.parent
|
||||
if export_type == "gguf" and raw_path.suffix.lower() == ".gguf"
|
||||
else raw_path
|
||||
)
|
||||
|
||||
allowed_root = allowed_root.resolve()
|
||||
delete_path = Path(os.path.abspath(str(target_path)))
|
||||
delete_path_is_symlink = delete_path.is_symlink()
|
||||
|
||||
if delete_path_is_symlink:
|
||||
if not _is_path_under_lexically(delete_path, allowed_root):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Model path is outside Studio storage",
|
||||
)
|
||||
if export_type == "gguf" and gguf_variant:
|
||||
target_path = delete_path.resolve()
|
||||
if not _is_path_under(target_path, allowed_root):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Model path is outside Studio storage",
|
||||
)
|
||||
else:
|
||||
target_path = delete_path
|
||||
else:
|
||||
target_path = target_path.resolve()
|
||||
|
||||
should_check_resolved_path = not delete_path_is_symlink or (
|
||||
export_type == "gguf" and gguf_variant
|
||||
)
|
||||
if should_check_resolved_path and not _is_path_under(target_path, allowed_root):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Model path is outside Studio storage",
|
||||
)
|
||||
if target_path == allowed_root:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Refusing to delete storage root",
|
||||
)
|
||||
if not target_path.exists() and not target_path.is_symlink():
|
||||
raise HTTPException(status_code = 404, detail = "Model not found on disk")
|
||||
|
||||
if source == "training":
|
||||
try:
|
||||
from core.training import get_training_backend
|
||||
|
||||
training_backend = get_training_backend()
|
||||
if training_backend.is_training_active():
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "Cannot delete trained models while training is running",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning("Could not check training status before delete: %s", e)
|
||||
raise HTTPException(
|
||||
status_code = 500,
|
||||
detail = "Could not verify training status before deleting",
|
||||
) from e
|
||||
|
||||
try:
|
||||
from routes.inference import get_llama_cpp_backend
|
||||
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
if (
|
||||
llama_backend.is_active
|
||||
and not llama_backend.is_loaded
|
||||
and llama_backend.model_identifier
|
||||
and _loaded_model_matches_deleted_path(
|
||||
llama_backend.model_identifier,
|
||||
target_path,
|
||||
)
|
||||
and (
|
||||
not gguf_variant
|
||||
or not llama_backend.hf_variant
|
||||
or llama_backend.hf_variant.lower() == gguf_variant.lower()
|
||||
)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "Cannot delete a model while it is loading",
|
||||
)
|
||||
if (
|
||||
llama_backend.is_loaded
|
||||
and llama_backend.model_identifier
|
||||
and _loaded_model_matches_deleted_path(
|
||||
llama_backend.model_identifier,
|
||||
target_path,
|
||||
)
|
||||
and (
|
||||
not gguf_variant
|
||||
or not llama_backend.hf_variant
|
||||
or llama_backend.hf_variant.lower() == gguf_variant.lower()
|
||||
)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Unload the model before deleting",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning("Could not check llama.cpp loaded model before delete: %s", e)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = "Could not verify model load status before deleting",
|
||||
) from e
|
||||
|
||||
try:
|
||||
inference_backend = get_inference_backend()
|
||||
loading_models = getattr(inference_backend, "loading_models", set())
|
||||
if any(
|
||||
_loading_model_matches_deleted_path(loading_model, target_path)
|
||||
for loading_model in loading_models
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "Cannot delete a model while it is loading",
|
||||
)
|
||||
if inference_backend.active_model_name:
|
||||
if _loaded_model_matches_deleted_path(
|
||||
inference_backend.active_model_name,
|
||||
target_path,
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Unload the model before deleting",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Could not check inference backend loaded model before delete: %s", e
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = "Could not verify model load status before deleting",
|
||||
) from e
|
||||
|
||||
try:
|
||||
if export_type == "gguf" and gguf_variant:
|
||||
if not target_path.is_dir():
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "GGUF variant deletion requires an export directory",
|
||||
)
|
||||
deleted_count, deleted_bytes = _delete_gguf_variant_files(
|
||||
target_path,
|
||||
gguf_variant,
|
||||
)
|
||||
if deleted_count == 0:
|
||||
raise HTTPException(
|
||||
status_code = 404,
|
||||
detail = f"Variant {gguf_variant} not found on disk",
|
||||
)
|
||||
try:
|
||||
if not any(target_path.iterdir()):
|
||||
target_path.rmdir()
|
||||
_prune_empty_parents(target_path, allowed_root)
|
||||
except OSError:
|
||||
pass
|
||||
logger.info(
|
||||
"Deleted %s GGUF file(s) for exported model at %s variant %s (%0.1f MB freed)",
|
||||
deleted_count,
|
||||
target_path,
|
||||
gguf_variant,
|
||||
deleted_bytes / (1024 * 1024),
|
||||
)
|
||||
return {
|
||||
"status": "deleted",
|
||||
"path": str(target_path),
|
||||
"gguf_variant": gguf_variant,
|
||||
}
|
||||
|
||||
if target_path.is_symlink() or target_path.is_file():
|
||||
target_path.unlink()
|
||||
else:
|
||||
shutil.rmtree(target_path)
|
||||
|
||||
if target_path.exists() or target_path.is_symlink():
|
||||
raise HTTPException(
|
||||
status_code = 500,
|
||||
detail = "Deletion incomplete; some files could not be removed",
|
||||
)
|
||||
|
||||
_prune_empty_parents(target_path, allowed_root)
|
||||
|
||||
logger.info("Deleted fine-tuned model at %s", target_path)
|
||||
return {"status": "deleted", "path": str(target_path)}
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error deleting fine-tuned model %s: %s",
|
||||
target_path,
|
||||
e,
|
||||
exc_info = True,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 500,
|
||||
detail = f"Failed to delete fine-tuned model: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@router.get("/loras/{lora_path:path}/base-model", response_model = LoRABaseModelResponse)
|
||||
async def get_lora_base_model(
|
||||
lora_path: str,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,12 @@ if str(backend_path) not in sys.path:
|
|||
# Import backend functions
|
||||
try:
|
||||
from core.training import get_training_backend
|
||||
from core.training.resume import (
|
||||
can_resume_run,
|
||||
get_resume_checkpoint_path,
|
||||
normalize_resume_output_dir,
|
||||
)
|
||||
from storage.studio_db import get_resumable_run_by_output_dir
|
||||
from utils.models.model_config import load_model_defaults
|
||||
from utils.paths import resolve_dataset_path
|
||||
except ImportError:
|
||||
|
|
@ -33,6 +39,12 @@ except ImportError:
|
|||
if str(parent_backend) not in sys.path:
|
||||
sys.path.insert(0, str(parent_backend))
|
||||
from core.training import get_training_backend
|
||||
from core.training.resume import (
|
||||
can_resume_run,
|
||||
get_resume_checkpoint_path,
|
||||
normalize_resume_output_dir,
|
||||
)
|
||||
from storage.studio_db import get_resumable_run_by_output_dir
|
||||
from utils.models.model_config import load_model_defaults
|
||||
from utils.paths import resolve_dataset_path
|
||||
|
||||
|
|
@ -152,6 +164,28 @@ async def start_training(
|
|||
request.local_eval_datasets = _validate_local_dataset_paths(
|
||||
request.local_eval_datasets, "Local eval dataset"
|
||||
)
|
||||
resume_output_dir: Optional[str] = None
|
||||
if request.resume_from_checkpoint:
|
||||
try:
|
||||
resume_output_dir = normalize_resume_output_dir(
|
||||
request.resume_from_checkpoint
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code = 400, detail = str(e))
|
||||
|
||||
resume_run = get_resumable_run_by_output_dir(resume_output_dir)
|
||||
if not resume_run or not can_resume_run(resume_run):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Resume checkpoint must belong to a stopped run with saved trainer state.",
|
||||
)
|
||||
resume_checkpoint = get_resume_checkpoint_path(resume_output_dir)
|
||||
if not resume_checkpoint:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Resume checkpoint must include saved trainer state.",
|
||||
)
|
||||
request.resume_from_checkpoint = resume_checkpoint
|
||||
|
||||
# Convert request to kwargs for backend
|
||||
training_kwargs = {
|
||||
|
|
@ -209,6 +243,8 @@ async def start_training(
|
|||
"wandb_project": request.wandb_project or "",
|
||||
"enable_tensorboard": request.enable_tensorboard,
|
||||
"tensorboard_dir": request.tensorboard_dir or "",
|
||||
"output_dir": resume_output_dir,
|
||||
"resume_from_checkpoint": request.resume_from_checkpoint,
|
||||
"trust_remote_code": request.trust_remote_code,
|
||||
"gpu_ids": request.gpu_ids,
|
||||
}
|
||||
|
|
@ -437,6 +473,9 @@ async def get_training_status(
|
|||
"loss": getattr(progress, "loss", None),
|
||||
"learning_rate": getattr(progress, "learning_rate", None),
|
||||
}
|
||||
output_dir = getattr(backend, "_output_dir", None)
|
||||
if output_dir:
|
||||
details["output_dir"] = output_dir
|
||||
|
||||
# Build metric history for chart recovery after SSE reconnection
|
||||
metric_history = None
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from fastapi import APIRouter, Depends, HTTPException, Query
|
|||
from loggers import get_logger
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from core.training.resume import can_resume_run
|
||||
from models import (
|
||||
TrainingRunDeleteResponse,
|
||||
TrainingRunDetailResponse,
|
||||
|
|
@ -34,7 +35,10 @@ async def list_training_runs(
|
|||
"""List training runs, newest first."""
|
||||
result = list_runs(limit = limit, offset = offset)
|
||||
return TrainingRunListResponse(
|
||||
runs = [TrainingRunSummary(**r) for r in result["runs"]],
|
||||
runs = [
|
||||
TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)})
|
||||
for r in result["runs"]
|
||||
],
|
||||
total = result["total"],
|
||||
)
|
||||
|
||||
|
|
@ -58,7 +62,12 @@ async def get_training_run_detail(
|
|||
metrics_data = get_run_metrics(run_id)
|
||||
|
||||
return TrainingRunDetailResponse(
|
||||
run = TrainingRunSummary(**{k: v for k, v in run.items() if k != "config_json"}),
|
||||
run = TrainingRunSummary(
|
||||
**{
|
||||
**{k: v for k, v in run.items() if k != "config_json"},
|
||||
"can_resume": can_resume_run(run),
|
||||
}
|
||||
),
|
||||
config = config,
|
||||
metrics = TrainingRunMetrics(**metrics_data),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -244,7 +244,7 @@ _shutdown_event = None
|
|||
|
||||
|
||||
def run_server(
|
||||
host: str = "0.0.0.0",
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 8888,
|
||||
frontend_path: Path = Path(__file__).resolve().parent.parent / "frontend" / "dist",
|
||||
silent: bool = False,
|
||||
|
|
@ -392,7 +392,11 @@ if __name__ == "__main__":
|
|||
pass
|
||||
|
||||
parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server")
|
||||
parser.add_argument("--host", default = "0.0.0.0", help = "Host to bind to")
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
default = "127.0.0.1",
|
||||
help = "Host to bind to (default: 127.0.0.1; use 0.0.0.0 for network/cloud access)",
|
||||
)
|
||||
parser.add_argument("--port", type = int, default = 8888, help = "Port to bind to")
|
||||
parser.add_argument(
|
||||
"--frontend",
|
||||
|
|
|
|||
33
studio/backend/state/tool_policy.py
Normal file
33
studio/backend/state/tool_policy.py
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""Process-level server-side tool policy.
|
||||
|
||||
Set by `unsloth run` at startup; consulted by the inference route gates.
|
||||
|
||||
None -> no CLI override (default). Per-request `enable_tools` is honored.
|
||||
True -> CLI forced tools on for every request.
|
||||
False -> CLI forced tools off for every request.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
_tool_policy: Optional[bool] = None
|
||||
|
||||
|
||||
def get_tool_policy() -> Optional[bool]:
|
||||
return _tool_policy
|
||||
|
||||
|
||||
def set_tool_policy(value: Optional[bool]) -> None:
|
||||
if value is not None and not isinstance(value, bool):
|
||||
raise TypeError(
|
||||
f"tool_policy must be Optional[bool], got {type(value).__name__}"
|
||||
)
|
||||
global _tool_policy
|
||||
_tool_policy = value
|
||||
|
||||
|
||||
def reset_tool_policy() -> None:
|
||||
global _tool_policy
|
||||
_tool_policy = None
|
||||
|
|
@ -267,10 +267,23 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
|
|||
total = conn.execute("SELECT COUNT(*) FROM training_runs").fetchone()[0]
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT id, status, model_name, dataset_name, started_at, ended_at,
|
||||
total_steps, final_step, final_loss, output_dir,
|
||||
duration_seconds, error_message, loss_sparkline
|
||||
FROM training_runs
|
||||
SELECT r.id, r.status, r.model_name, r.dataset_name, r.started_at,
|
||||
r.ended_at, r.total_steps, r.final_step, r.final_loss,
|
||||
r.output_dir, r.duration_seconds, r.error_message,
|
||||
r.loss_sparkline,
|
||||
CASE
|
||||
WHEN r.status = 'stopped'
|
||||
AND r.output_dir IS NOT NULL
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM training_runs newer
|
||||
WHERE newer.output_dir = r.output_dir
|
||||
AND newer.status IN ('stopped', 'completed')
|
||||
AND newer.started_at > r.started_at
|
||||
)
|
||||
THEN 1 ELSE 0
|
||||
END AS resumed_later
|
||||
FROM training_runs r
|
||||
ORDER BY started_at DESC
|
||||
LIMIT ? OFFSET ?
|
||||
""",
|
||||
|
|
@ -297,7 +310,26 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
|
|||
def get_run(id: str) -> Optional[dict]:
|
||||
conn = get_connection()
|
||||
try:
|
||||
row = conn.execute("SELECT * FROM training_runs WHERE id = ?", (id,)).fetchone()
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT r.*,
|
||||
CASE
|
||||
WHEN r.status = 'stopped'
|
||||
AND r.output_dir IS NOT NULL
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM training_runs newer
|
||||
WHERE newer.output_dir = r.output_dir
|
||||
AND newer.status IN ('stopped', 'completed')
|
||||
AND newer.started_at > r.started_at
|
||||
)
|
||||
THEN 1 ELSE 0
|
||||
END AS resumed_later
|
||||
FROM training_runs r
|
||||
WHERE r.id = ?
|
||||
""",
|
||||
(id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
run = dict(row)
|
||||
|
|
@ -313,6 +345,45 @@ def get_run(id: str) -> Optional[dict]:
|
|||
conn.close()
|
||||
|
||||
|
||||
def get_resumable_run_by_output_dir(output_dir: str) -> Optional[dict]:
|
||||
conn = get_connection()
|
||||
try:
|
||||
row = conn.execute(
|
||||
"""
|
||||
SELECT r.*,
|
||||
0 AS resumed_later
|
||||
FROM training_runs r
|
||||
WHERE r.output_dir = ?
|
||||
AND r.status = 'stopped'
|
||||
AND NOT EXISTS (
|
||||
SELECT 1
|
||||
FROM training_runs newer
|
||||
WHERE newer.output_dir = r.output_dir
|
||||
AND newer.status IN ('stopped', 'completed')
|
||||
AND newer.started_at > r.started_at
|
||||
)
|
||||
ORDER BY r.started_at DESC
|
||||
LIMIT 1
|
||||
""",
|
||||
(output_dir,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
run = dict(row)
|
||||
sparkline = run.get("loss_sparkline")
|
||||
if sparkline:
|
||||
try:
|
||||
run["loss_sparkline"] = json.loads(sparkline)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
logger.debug(
|
||||
"Failed to parse loss_sparkline for output_dir %s", output_dir
|
||||
)
|
||||
run["loss_sparkline"] = None
|
||||
return run
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
def get_run_metrics(id: str) -> dict:
|
||||
"""Return metric arrays for a run, using paired step arrays per metric."""
|
||||
conn = get_connection()
|
||||
|
|
|
|||
91
studio/backend/tests/test_data_recipe_github_progress.py
Normal file
91
studio/backend/tests/test_data_recipe_github_progress.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from core.data_recipe.jobs.parse import apply_update, parse_log_message
|
||||
from core.data_recipe.jobs.types import Job
|
||||
from routes.data_recipe.validate import _GITHUB_VALIDATE_NOTE, validate
|
||||
from models.data_recipe import RecipePayload
|
||||
|
||||
|
||||
def test_github_page_log_updates_source_progress_without_cursor():
|
||||
job = Job(job_id = "job-1")
|
||||
job.source_progress_estimated_total = 200
|
||||
|
||||
update = parse_log_message(
|
||||
"[unslothai/unsloth] issues page 2 (+15) cursor=abc123 remaining=2960"
|
||||
)
|
||||
|
||||
assert update is not None
|
||||
apply_update(job, update)
|
||||
|
||||
progress = job.source_progress
|
||||
assert progress is not None
|
||||
assert progress.source == "github"
|
||||
assert progress.status == "fetching"
|
||||
assert progress.repo == "unslothai/unsloth"
|
||||
assert progress.resource == "issues"
|
||||
assert progress.page == 2
|
||||
assert progress.page_items == 15
|
||||
assert progress.fetched_items == 15
|
||||
assert progress.estimated_total == 200
|
||||
assert progress.rate_remaining == 2960
|
||||
assert progress.message is not None
|
||||
assert "cursor" not in progress.message
|
||||
assert "abc123" not in progress.message
|
||||
|
||||
|
||||
def test_github_rate_limit_log_updates_source_progress():
|
||||
job = Job(job_id = "job-1")
|
||||
|
||||
update = parse_log_message("Rate limit hit. Sleeping 123s until reset.")
|
||||
|
||||
assert update is not None
|
||||
apply_update(job, update)
|
||||
|
||||
progress = job.source_progress
|
||||
assert progress is not None
|
||||
assert progress.status == "rate_limited"
|
||||
assert progress.retry_after_sec == 123
|
||||
assert "resume automatically" in (progress.message or "")
|
||||
|
||||
|
||||
def test_github_real_sample_prs_and_trial_limit_are_parsed():
|
||||
job = Job(job_id = "job-1")
|
||||
|
||||
for message in (
|
||||
"[unslothai/unsloth] PRs page 4 (+25) cursor=abc123 remaining=4983",
|
||||
"Trial limit reached for PRs (100)",
|
||||
):
|
||||
update = parse_log_message(message)
|
||||
assert update is not None
|
||||
apply_update(job, update)
|
||||
|
||||
progress = job.source_progress
|
||||
assert progress is not None
|
||||
assert progress.repo == "unslothai/unsloth"
|
||||
assert progress.resource == "pulls"
|
||||
assert progress.page == 4
|
||||
assert progress.fetched_items == 25
|
||||
assert progress.rate_remaining == 4983
|
||||
assert progress.message == "GitHub pulls trial limit reached (100)."
|
||||
|
||||
|
||||
def test_github_validate_skips_live_access_with_honest_note():
|
||||
response = validate(
|
||||
RecipePayload(
|
||||
recipe = {
|
||||
"seed_config": {
|
||||
"source": {
|
||||
"seed_type": "github_repo",
|
||||
"repos": ["unslothai/unsloth"],
|
||||
"item_types": ["issues"],
|
||||
"limit": 1,
|
||||
}
|
||||
},
|
||||
"columns": [{"column_type": "expression", "name": "x", "expr": "1"}],
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert response.valid is True
|
||||
assert response.raw_detail == _GITHUB_VALIDATE_NOTE
|
||||
|
|
@ -246,7 +246,10 @@ def test_desktop_session_uses_real_admin_identity_for_api_keys():
|
|||
assert [row["name"] for row in rows] == ["desktop"]
|
||||
|
||||
|
||||
def test_local_recipe_token_preserves_desktop_marker(loaded_local_model):
|
||||
def test_local_recipe_token_authenticates_as_admin_for_desktop_user(loaded_local_model):
|
||||
# _inject_local_providers mints an internal sk-unsloth-* API key (not a
|
||||
# forwarded JWT). The unified API-key path validates as the real admin
|
||||
# user regardless of whether the incoming session was desktop or web.
|
||||
from auth.authentication import create_access_token, get_current_subject
|
||||
|
||||
seed_user(must_change_password = True)
|
||||
|
|
@ -260,13 +263,7 @@ def test_local_recipe_token_preserves_desktop_marker(loaded_local_model):
|
|||
jobs_route._inject_local_providers(recipe, local_recipe_request(incoming_token))
|
||||
|
||||
local_token = recipe["model_providers"][0]["api_key"]
|
||||
payload = jwt.decode(
|
||||
local_token,
|
||||
storage.get_jwt_secret(storage.DEFAULT_ADMIN_USERNAME),
|
||||
algorithms = ["HS256"],
|
||||
)
|
||||
assert payload["sub"] == storage.DEFAULT_ADMIN_USERNAME
|
||||
assert payload["desktop"] is True
|
||||
assert local_token.startswith(storage.API_KEY_PREFIX)
|
||||
credentials = HTTPAuthorizationCredentials(
|
||||
scheme = "Bearer",
|
||||
credentials = local_token,
|
||||
|
|
@ -276,8 +273,10 @@ def test_local_recipe_token_preserves_desktop_marker(loaded_local_model):
|
|||
)
|
||||
|
||||
|
||||
def test_local_recipe_token_keeps_web_marker_absent(loaded_local_model):
|
||||
from auth.authentication import create_access_token
|
||||
def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_model):
|
||||
# Mirror of the desktop variant: API-key issuance is identical for web
|
||||
# and desktop incoming tokens; auth via get_current_subject works the same.
|
||||
from auth.authentication import create_access_token, get_current_subject
|
||||
|
||||
seed_user(must_change_password = False)
|
||||
jobs_route = data_recipe_jobs_module()
|
||||
|
|
@ -287,13 +286,14 @@ def test_local_recipe_token_keeps_web_marker_absent(loaded_local_model):
|
|||
jobs_route._inject_local_providers(recipe, local_recipe_request(incoming_token))
|
||||
|
||||
local_token = recipe["model_providers"][0]["api_key"]
|
||||
payload = jwt.decode(
|
||||
local_token,
|
||||
storage.get_jwt_secret(storage.DEFAULT_ADMIN_USERNAME),
|
||||
algorithms = ["HS256"],
|
||||
assert local_token.startswith(storage.API_KEY_PREFIX)
|
||||
credentials = HTTPAuthorizationCredentials(
|
||||
scheme = "Bearer",
|
||||
credentials = local_token,
|
||||
)
|
||||
assert (
|
||||
asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
|
||||
)
|
||||
assert payload["sub"] == storage.DEFAULT_ADMIN_USERNAME
|
||||
assert "desktop" not in payload
|
||||
|
||||
|
||||
def test_desktop_login_rejects_invalid_secret():
|
||||
|
|
@ -381,6 +381,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
|
|||
datasets_router = APIRouter(),
|
||||
export_router = APIRouter(),
|
||||
inference_router = APIRouter(),
|
||||
inference_studio_router = APIRouter(),
|
||||
models_router = APIRouter(),
|
||||
training_history_router = APIRouter(),
|
||||
training_router = APIRouter(),
|
||||
|
|
|
|||
|
|
@ -746,7 +746,15 @@ class TestRouteErrors(unittest.TestCase):
|
|||
):
|
||||
with self.assertRaises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
inference_route.load_model(request, current_subject = "test-user")
|
||||
inference_route.load_model(
|
||||
request,
|
||||
SimpleNamespace(
|
||||
app = SimpleNamespace(
|
||||
state = SimpleNamespace(llama_parallel_slots = 1),
|
||||
),
|
||||
),
|
||||
current_subject = "test-user",
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(exc_info.exception.status_code, 400)
|
||||
|
|
@ -886,7 +894,15 @@ class TestRouteErrors(unittest.TestCase):
|
|||
):
|
||||
with self.assertRaises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
inference_route.load_model(request, current_subject = "test-user")
|
||||
inference_route.load_model(
|
||||
request,
|
||||
SimpleNamespace(
|
||||
app = SimpleNamespace(
|
||||
state = SimpleNamespace(llama_parallel_slots = 1),
|
||||
),
|
||||
),
|
||||
current_subject = "test-user",
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(exc_info.exception.status_code, 400)
|
||||
|
|
@ -942,7 +958,15 @@ class TestRouteErrors(unittest.TestCase):
|
|||
):
|
||||
with self.assertRaises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
inference_route.load_model(request, current_subject = "test-user")
|
||||
inference_route.load_model(
|
||||
request,
|
||||
SimpleNamespace(
|
||||
app = SimpleNamespace(
|
||||
state = SimpleNamespace(llama_parallel_slots = 1),
|
||||
),
|
||||
),
|
||||
current_subject = "test-user",
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(exc_info.exception.status_code, 400)
|
||||
|
|
@ -1025,6 +1049,182 @@ class TestMinGpuVram(unittest.TestCase):
|
|||
|
||||
|
||||
class TestPerGpuFitGuardAllCounts(unittest.TestCase):
|
||||
def test_training_estimate_resolves_attention_without_raising(self):
|
||||
with (
|
||||
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
||||
patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (8 * (1024**3), "config"),
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
||||
return_value = "unsloth/test",
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
||||
return_value = SimpleNamespace(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 128256,
|
||||
tie_word_embeddings = False,
|
||||
),
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
|
||||
return_value = "eager",
|
||||
),
|
||||
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
|
||||
):
|
||||
_, metadata = estimate_required_model_memory_gb(
|
||||
"unsloth/test",
|
||||
training_type = "LoRA/QLoRA",
|
||||
load_in_4bit = True,
|
||||
)
|
||||
|
||||
self.assertEqual(metadata.get("estimation_mode"), "detailed")
|
||||
self.assertEqual(metadata.get("attention_implementation"), "eager")
|
||||
|
||||
def test_training_estimate_falls_back_when_attention_resolution_fails(self):
|
||||
with (
|
||||
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
||||
patch(
|
||||
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
||||
return_value = (8 * (1024**3), "config"),
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
||||
return_value = "unsloth/test",
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
||||
return_value = SimpleNamespace(
|
||||
hidden_size = 4096,
|
||||
num_hidden_layers = 32,
|
||||
num_attention_heads = 32,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 14336,
|
||||
vocab_size = 128256,
|
||||
tie_word_embeddings = False,
|
||||
),
|
||||
),
|
||||
patch(
|
||||
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
|
||||
side_effect = RuntimeError("attention unavailable"),
|
||||
),
|
||||
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
|
||||
):
|
||||
_, metadata = estimate_required_model_memory_gb(
|
||||
"unsloth/test",
|
||||
training_type = "LoRA/QLoRA",
|
||||
load_in_4bit = True,
|
||||
)
|
||||
|
||||
self.assertEqual(metadata.get("estimation_mode"), "detailed")
|
||||
self.assertEqual(
|
||||
metadata.get("attention_implementation"),
|
||||
"eager",
|
||||
)
|
||||
|
||||
def test_attention_resolver_does_not_mutate_loaded_config(self):
|
||||
from utils.hardware import hardware as hardware_module
|
||||
|
||||
config = SimpleNamespace(
|
||||
hidden_size = 1024,
|
||||
num_hidden_layers = 2,
|
||||
num_attention_heads = 8,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 2048,
|
||||
vocab_size = 1024,
|
||||
tie_word_embeddings = True,
|
||||
)
|
||||
|
||||
def _stub_resolver(model_class, cfg):
|
||||
cfg._attn_implementation = "eager"
|
||||
return "eager"
|
||||
|
||||
with patch(
|
||||
"unsloth.models._utils.resolve_attention_implementation",
|
||||
side_effect = _stub_resolver,
|
||||
):
|
||||
hardware_module._determine_attention_impl_for_gpu_estimate(config)
|
||||
|
||||
self.assertFalse(hasattr(config, "_attn_implementation"))
|
||||
|
||||
def test_attention_resolver_handles_missing_model_mapping(self):
|
||||
from utils.hardware import hardware as hardware_module
|
||||
|
||||
config = SimpleNamespace(
|
||||
hidden_size = 1024,
|
||||
num_hidden_layers = 2,
|
||||
num_attention_heads = 8,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 2048,
|
||||
vocab_size = 1024,
|
||||
tie_word_embeddings = True,
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def _stub_resolver(model_class, cfg):
|
||||
captured["model_class"] = model_class
|
||||
return "eager"
|
||||
|
||||
from transformers import AutoModel, AutoModelForCausalLM
|
||||
|
||||
with (
|
||||
patch.object(AutoModelForCausalLM, "_model_mapping", new = None),
|
||||
patch.object(AutoModel, "_model_mapping", new = None),
|
||||
patch(
|
||||
"unsloth.models._utils.resolve_attention_implementation",
|
||||
side_effect = _stub_resolver,
|
||||
),
|
||||
):
|
||||
result = hardware_module._determine_attention_impl_for_gpu_estimate(config)
|
||||
|
||||
self.assertEqual(result, "eager")
|
||||
self.assertIsNone(captured["model_class"])
|
||||
|
||||
def test_attention_resolver_does_not_mutate_nested_text_config(self):
|
||||
from utils.hardware import hardware as hardware_module
|
||||
|
||||
text_config = SimpleNamespace(
|
||||
hidden_size = 1024,
|
||||
num_hidden_layers = 2,
|
||||
num_attention_heads = 8,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 2048,
|
||||
vocab_size = 1024,
|
||||
tie_word_embeddings = True,
|
||||
)
|
||||
config = SimpleNamespace(
|
||||
hidden_size = 1024,
|
||||
num_hidden_layers = 2,
|
||||
num_attention_heads = 8,
|
||||
num_key_value_heads = 8,
|
||||
intermediate_size = 2048,
|
||||
vocab_size = 1024,
|
||||
tie_word_embeddings = True,
|
||||
text_config = text_config,
|
||||
)
|
||||
|
||||
def _stub_resolver(model_class, cfg):
|
||||
cfg._attn_implementation = "eager"
|
||||
inner = getattr(cfg, "text_config", None)
|
||||
if inner is not None:
|
||||
inner._attn_implementation = "eager"
|
||||
return "eager"
|
||||
|
||||
with patch(
|
||||
"unsloth.models._utils.resolve_attention_implementation",
|
||||
side_effect = _stub_resolver,
|
||||
):
|
||||
hardware_module._determine_attention_impl_for_gpu_estimate(config)
|
||||
|
||||
self.assertFalse(hasattr(config, "_attn_implementation"))
|
||||
self.assertFalse(hasattr(text_config, "_attn_implementation"))
|
||||
|
||||
def test_min_per_gpu_generated_for_all_visible_counts(self):
|
||||
with (
|
||||
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
||||
|
|
@ -1101,3 +1301,123 @@ class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase):
|
|||
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU):
|
||||
with self.assertRaisesRegex(ValueError, "only supported on CUDA"):
|
||||
prepare_gpu_selection([0], model_name = "unsloth/test")
|
||||
|
||||
|
||||
class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase):
|
||||
def _run(
|
||||
self,
|
||||
model_path,
|
||||
*,
|
||||
config_bytes,
|
||||
local_bytes,
|
||||
safetensors_params = None,
|
||||
config = object(),
|
||||
):
|
||||
from utils.hardware import hardware as hardware_module
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_resolve_model_identifier_for_gpu_estimate",
|
||||
return_value = model_path,
|
||||
),
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_get_hf_safetensors_total_params",
|
||||
return_value = safetensors_params,
|
||||
),
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_load_config_for_gpu_estimate",
|
||||
return_value = config,
|
||||
),
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_estimate_fp16_model_size_bytes_from_config",
|
||||
return_value = config_bytes,
|
||||
),
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_get_local_weight_size_bytes",
|
||||
return_value = local_bytes,
|
||||
),
|
||||
):
|
||||
return hardware_module.estimate_fp16_model_size_bytes(model_path)
|
||||
|
||||
def test_local_weight_bytes_preferred_when_larger_than_config(self):
|
||||
bytes_, src = self._run(
|
||||
"/local/vlm",
|
||||
config_bytes = 2 * (1 << 30),
|
||||
local_bytes = 20 * (1 << 30),
|
||||
)
|
||||
self.assertEqual(bytes_, 20 * (1 << 30))
|
||||
self.assertEqual(src, "weight_bytes")
|
||||
|
||||
def test_config_bytes_preferred_when_larger_than_local(self):
|
||||
bytes_, src = self._run(
|
||||
"/local/text-only",
|
||||
config_bytes = 20 * (1 << 30),
|
||||
local_bytes = 2 * (1 << 30),
|
||||
)
|
||||
self.assertEqual(bytes_, 20 * (1 << 30))
|
||||
self.assertEqual(src, "config")
|
||||
|
||||
def test_config_bytes_returned_when_no_local_weights(self):
|
||||
bytes_, src = self._run(
|
||||
"/local/no-weights",
|
||||
config_bytes = 5 * (1 << 30),
|
||||
local_bytes = None,
|
||||
)
|
||||
self.assertEqual(bytes_, 5 * (1 << 30))
|
||||
self.assertEqual(src, "config")
|
||||
|
||||
def test_local_bytes_returned_when_config_resolution_fails(self):
|
||||
bytes_, src = self._run(
|
||||
"/local/no-config",
|
||||
config_bytes = None,
|
||||
local_bytes = 7 * (1 << 30),
|
||||
config = None,
|
||||
)
|
||||
self.assertEqual(bytes_, 7 * (1 << 30))
|
||||
self.assertEqual(src, "weight_bytes")
|
||||
|
||||
def test_equal_local_and_config_keeps_config_label(self):
|
||||
# why: tie-breaker is "local must be strictly larger" so an exact
|
||||
# match keeps the config-derived path.
|
||||
same = 8 * (1 << 30)
|
||||
bytes_, src = self._run(
|
||||
"/local/equal",
|
||||
config_bytes = same,
|
||||
local_bytes = same,
|
||||
)
|
||||
self.assertEqual(bytes_, same)
|
||||
self.assertEqual(src, "config")
|
||||
|
||||
def test_remote_safetensors_path_unaffected_by_local_weights(self):
|
||||
from utils.hardware import hardware as hardware_module
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_resolve_model_identifier_for_gpu_estimate",
|
||||
return_value = "owner/repo",
|
||||
),
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_get_hf_safetensors_total_params",
|
||||
return_value = 1_000_000_000,
|
||||
),
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_load_config_for_gpu_estimate",
|
||||
) as mock_load,
|
||||
patch.object(
|
||||
hardware_module,
|
||||
"_get_local_weight_size_bytes",
|
||||
) as mock_local,
|
||||
):
|
||||
bytes_, src = hardware_module.estimate_fp16_model_size_bytes("owner/repo")
|
||||
self.assertEqual(bytes_, 2 * 1_000_000_000)
|
||||
self.assertEqual(src, "safetensors")
|
||||
mock_load.assert_not_called()
|
||||
mock_local.assert_not_called()
|
||||
|
|
|
|||
98
studio/backend/tests/test_host_defaults.py
Normal file
98
studio/backend/tests/test_host_defaults.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests that Unsloth Studio defaults to 127.0.0.1 (loopback) not 0.0.0.0.
|
||||
|
||||
Uses AST parsing to inspect source-level defaults without requiring the
|
||||
full studio venv (run.py has heavy dependencies like structlog/uvicorn).
|
||||
"""
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
_RUN_PY = Path(__file__).resolve().parent.parent / "run.py"
|
||||
|
||||
|
||||
def _parse_function_param_defaults(source: str, func_name: str) -> dict:
|
||||
"""Return {param_name: default_value} for a named function in *source*.
|
||||
|
||||
Only handles ast.Constant defaults (strings, ints, bools).
|
||||
"""
|
||||
tree = ast.parse(source)
|
||||
for node in ast.walk(tree):
|
||||
if (
|
||||
isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
||||
and node.name == func_name
|
||||
):
|
||||
result = {}
|
||||
all_args = node.args.args
|
||||
defaults = node.args.defaults
|
||||
# Defaults are right-aligned against the args list
|
||||
offset = len(all_args) - len(defaults)
|
||||
for i, default in enumerate(defaults):
|
||||
arg_name = all_args[offset + i].arg
|
||||
if isinstance(default, ast.Constant):
|
||||
result[arg_name] = default.value
|
||||
return result
|
||||
return {}
|
||||
|
||||
|
||||
def _parse_argparse_add_argument_default(source: str, option_name: str):
|
||||
"""Return the 'default' kwarg value for add_argument(option_name, ...) in *source*.
|
||||
|
||||
Walks the entire module so the call can live in __main__ or in a helper
|
||||
function — only handles ast.Constant defaults.
|
||||
"""
|
||||
tree = ast.parse(source)
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.Call):
|
||||
continue
|
||||
func = node.func
|
||||
if not (isinstance(func, ast.Attribute) and func.attr == "add_argument"):
|
||||
continue
|
||||
if not node.args:
|
||||
continue
|
||||
first_arg = node.args[0]
|
||||
if not (isinstance(first_arg, ast.Constant) and first_arg.value == option_name):
|
||||
continue
|
||||
for kw in node.keywords:
|
||||
if kw.arg == "default" and isinstance(kw.value, ast.Constant):
|
||||
return kw.value.value
|
||||
return None
|
||||
|
||||
|
||||
def test_run_server_default_host_is_loopback():
|
||||
"""run_server() parameter default for 'host' must be 127.0.0.1, not 0.0.0.0.
|
||||
|
||||
Binding to 0.0.0.0 by default exposes the service on all network
|
||||
interfaces, contradicting the documented "privacy first / 100% local"
|
||||
guarantee. Loopback (127.0.0.1) is the least-permissive default;
|
||||
users who need network access can pass -H 0.0.0.0 explicitly.
|
||||
"""
|
||||
source = _RUN_PY.read_text()
|
||||
defaults = _parse_function_param_defaults(source, "run_server")
|
||||
assert (
|
||||
"host" in defaults
|
||||
), "run_server() must have a 'host' parameter with a default"
|
||||
host_default = defaults["host"]
|
||||
assert host_default == "127.0.0.1", (
|
||||
f"run_server() host default must be '127.0.0.1' (loopback) "
|
||||
f"but got '{host_default}'. Binding to '{host_default}' by default "
|
||||
f"exposes the service beyond localhost."
|
||||
)
|
||||
|
||||
|
||||
def test_argparse_default_host_is_loopback():
|
||||
"""argparse --host add_argument default must be 127.0.0.1.
|
||||
|
||||
When run.py is invoked directly (python run.py), the argparse default
|
||||
should match the function default so direct execution is equally safe.
|
||||
"""
|
||||
source = _RUN_PY.read_text()
|
||||
host_default = _parse_argparse_add_argument_default(source, "--host")
|
||||
assert (
|
||||
host_default is not None
|
||||
), "Could not find add_argument('--host', ...) in run.py"
|
||||
assert (
|
||||
host_default == "127.0.0.1"
|
||||
), f"run.py argparse --host default must be '127.0.0.1', got '{host_default}'"
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -114,9 +114,13 @@ def _make_backend(
|
|||
inst._kv_value_length = kv_value_length
|
||||
inst._kv_lora_rank = None
|
||||
inst._sliding_window = None
|
||||
inst._sliding_window_pattern = None
|
||||
inst._ssm_inner_size = None
|
||||
inst._full_attention_interval = None
|
||||
inst._key_length_mla = None
|
||||
inst._n_kv_heads_by_layer = None
|
||||
inst._kv_key_length_swa = None
|
||||
inst._kv_value_length_swa = None
|
||||
return inst
|
||||
|
||||
|
||||
|
|
@ -137,7 +141,7 @@ def _drive(
|
|||
model_size = int(model_gib * GIB)
|
||||
cache_type_kv = None
|
||||
|
||||
def fake_estimate(n_ctx_, _type = None):
|
||||
def fake_estimate(n_ctx_, _type = None, **_kwargs):
|
||||
return 0 if n_ctx_ <= 0 else n_ctx_ * kv_per_token_bytes
|
||||
|
||||
inst._estimate_kv_cache_bytes = fake_estimate
|
||||
|
|
|
|||
|
|
@ -99,9 +99,13 @@ def _make_backend(native_ctx = 131072):
|
|||
inst._kv_value_length = 128
|
||||
inst._kv_lora_rank = None
|
||||
inst._sliding_window = None
|
||||
inst._sliding_window_pattern = None
|
||||
inst._ssm_inner_size = None
|
||||
inst._full_attention_interval = None
|
||||
inst._key_length_mla = None
|
||||
inst._n_kv_heads_by_layer = None
|
||||
inst._kv_key_length_swa = None
|
||||
inst._kv_value_length_swa = None
|
||||
return inst
|
||||
|
||||
|
||||
|
|
@ -114,7 +118,7 @@ def _compute_max_available_ctx(native_ctx, model_gib, gpus, kv_per_token_bytes =
|
|||
model_size = int(model_gib * GIB)
|
||||
|
||||
inst._estimate_kv_cache_bytes = (
|
||||
lambda n, _t = None: 0 if n <= 0 else n * kv_per_token_bytes
|
||||
lambda n, _t = None, **_kw: 0 if n <= 0 else n * kv_per_token_bytes
|
||||
)
|
||||
inst._can_estimate_kv = lambda: True
|
||||
|
||||
|
|
|
|||
189
studio/backend/tests/test_llama_server_args.py
Normal file
189
studio/backend/tests/test_llama_server_args.py
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit tests for the llama-server pass-through args validator.
|
||||
|
||||
The validator is the security boundary between user-supplied CLI / HTTP
|
||||
input and the llama-server subprocess command. These tests pin the
|
||||
denylist behavior so the boundary doesn't quietly regress when new
|
||||
managed flags are added.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from core.inference.llama_server_args import (
|
||||
is_managed_flag,
|
||||
validate_extra_args,
|
||||
)
|
||||
|
||||
|
||||
# ── Pass-through (allowed) ───────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args",
|
||||
[
|
||||
# Sampling
|
||||
["--top-k", "20"],
|
||||
["--top-p", "0.9", "--min-p", "0.05"],
|
||||
["--seed", "-1"], # negative value, not a flag
|
||||
["--temp", "0.0"],
|
||||
["--repeat-penalty", "1.05"],
|
||||
["--mirostat", "2", "--mirostat-lr", "0.1"],
|
||||
["--xtc-probability", "0.05", "--xtc-threshold", "0.1"],
|
||||
["--dry-multiplier", "0.5"],
|
||||
# Tier-2 knobs that map to LoadRequest fields
|
||||
["--cache-type-k", "q8_0"],
|
||||
["--cache-type-v", "q8_0"],
|
||||
["--chat-template-file", "/tmp/tpl.jinja"],
|
||||
["--chat-template-kwargs", '{"reasoning_effort":"high"}'],
|
||||
["--spec-type", "ngram-mod"],
|
||||
["--spec-default"],
|
||||
# Reasoning controls
|
||||
["--reasoning-format", "deepseek"],
|
||||
["-rea", "auto"],
|
||||
# Soft-managed flags the user may want to override on the CLI;
|
||||
# llama.cpp's last-wins parsing means these win over Studio's
|
||||
# auto-set version.
|
||||
["-c", "131072"],
|
||||
["--ctx-size", "8192"],
|
||||
["--parallel", "1"],
|
||||
["-np", "8"],
|
||||
["--flash-attn", "off"],
|
||||
["-fa", "on"],
|
||||
["--no-context-shift"],
|
||||
["--context-shift"],
|
||||
["--jinja"],
|
||||
["--no-jinja"],
|
||||
["-ngl", "-1"],
|
||||
["--gpu-layers", "32"],
|
||||
["-t", "16"],
|
||||
["--threads", "32"],
|
||||
["-fit", "off"],
|
||||
["--fit", "on"],
|
||||
["--fit-ctx", "8192"],
|
||||
],
|
||||
)
|
||||
def test_pass_through_allowed(args):
|
||||
assert validate_extra_args(args) == args
|
||||
|
||||
|
||||
def test_none_returns_empty_list():
|
||||
assert validate_extra_args(None) == []
|
||||
|
||||
|
||||
def test_empty_list_returns_empty_list():
|
||||
assert validate_extra_args([]) == []
|
||||
|
||||
|
||||
def test_value_with_equals_form_passes_through():
|
||||
assert validate_extra_args(["--top-k=20"]) == ["--top-k=20"]
|
||||
|
||||
|
||||
def test_non_flag_token_passes_through():
|
||||
# A bare positional value (not preceded by a flag) is preserved
|
||||
# verbatim. llama-server may reject it, but that's not our job.
|
||||
assert validate_extra_args(["foo"]) == ["foo"]
|
||||
|
||||
|
||||
# ── Denylist (rejected) ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"denied",
|
||||
[
|
||||
# Model identity
|
||||
"-m",
|
||||
"--model",
|
||||
"-hf",
|
||||
"-hfr",
|
||||
"--hf-repo",
|
||||
"-hff",
|
||||
"--hf-file",
|
||||
"-hft",
|
||||
"--hf-token",
|
||||
"-mm",
|
||||
"--mmproj",
|
||||
"--mmproj-url",
|
||||
# Networking (Studio binds + proxies)
|
||||
"--host",
|
||||
"--port",
|
||||
"--path",
|
||||
"--api-prefix",
|
||||
"--reuse-port",
|
||||
# Auth / TLS
|
||||
"--api-key",
|
||||
"--api-key-file",
|
||||
"--ssl-key-file",
|
||||
"--ssl-cert-file",
|
||||
# Single-model server
|
||||
"--webui",
|
||||
"--no-webui",
|
||||
"--models-dir",
|
||||
"--models-max",
|
||||
],
|
||||
)
|
||||
def test_denylist_rejects_all_aliases(denied):
|
||||
with pytest.raises(ValueError, match = denied):
|
||||
validate_extra_args([denied, "value"])
|
||||
|
||||
|
||||
def test_denylist_rejects_equals_form():
|
||||
with pytest.raises(ValueError, match = "--port"):
|
||||
validate_extra_args(["--port=9000"])
|
||||
|
||||
|
||||
def test_denylist_rejects_short_form_when_long_is_denied():
|
||||
# -m is the short form of the hard-denied --model; rejecting only
|
||||
# the long form would leave a trivial bypass.
|
||||
with pytest.raises(ValueError, match = "-m"):
|
||||
validate_extra_args(["-m", "/some/other/path.gguf"])
|
||||
|
||||
|
||||
def test_denylist_message_names_offending_flag():
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
validate_extra_args(["--top-k", "20", "--api-key", "secret"])
|
||||
assert "--api-key" in str(excinfo.value)
|
||||
|
||||
|
||||
def test_first_denied_flag_short_circuits():
|
||||
# Validation stops at the first denied flag; later denied flags
|
||||
# in the same call don't matter for behaviour, but the message
|
||||
# should name the first one we hit.
|
||||
with pytest.raises(ValueError, match = "--port"):
|
||||
validate_extra_args(["--port", "1", "--host", "x"])
|
||||
|
||||
|
||||
# ── Numeric values that look flag-ish ─────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["-1", "-0.5", "-42", "-.5"])
|
||||
def test_negative_number_value_is_not_flag(value):
|
||||
# ``--seed -1`` is a value, not a flag. Validator must not try
|
||||
# to look up "-1" in the denylist.
|
||||
assert validate_extra_args(["--seed", value]) == ["--seed", value]
|
||||
|
||||
|
||||
# ── is_managed_flag helper ───────────────────────────────────────────
|
||||
|
||||
|
||||
def test_is_managed_flag_true_for_denied():
|
||||
assert is_managed_flag("--port") is True
|
||||
assert is_managed_flag("--api-key") is True
|
||||
assert is_managed_flag("-m") is True
|
||||
assert is_managed_flag("--model") is True
|
||||
|
||||
|
||||
def test_is_managed_flag_false_for_pass_through():
|
||||
assert is_managed_flag("--top-k") is False
|
||||
assert is_managed_flag("--cache-type-k") is False
|
||||
assert is_managed_flag("--chat-template-file") is False
|
||||
# Soft-managed flags pass through (last-wins override)
|
||||
assert is_managed_flag("-c") is False
|
||||
assert is_managed_flag("--ctx-size") is False
|
||||
assert is_managed_flag("--parallel") is False
|
||||
assert is_managed_flag("--flash-attn") is False
|
||||
assert is_managed_flag("-ngl") is False
|
||||
assert is_managed_flag("--threads") is False
|
||||
|
|
@ -320,11 +320,23 @@ class TestPydanticModels:
|
|||
"""Field exists in InferenceStatusResponse.model_fields."""
|
||||
assert "native_context_length" in InferenceStatusResponse.model_fields
|
||||
|
||||
def test_status_response_has_chat_template_field(self):
|
||||
"""Status includes chat_template so the UI can rehydrate after refresh."""
|
||||
assert "chat_template" in InferenceStatusResponse.model_fields
|
||||
|
||||
def test_status_response_defaults_none(self):
|
||||
"""Omitting native_context_length defaults to None."""
|
||||
resp = InferenceStatusResponse()
|
||||
assert resp.native_context_length is None
|
||||
|
||||
def test_status_response_chat_template_roundtrip(self):
|
||||
"""chat_template serializes and validates as part of status."""
|
||||
resp = InferenceStatusResponse(chat_template = "{{ messages }}")
|
||||
roundtripped = InferenceStatusResponse.model_validate_json(
|
||||
resp.model_dump_json()
|
||||
)
|
||||
assert roundtripped.chat_template == "{{ messages }}"
|
||||
|
||||
def test_roundtrip_preserves_value(self):
|
||||
"""model_validate_json(model_dump_json()) round-trips."""
|
||||
resp = LoadResponse(
|
||||
|
|
|
|||
|
|
@ -144,13 +144,14 @@ class TestChatMessageToolRoles:
|
|||
|
||||
# ── Role-aware content requirements ────────────────────────────
|
||||
|
||||
def test_user_empty_content_rejected(self):
|
||||
with pytest.raises(ValidationError):
|
||||
ChatMessage(role = "user", content = "")
|
||||
@pytest.mark.parametrize("role", ["user", "system"])
|
||||
def test_empty_string_content_allowed(self, role):
|
||||
msg = ChatMessage(role = role, content = "")
|
||||
assert msg.content == ""
|
||||
|
||||
def test_system_empty_content_rejected(self):
|
||||
def test_user_missing_content_rejected(self):
|
||||
with pytest.raises(ValidationError):
|
||||
ChatMessage(role = "system", content = "")
|
||||
ChatMessage(role = "user")
|
||||
|
||||
def test_user_empty_list_content_rejected(self):
|
||||
with pytest.raises(ValidationError):
|
||||
|
|
@ -226,6 +227,14 @@ class TestChatCompletionRequestToolFields:
|
|||
assert len(req.tools) == 1
|
||||
assert req.tools[0]["function"]["name"] == "get_weather"
|
||||
|
||||
def test_image_base64_allows_empty_user_text(self):
|
||||
req = ChatCompletionRequest(
|
||||
messages = [{"role": "user", "content": ""}],
|
||||
image_base64 = "aW1hZ2U=",
|
||||
)
|
||||
assert req.messages[0].content == ""
|
||||
assert req.image_base64 == "aW1hZ2U="
|
||||
|
||||
def test_tool_choice_string_auto(self):
|
||||
assert self._make(tool_choice = "auto").tool_choice == "auto"
|
||||
|
||||
|
|
|
|||
56
studio/backend/tests/test_tool_policy_gates.py
Normal file
56
studio/backend/tests/test_tool_policy_gates.py
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""
|
||||
Tests for `_effective_enable_tools` -- the helper that folds the
|
||||
process-level `tool_policy` over a request's `enable_tools` field.
|
||||
|
||||
Truth table (policy x payload.enable_tools -> effective):
|
||||
policy=None + payload=None -> None
|
||||
policy=None + payload=True -> True
|
||||
policy=None + payload=False -> False
|
||||
policy=True + payload=* -> True
|
||||
policy=False + payload=* -> False
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
_backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
sys.path.insert(0, _backend)
|
||||
|
||||
import pytest
|
||||
|
||||
from routes.inference import _effective_enable_tools
|
||||
from state.tool_policy import reset_tool_policy, set_tool_policy
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _reset():
|
||||
reset_tool_policy()
|
||||
yield
|
||||
reset_tool_policy()
|
||||
|
||||
|
||||
def _payload(value):
|
||||
return SimpleNamespace(enable_tools = value)
|
||||
|
||||
|
||||
class TestEffectiveEnableTools:
|
||||
@pytest.mark.parametrize(
|
||||
"payload_value,expected",
|
||||
[(None, None), (True, True), (False, False)],
|
||||
)
|
||||
def test_no_policy_falls_through_to_payload(self, payload_value, expected):
|
||||
assert _effective_enable_tools(_payload(payload_value)) == expected
|
||||
|
||||
@pytest.mark.parametrize("payload_value", [None, True, False])
|
||||
def test_policy_true_overrides_any_payload(self, payload_value):
|
||||
set_tool_policy(True)
|
||||
assert _effective_enable_tools(_payload(payload_value)) is True
|
||||
|
||||
@pytest.mark.parametrize("payload_value", [None, True, False])
|
||||
def test_policy_false_overrides_any_payload(self, payload_value):
|
||||
set_tool_policy(False)
|
||||
assert _effective_enable_tools(_payload(payload_value)) is False
|
||||
59
studio/backend/tests/test_tool_policy_state.py
Normal file
59
studio/backend/tests/test_tool_policy_state.py
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||||
|
||||
"""
|
||||
Tests for the process-level server-side tool policy used by `unsloth run`.
|
||||
|
||||
The policy has three states:
|
||||
None -> no CLI override (default; honor per-request enable_tools)
|
||||
True -> CLI forced tools on
|
||||
False -> CLI forced tools off
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
_backend = os.path.join(os.path.dirname(__file__), "..")
|
||||
sys.path.insert(0, _backend)
|
||||
|
||||
import pytest
|
||||
|
||||
from state.tool_policy import (
|
||||
get_tool_policy,
|
||||
reset_tool_policy,
|
||||
set_tool_policy,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _reset():
|
||||
reset_tool_policy()
|
||||
yield
|
||||
reset_tool_policy()
|
||||
|
||||
|
||||
class TestToolPolicy:
|
||||
def test_default_is_none(self):
|
||||
assert get_tool_policy() is None
|
||||
|
||||
def test_set_true_then_get(self):
|
||||
set_tool_policy(True)
|
||||
assert get_tool_policy() is True
|
||||
|
||||
def test_set_false_then_get(self):
|
||||
set_tool_policy(False)
|
||||
assert get_tool_policy() is False
|
||||
|
||||
def test_set_none_clears(self):
|
||||
set_tool_policy(True)
|
||||
set_tool_policy(None)
|
||||
assert get_tool_policy() is None
|
||||
|
||||
def test_reset_clears(self):
|
||||
set_tool_policy(False)
|
||||
reset_tool_policy()
|
||||
assert get_tool_policy() is None
|
||||
|
||||
def test_rejects_non_optional_bool(self):
|
||||
with pytest.raises(TypeError):
|
||||
set_tool_policy("true") # type: ignore[arg-type]
|
||||
|
|
@ -133,6 +133,22 @@ def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch)
|
|||
)
|
||||
|
||||
|
||||
def test_causal_conv1d_fast_path_includes_qwen3_6_variants(monkeypatch):
|
||||
install_mock = mock.Mock(return_value = True)
|
||||
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
||||
|
||||
worker._ensure_causal_conv1d_fast_path(
|
||||
event_queue = [],
|
||||
model_name = "unsloth/Qwen3.6-4B",
|
||||
)
|
||||
worker._ensure_causal_conv1d_fast_path(
|
||||
event_queue = [],
|
||||
model_name = "unsloth/Qwen3_6-4B",
|
||||
)
|
||||
|
||||
assert install_mock.call_count == 2
|
||||
|
||||
|
||||
def test_mamba_ssm_path_preserves_wheel_first_install_args(monkeypatch):
|
||||
install_mock = mock.Mock(return_value = True)
|
||||
monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -364,6 +364,10 @@ TEMPLATE_TO_MODEL_MAPPER = {
|
|||
"unsloth/Qwen3-4B-Thinking-2507-bnb-4bit",
|
||||
"unsloth/Qwen3-30B-A3B-Thinking-2507",
|
||||
"Qwen/Qwen3-30B-A3B-Thinking-2507",
|
||||
"Qwen/Qwen3.6-35B-A3B",
|
||||
"unsloth/Qwen3.6-35B-A3B",
|
||||
"Qwen/Qwen3.6-27B",
|
||||
"unsloth/Qwen3.6-27B",
|
||||
),
|
||||
"qwen3.5": (
|
||||
"unsloth/Qwen3.5-0.8B",
|
||||
|
|
|
|||
|
|
@ -33,7 +33,13 @@ Non-quantizable = 2*H*L + V*H + (V*H if not tie_embeddings else 0)
|
|||
| QLoRA 4-bit | `Quantizable * 2 / 3.2 + Non-quantizable * 2` |
|
||||
| LoRA / Full fp16 | `(Quantizable + Non-quantizable) * 2` |
|
||||
|
||||
The 3.2 factor (`16/5`) accounts for BNB NF4 blockwise scales.
|
||||
The 3.2 factor (`16/5`) accounts for BNB NF4 blockwise scales. Repos whose
|
||||
quantization config enables `bnb_4bit_use_double_quant` use a tighter, still
|
||||
conservative 3.6 factor for the quantized portion of the weights.
|
||||
When a 4-bit config has `llm_int8_skip_modules` entries that point to language
|
||||
model layers or submodules, those quantizable weights are charged at fp16
|
||||
instead of NF4. Generic embedding and multimodal skip names are already covered
|
||||
by non-quantizable terms or excluded from text training weights.
|
||||
|
||||
## 2. LoRA Adapters
|
||||
|
||||
|
|
@ -53,6 +59,18 @@ MLP modules multiply by `E` for MoE.
|
|||
LoRA_bytes = sum(A + B per selected module) * L * 2
|
||||
```
|
||||
|
||||
`all-linear` is treated as all known text linear modules in the table above.
|
||||
The estimator deliberately does not infer multimodal or vision-tower LoRA
|
||||
modules from config shapes; those modules vary too much across VLM families for
|
||||
a generic config formula.
|
||||
|
||||
Some decoder configs expose layer-shape fields such as `layer_types`,
|
||||
`head_dim`, `global_head_dim`, `num_global_key_value_heads`, `attention_k_eq_v`,
|
||||
`num_kv_shared_layers`, `use_double_wide_mlp`, `vocab_size_per_layer_input`, and
|
||||
`hidden_size_per_layer_input`. When those fields are present, the estimator
|
||||
derives text weight and LoRA counts from the per-layer shapes instead of
|
||||
assuming every layer has the same seven projection modules.
|
||||
|
||||
## 3. Optimizer States (calibrated)
|
||||
|
||||
| Optimizer | Bytes/param | Notes |
|
||||
|
|
@ -77,6 +95,21 @@ Per-layer (from `unsloth_zoo/vllm_utils.py`):
|
|||
Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25
|
||||
```
|
||||
|
||||
When the resolved attention implementation is none of `flash_attention_2`,
|
||||
`sdpa`, or `flex_attention` (PyTorch SDPA dispatches to flash or
|
||||
memory-efficient kernels and FlexAttention is also a memory-efficient
|
||||
kernel, all of which are O(n) in memory), activation memory also includes
|
||||
a quadratic attention-score/workspace estimate:
|
||||
|
||||
```
|
||||
Non_flash_attention = B * num_attention_heads * S^2 * 2 * 12.0 * effective_layers
|
||||
Activations = max(Per_layer_with_gc, Non_flash_attention)
|
||||
```
|
||||
|
||||
Studio resolves the attention implementation with Unsloth's
|
||||
`resolve_attention_implementation` helper and uses that result directly. The
|
||||
estimator does not duplicate model-family attention policy.
|
||||
|
||||
| GC Mode | Full FT | LoRA/QLoRA |
|
||||
|---------|---------|------------|
|
||||
| none | `L` layers | `L` layers |
|
||||
|
|
@ -85,13 +118,33 @@ Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25
|
|||
|
||||
## 6. Floors
|
||||
|
||||
Gradients and activations have minimum floors at **15% of model weight memory** to account for autograd overhead, attention score matrices, NCCL buffers, mixed-precision scaling, and PyTorch fragmentation.
|
||||
Activations use the computed formula directly:
|
||||
|
||||
```
|
||||
gradient_bytes = max(computed, weights * 0.15)
|
||||
activation_bytes = max(computed, weights * 0.15 * B/2)
|
||||
activation_bytes = computed_activation_bytes
|
||||
```
|
||||
|
||||
Full fine-tuning keeps the gradient floor at **15% of model weight memory** to
|
||||
account for autograd overhead, NCCL buffers, mixed-precision scaling, and
|
||||
PyTorch fragmentation:
|
||||
|
||||
```
|
||||
gradient_bytes = max(computed_gradient_bytes, weights * 0.15)
|
||||
```
|
||||
|
||||
For LoRA/QLoRA, the base model is frozen, so the weight-derived gradient floor
|
||||
is capped by trainable-state and live-activation scale:
|
||||
|
||||
```
|
||||
raw_gradient_bytes = trainable_params * 2
|
||||
gradient_floor = min(weights * 0.15, max(computed_activation_bytes, optimizer_bytes))
|
||||
gradient_bytes = max(raw_gradient_bytes, gradient_floor)
|
||||
```
|
||||
|
||||
This prevents frozen quantized model size from dominating gradient/state
|
||||
overhead when the measured runtime footprint is governed by LoRA optimizer
|
||||
states and live activations.
|
||||
|
||||
## 7. CUDA Overhead
|
||||
|
||||
**1.4 GB** fixed — CUDA driver + PyTorch runtime, calibrated on RTX 5070 Ti.
|
||||
|
|
@ -106,34 +159,6 @@ usable_gb = free[gpu_0] + sum(free[gpu_i] * 0.85 for i in 1..N)
|
|||
|
||||
---
|
||||
|
||||
## Reference Table (bsz=2, seq=2048, rank=16, GC=unsloth, adamw_8bit)
|
||||
|
||||
| Model | Weights | LoRA | Optim | Grad | Act | CUDA | Total |
|
||||
|-------|---------|------|-------|------|-----|------|-------|
|
||||
| 0.5B QLoRA | 0.5 | 0.0 | 0.0 | 0.1 | 0.1 | 1.4 | **2.1** |
|
||||
| 1B QLoRA | 1.1 | 0.0 | 0.0 | 0.2 | 0.2 | 1.4 | **2.9** |
|
||||
| 3B QLoRA | 2.4 | 0.0 | 0.1 | 0.5 | 0.5 | 1.4 | **4.9** |
|
||||
| 8B QLoRA | 6.0 | 0.1 | 0.2 | 1.2 | 1.2 | 1.4 | **10.1** |
|
||||
| 8B LoRA fp16 | 15.0 | 0.1 | 0.2 | 3.0 | 3.0 | 1.4 | **22.6** |
|
||||
| 8B Full FT | 15.0 | — | 29.9 | 15.0 | 3.0 | 1.4 | **64.2** |
|
||||
| 32B LoRA fp16 | 61.0 | 0.2 | 0.5 | 12.2 | 12.2 | 1.4 | **87.6** |
|
||||
| 72B QLoRA | 45.5 | 0.4 | 0.8 | 9.1 | 9.1 | 1.4 | **66.3** |
|
||||
|
||||
## E2E Validation (Llama-3.2-1B, B200 emulating 24GB)
|
||||
|
||||
| Config | Estimated | Actual (nvsmi) | Error |
|
||||
|--------|----------|----------------|-------|
|
||||
| QLoRA bsz=2 seq=512 | 2.55 GB | 2.65 GB | -3.7% |
|
||||
| QLoRA bsz=2 seq=2048 | 2.60 GB | 2.65 GB | -1.8% |
|
||||
| QLoRA bsz=4 seq=2048 | 2.65 GB | 2.65 GB | +0.0% |
|
||||
| LoRA fp16 bsz=2 | 3.84 GB | 3.88 GB | -1.0% |
|
||||
| Full FT adamw_8bit | 10.89 GB | 10.80 GB | +0.8% |
|
||||
| Full FT adamw_torch | 13.19 GB | 12.93 GB | +2.0% |
|
||||
|
||||
*Note: e2e numbers predate the 15% floors, which add safety margin on top.*
|
||||
|
||||
---
|
||||
|
||||
## Parameter Flow
|
||||
|
||||
```
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import subprocess
|
|||
from typing import Any, Optional
|
||||
|
||||
from loggers import get_logger
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
|
@ -28,6 +29,7 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = timeout,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
logger.warning("amd-smi query failed: %s", e)
|
||||
|
|
|
|||
|
|
@ -774,6 +774,34 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
|
|||
return None
|
||||
|
||||
|
||||
def _determine_attention_impl_for_gpu_estimate(config) -> str:
|
||||
import copy as _copy
|
||||
|
||||
from unsloth.models._utils import resolve_attention_implementation
|
||||
from transformers import AutoModel, AutoModelForCausalLM
|
||||
|
||||
# why: resolve_attention_implementation calls _set_attn_impl which writes
|
||||
# _attn_implementation onto the config; PreTrainedConfig's setter walks
|
||||
# `sub_configs` and propagates to nested text_config / sub-configs, so a
|
||||
# shallow copy still mutates those shared inner objects on the cached
|
||||
# config returned by _load_config_for_gpu_estimate. Deepcopy isolates them.
|
||||
config_copy = _copy.deepcopy(config)
|
||||
|
||||
model_class = None
|
||||
for auto_model in (AutoModelForCausalLM, AutoModel):
|
||||
mapping = getattr(auto_model, "_model_mapping", None)
|
||||
if mapping is None:
|
||||
continue
|
||||
try:
|
||||
if config_copy.__class__ in mapping:
|
||||
model_class = mapping[config_copy.__class__]
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
return resolve_attention_implementation(model_class, config_copy)
|
||||
|
||||
|
||||
def _estimate_fp16_model_size_bytes_from_config(config) -> Optional[int]:
|
||||
from .vram_estimation import extract_arch_config, compute_total_params
|
||||
|
||||
|
|
@ -844,12 +872,21 @@ def estimate_fp16_model_size_bytes(
|
|||
return int(total_params * 2), "safetensors"
|
||||
|
||||
config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token)
|
||||
config_bytes: Optional[int] = None
|
||||
if config is not None:
|
||||
config_bytes = _estimate_fp16_model_size_bytes_from_config(config)
|
||||
if config_bytes is not None:
|
||||
return config_bytes, "config"
|
||||
|
||||
local_bytes = _get_local_weight_size_bytes(estimate_model)
|
||||
|
||||
# why: config-derived bytes cover only the text tower; local safetensors
|
||||
# include vision/audio towers. Take the larger so the multimodal
|
||||
# extra_bytes correction can fire.
|
||||
if config_bytes is not None and local_bytes is not None:
|
||||
if local_bytes > config_bytes:
|
||||
return local_bytes, "weight_bytes"
|
||||
return config_bytes, "config"
|
||||
if config_bytes is not None:
|
||||
return config_bytes, "config"
|
||||
if local_bytes is not None:
|
||||
return local_bytes, "weight_bytes"
|
||||
|
||||
|
|
@ -877,6 +914,9 @@ def estimate_required_model_memory_gb(
|
|||
TrainingVramConfig,
|
||||
extract_arch_config,
|
||||
estimate_training_vram,
|
||||
compute_total_params,
|
||||
compute_optimizer_bytes,
|
||||
compute_gradient_bytes,
|
||||
CUDA_OVERHEAD_BYTES,
|
||||
QUANT_4BIT_FACTOR,
|
||||
DEFAULT_TARGET_MODULES,
|
||||
|
|
@ -926,13 +966,44 @@ def estimate_required_model_memory_gb(
|
|||
model_name, hf_token = hf_token
|
||||
)
|
||||
config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token)
|
||||
if config is not None:
|
||||
try:
|
||||
vram_config.attention_implementation = (
|
||||
_determine_attention_impl_for_gpu_estimate(config)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Could not resolve attention implementation for '%s': %s",
|
||||
estimate_model,
|
||||
e,
|
||||
)
|
||||
# why: if we cannot prove flash attention is usable, charge the
|
||||
# quadratic non-flash activation path so GPU selection stays
|
||||
# conservative.
|
||||
vram_config.attention_implementation = "eager"
|
||||
arch = extract_arch_config(config) if config is not None else None
|
||||
|
||||
if arch is not None:
|
||||
breakdown = estimate_training_vram(arch, vram_config)
|
||||
# why: extract_arch_config only sees text_config; safetensors include
|
||||
# vision/audio tower bytes that the text-arch fp16 total misses.
|
||||
arch_fp16_bytes = compute_total_params(arch) * 2
|
||||
extra_bytes = max(0, int(model_size_bytes) - arch_fp16_bytes)
|
||||
if extra_bytes > 0:
|
||||
breakdown.model_weights += extra_bytes
|
||||
if training_method == "full":
|
||||
# why: full fine-tuning makes the extra (vision/audio) params
|
||||
# trainable; optimizer + gradient bytes scale with them too.
|
||||
extra_params = extra_bytes // 2
|
||||
breakdown.optimizer_states += compute_optimizer_bytes(
|
||||
extra_params,
|
||||
vram_config.optimizer,
|
||||
)
|
||||
breakdown.gradients += compute_gradient_bytes(extra_params)
|
||||
required_gb = breakdown.total / (1024**3)
|
||||
metadata["required_gb"] = round(required_gb, 3)
|
||||
metadata["estimation_mode"] = "detailed"
|
||||
metadata["attention_implementation"] = vram_config.attention_implementation
|
||||
metadata["vram_breakdown"] = breakdown.to_gb_dict()
|
||||
max_gpus = max(1, get_visible_gpu_count())
|
||||
for n_gpus in range(1, max_gpus + 1):
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import Any, Optional
|
|||
|
||||
from loggers import get_logger
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -65,6 +66,7 @@ def get_physical_gpu_count() -> Optional[int]:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode == 0 and result.stdout.strip():
|
||||
|
|
@ -90,6 +92,7 @@ def get_primary_gpu_utilization() -> dict[str, Any]:
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
|
|
@ -141,6 +144,7 @@ def get_visible_gpu_utilization(
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 5,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
|
|
@ -227,6 +231,7 @@ def get_backend_visible_gpu_info(
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -32,6 +32,7 @@ import threading
|
|||
import yaml
|
||||
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -583,6 +584,7 @@ def _is_vision_model_subprocess(
|
|||
capture_output = True,
|
||||
text = True,
|
||||
timeout = 60,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
|
||||
|
|
@ -1227,9 +1229,11 @@ def _resolve_gguf_dir(p: Path) -> Optional[Path]:
|
|||
return p
|
||||
if p.is_file() and p.suffix.lower() == ".gguf":
|
||||
parent = p.parent
|
||||
if (parent / "config.json").exists() or (
|
||||
parent / "adapter_config.json"
|
||||
).exists():
|
||||
if (
|
||||
(parent / "config.json").exists()
|
||||
or (parent / "adapter_config.json").exists()
|
||||
or (parent / "export_metadata.json").exists()
|
||||
):
|
||||
return parent
|
||||
return None
|
||||
|
||||
|
|
|
|||
406
studio/backend/utils/native_path_leases.py
Normal file
406
studio/backend/utils/native_path_leases.py
Normal file
|
|
@ -0,0 +1,406 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Verification for Tauri native path signed grants.
|
||||
|
||||
Rust signs compact ``base64url(payload_json).base64url(hmac)`` grants. The
|
||||
frontend can see and forward the grant, but cannot change it without breaking
|
||||
the HMAC. The backend verifies the original payload segment bytes, then
|
||||
re-stats the path before any native read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import stat as _stat_module
|
||||
import threading
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, Iterator, Mapping
|
||||
|
||||
LEASE_SECRET_ENV = "UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET"
|
||||
_MAX_NATIVE_PATH_REDACTIONS = 100
|
||||
_MAX_NATIVE_PATH_LABELS = 10_000
|
||||
_MIN_LEASE_SECRET_BYTES = 32
|
||||
|
||||
_REPLAY_LOCK = threading.Lock()
|
||||
_USED_NONCES: dict[str, int] = {}
|
||||
_REDACTION_LOCK = threading.Lock()
|
||||
_NATIVE_PATH_REDACTIONS: list[str] = []
|
||||
_NATIVE_PATH_LABELS: dict[str, str] = {}
|
||||
_NATIVE_PATH_ENV_LOCK = threading.Lock()
|
||||
_SECRET_INIT_LOCK = threading.Lock()
|
||||
_CACHED_LEASE_SECRET: bytes | None = None
|
||||
_SCRUB_REFCOUNT = 0
|
||||
_SCRUB_SAVED_SECRET: str | None = None
|
||||
|
||||
|
||||
class NativePathLeaseError(ValueError):
|
||||
"""Raised when a native path grant is missing, invalid, or unsafe."""
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class NativePathGrant:
|
||||
operation: str
|
||||
canonical_path: Path
|
||||
path_kind: str
|
||||
path_type: str
|
||||
source_kind: str
|
||||
token_id_hash: str
|
||||
display_label: str
|
||||
expires_at_ms: int
|
||||
size_bytes: int | None
|
||||
modified_ms: int | None
|
||||
|
||||
|
||||
def native_path_leases_supported() -> bool:
|
||||
try:
|
||||
_decode_secret()
|
||||
except NativePathLeaseError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def child_env_without_native_path_secret(
|
||||
env: Mapping[str, str] | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""Return a child-process env with the native path lease secret removed."""
|
||||
|
||||
if env is None:
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
cleaned = dict(os.environ)
|
||||
else:
|
||||
cleaned = dict(env)
|
||||
cleaned.pop(LEASE_SECRET_ENV, None)
|
||||
return cleaned
|
||||
|
||||
|
||||
def run_without_native_path_secret(
|
||||
target: Callable[..., Any],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Run a multiprocessing child target without the native path lease secret."""
|
||||
|
||||
global _CACHED_LEASE_SECRET, _SCRUB_SAVED_SECRET
|
||||
os.environ.pop(LEASE_SECRET_ENV, None)
|
||||
_CACHED_LEASE_SECRET = None
|
||||
_SCRUB_SAVED_SECRET = None
|
||||
return target(*args, **kwargs)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def native_path_secret_removed_for_child_start() -> Iterator[None]:
|
||||
global _SCRUB_REFCOUNT, _SCRUB_SAVED_SECRET, _CACHED_LEASE_SECRET
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
if _SCRUB_REFCOUNT == 0:
|
||||
_SCRUB_SAVED_SECRET = os.environ.pop(LEASE_SECRET_ENV, None)
|
||||
_CACHED_LEASE_SECRET = None
|
||||
_SCRUB_REFCOUNT += 1
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
_SCRUB_REFCOUNT -= 1
|
||||
if _SCRUB_REFCOUNT == 0 and _SCRUB_SAVED_SECRET is not None:
|
||||
os.environ[LEASE_SECRET_ENV] = _SCRUB_SAVED_SECRET
|
||||
_SCRUB_SAVED_SECRET = None
|
||||
|
||||
|
||||
def verify_native_path_lease(
|
||||
lease: str | None,
|
||||
*,
|
||||
operation: str,
|
||||
expected_kind: str | None = None,
|
||||
expected_path_type: str | None = None,
|
||||
allowed_suffixes: Iterable[str] | None = None,
|
||||
) -> NativePathGrant:
|
||||
if not lease:
|
||||
raise NativePathLeaseError("Native path grant is required.")
|
||||
|
||||
secret = _decode_secret()
|
||||
payload_b64, signature_b64 = _split_lease(lease)
|
||||
expected_signature = hmac.new(
|
||||
secret,
|
||||
payload_b64.encode("ascii"),
|
||||
hashlib.sha256,
|
||||
).digest()
|
||||
supplied_signature = _b64decode(signature_b64)
|
||||
if not hmac.compare_digest(expected_signature, supplied_signature):
|
||||
raise NativePathLeaseError("Native path grant signature is invalid.")
|
||||
|
||||
payload = _decode_payload(payload_b64)
|
||||
_validate_payload(payload, operation = operation, expected_kind = expected_kind)
|
||||
|
||||
path = Path(str(payload["canonical_path"]))
|
||||
_reject_network_or_device_path(path)
|
||||
try:
|
||||
signed_lstat = os.lstat(path)
|
||||
except OSError as exc:
|
||||
raise NativePathLeaseError("Native path is no longer accessible.") from exc
|
||||
if _stat_module.S_ISLNK(signed_lstat.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a regular file.")
|
||||
try:
|
||||
resolved = path.resolve(strict = True)
|
||||
except OSError as exc:
|
||||
raise NativePathLeaseError("Native path is no longer accessible.") from exc
|
||||
_reject_network_or_device_path(resolved)
|
||||
if not _same_native_path(resolved, path):
|
||||
raise NativePathLeaseError(
|
||||
"Native path grant no longer resolves to the selected path."
|
||||
)
|
||||
|
||||
grant = NativePathGrant(
|
||||
operation = str(payload["operation"]),
|
||||
canonical_path = resolved,
|
||||
path_kind = str(payload["path_kind"]),
|
||||
path_type = str(payload["path_type"]),
|
||||
source_kind = str(payload["source_kind"]),
|
||||
token_id_hash = str(payload["token_id_hash"]),
|
||||
display_label = str(payload.get("display_label") or resolved.name),
|
||||
expires_at_ms = _required_int(payload, "expires_at_ms"),
|
||||
size_bytes = _optional_int(payload.get("size_bytes")),
|
||||
modified_ms = _optional_int(payload.get("modified_ms")),
|
||||
)
|
||||
|
||||
if expected_path_type and grant.path_type != expected_path_type:
|
||||
raise NativePathLeaseError("Native path grant has the wrong path type.")
|
||||
suffixes = tuple(s.lower() for s in (allowed_suffixes or ()))
|
||||
if suffixes and resolved.suffix.lower() not in suffixes:
|
||||
raise NativePathLeaseError("Native path grant has an unsupported file type.")
|
||||
|
||||
_validate_current_stat(grant)
|
||||
_consume_nonce(str(payload["nonce"]), grant.expires_at_ms)
|
||||
_remember_native_path_for_redaction(str(resolved), grant.display_label)
|
||||
return grant
|
||||
|
||||
|
||||
def display_label_for_native_path(value: str | None) -> str | None:
|
||||
if not value:
|
||||
return value
|
||||
with _REDACTION_LOCK:
|
||||
return _NATIVE_PATH_LABELS.get(value, value)
|
||||
|
||||
|
||||
def is_registered_native_path_label(path_value: str | None, label: str | None) -> bool:
|
||||
if not path_value or not label:
|
||||
return False
|
||||
with _REDACTION_LOCK:
|
||||
return _NATIVE_PATH_LABELS.get(path_value) == label
|
||||
|
||||
|
||||
def redact_native_paths(value: str) -> str:
|
||||
with _REDACTION_LOCK:
|
||||
paths = sorted(_NATIVE_PATH_REDACTIONS, key = len, reverse = True)
|
||||
redacted = value
|
||||
for path in paths:
|
||||
for variant in {path, path.replace("/", "\\"), path.replace("\\", "/")}:
|
||||
if variant:
|
||||
redacted = redacted.replace(variant, "<native_path>")
|
||||
return redacted
|
||||
|
||||
|
||||
def _decode_secret() -> bytes:
|
||||
global _CACHED_LEASE_SECRET
|
||||
if _CACHED_LEASE_SECRET is not None:
|
||||
return _CACHED_LEASE_SECRET
|
||||
with _SECRET_INIT_LOCK:
|
||||
if _CACHED_LEASE_SECRET is not None:
|
||||
return _CACHED_LEASE_SECRET
|
||||
with _NATIVE_PATH_ENV_LOCK:
|
||||
encoded = os.environ.get(LEASE_SECRET_ENV)
|
||||
if encoded is None and _SCRUB_SAVED_SECRET is not None:
|
||||
encoded = _SCRUB_SAVED_SECRET
|
||||
if not encoded:
|
||||
raise NativePathLeaseError(
|
||||
"Native path grants require the managed desktop backend."
|
||||
)
|
||||
try:
|
||||
secret = _b64decode(encoded)
|
||||
except Exception as exc:
|
||||
raise NativePathLeaseError("Native path grant secret is invalid.") from exc
|
||||
if len(secret) < _MIN_LEASE_SECRET_BYTES:
|
||||
raise NativePathLeaseError("Native path grant secret is invalid.")
|
||||
_CACHED_LEASE_SECRET = secret
|
||||
return secret
|
||||
|
||||
|
||||
def _split_lease(lease: str) -> tuple[str, str]:
|
||||
if not isinstance(lease, str):
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.")
|
||||
try:
|
||||
lease.encode("ascii")
|
||||
except UnicodeEncodeError as exc:
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.") from exc
|
||||
parts = lease.split(".")
|
||||
if len(parts) != 2 or not parts[0] or not parts[1]:
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.")
|
||||
return parts[0], parts[1]
|
||||
|
||||
|
||||
def _decode_payload(payload_b64: str) -> dict[str, Any]:
|
||||
try:
|
||||
payload = json.loads(_b64decode(payload_b64).decode("utf-8"))
|
||||
except Exception as exc:
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.") from exc
|
||||
if not isinstance(payload, dict):
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.")
|
||||
return payload
|
||||
|
||||
|
||||
def _validate_payload(
|
||||
payload: dict[str, Any], *, operation: str, expected_kind: str | None
|
||||
) -> None:
|
||||
required = (
|
||||
"version",
|
||||
"operation",
|
||||
"canonical_path",
|
||||
"path_kind",
|
||||
"path_type",
|
||||
"source_kind",
|
||||
"token_id_hash",
|
||||
"issued_at_ms",
|
||||
"expires_at_ms",
|
||||
"nonce",
|
||||
)
|
||||
missing = [key for key in required if key not in payload]
|
||||
if missing:
|
||||
raise NativePathLeaseError(
|
||||
"Native path grant payload is missing required fields."
|
||||
)
|
||||
if _required_int(payload, "version") != 1:
|
||||
raise NativePathLeaseError("Native path grant version is unsupported.")
|
||||
if payload["operation"] != operation:
|
||||
raise NativePathLeaseError("Native path grant operation is invalid.")
|
||||
if expected_kind and payload["path_kind"] != expected_kind:
|
||||
raise NativePathLeaseError("Native path grant kind is invalid.")
|
||||
now_ms = int(time.time() * 1000)
|
||||
issued_at_ms = _required_int(payload, "issued_at_ms")
|
||||
expires_at_ms = _required_int(payload, "expires_at_ms")
|
||||
if issued_at_ms >= expires_at_ms:
|
||||
raise NativePathLeaseError("Native path grant timestamps are inconsistent.")
|
||||
if expires_at_ms <= now_ms:
|
||||
raise NativePathLeaseError("Native path grant has expired.")
|
||||
if issued_at_ms > now_ms + 30_000:
|
||||
raise NativePathLeaseError("Native path grant issue time is invalid.")
|
||||
for key in ("canonical_path", "nonce", "token_id_hash", "display_label"):
|
||||
raw = payload.get(key)
|
||||
if raw is None:
|
||||
continue
|
||||
if "\x00" in str(raw):
|
||||
raise NativePathLeaseError("Native path grant contains invalid characters.")
|
||||
|
||||
|
||||
def _validate_current_stat(grant: NativePathGrant) -> None:
|
||||
try:
|
||||
st = os.lstat(grant.canonical_path)
|
||||
except OSError as exc:
|
||||
raise NativePathLeaseError("Native path is no longer accessible.") from exc
|
||||
if _stat_module.S_ISLNK(st.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a regular file.")
|
||||
if grant.path_type == "file":
|
||||
if not _stat_module.S_ISREG(st.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a regular file.")
|
||||
elif grant.path_type == "directory":
|
||||
if not _stat_module.S_ISDIR(st.st_mode):
|
||||
raise NativePathLeaseError("Native path is no longer a directory.")
|
||||
else:
|
||||
raise NativePathLeaseError("Native path grant has an unsupported path type.")
|
||||
|
||||
if grant.size_bytes is not None and st.st_size != grant.size_bytes:
|
||||
raise NativePathLeaseError("Native path changed after it was selected.")
|
||||
current_modified_ms = int(st.st_mtime_ns // 1_000_000)
|
||||
if grant.modified_ms is not None and current_modified_ms != grant.modified_ms:
|
||||
raise NativePathLeaseError("Native path changed after it was selected.")
|
||||
|
||||
|
||||
def _consume_nonce(nonce: str, expires_at_ms: int) -> None:
|
||||
now_ms = int(time.time() * 1000)
|
||||
with _REPLAY_LOCK:
|
||||
for key, expiry in list(_USED_NONCES.items()):
|
||||
if expiry <= now_ms:
|
||||
_USED_NONCES.pop(key, None)
|
||||
if nonce in _USED_NONCES:
|
||||
raise NativePathLeaseError("Native path grant was already used.")
|
||||
_USED_NONCES[nonce] = expires_at_ms
|
||||
|
||||
|
||||
def _remember_native_path_for_redaction(path: str, display_label: str) -> None:
|
||||
with _REDACTION_LOCK:
|
||||
_NATIVE_PATH_LABELS[path] = display_label
|
||||
if len(_NATIVE_PATH_LABELS) > _MAX_NATIVE_PATH_LABELS:
|
||||
excess = len(_NATIVE_PATH_LABELS) - _MAX_NATIVE_PATH_LABELS
|
||||
for stale_path in list(_NATIVE_PATH_LABELS.keys())[:excess]:
|
||||
_NATIVE_PATH_LABELS.pop(stale_path, None)
|
||||
if path in _NATIVE_PATH_REDACTIONS:
|
||||
return
|
||||
_NATIVE_PATH_REDACTIONS.append(path)
|
||||
del _NATIVE_PATH_REDACTIONS[:-_MAX_NATIVE_PATH_REDACTIONS]
|
||||
|
||||
|
||||
def _reject_network_or_device_path(path: Path) -> None:
|
||||
text = str(path)
|
||||
if os.name == "nt":
|
||||
normalized = text.replace("/", "\\").lower()
|
||||
if normalized.startswith("\\\\?\\"):
|
||||
rest = normalized[4:]
|
||||
is_local_drive = len(rest) >= 3 and rest[0].isalpha() and rest[1:3] == ":\\"
|
||||
if not is_local_drive:
|
||||
raise NativePathLeaseError(
|
||||
"Network paths are not supported for native grants."
|
||||
)
|
||||
elif normalized.startswith("\\\\"):
|
||||
raise NativePathLeaseError(
|
||||
"Network paths are not supported for native grants."
|
||||
)
|
||||
if os.name != "nt":
|
||||
for root in ("/dev", "/proc", "/sys"):
|
||||
if path.is_relative_to(root):
|
||||
raise NativePathLeaseError(
|
||||
"Device and virtual filesystem paths are not supported."
|
||||
)
|
||||
if "\x00" in text:
|
||||
raise NativePathLeaseError("Native path contains invalid characters.")
|
||||
|
||||
|
||||
def _b64decode(value: str) -> bytes:
|
||||
try:
|
||||
padding = "=" * (-len(value) % 4)
|
||||
return base64.urlsafe_b64decode((value + padding).encode("ascii"))
|
||||
except (UnicodeEncodeError, binascii.Error, ValueError) as exc:
|
||||
raise NativePathLeaseError("Native path grant has an invalid format.") from exc
|
||||
|
||||
|
||||
def _same_native_path(resolved: Path, signed: Path) -> bool:
|
||||
try:
|
||||
return resolved.samefile(signed)
|
||||
except OSError:
|
||||
return os.path.normcase(str(resolved)) == os.path.normcase(str(signed))
|
||||
|
||||
|
||||
def _optional_int(value: Any) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.") from exc
|
||||
|
||||
|
||||
def _required_int(payload: dict[str, Any], key: str) -> int:
|
||||
raw = payload.get(key)
|
||||
if raw is None:
|
||||
raise NativePathLeaseError(
|
||||
"Native path grant payload is missing required fields."
|
||||
)
|
||||
try:
|
||||
return int(raw)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise NativePathLeaseError("Native path grant payload is invalid.") from exc
|
||||
|
|
@ -36,6 +36,7 @@ import subprocess
|
|||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
from utils.subprocess_compat import (
|
||||
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
|
||||
)
|
||||
|
|
@ -63,6 +64,7 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
|
|||
TRANSFORMERS_550_MODEL_SUBSTRINGS: tuple[str, ...] = (
|
||||
"gemma-4", # Gemma-4 (E2B-it, E4B-it, 31B-it, 26B-A4B-it)
|
||||
"gemma4", # Gemma-4 alternate naming
|
||||
"qwen3.6",
|
||||
)
|
||||
|
||||
# Architecture classes / model_type values that require transformers 5.5.0.
|
||||
|
|
@ -503,6 +505,7 @@ def _install_to_dir(pkg: str, target_dir: str) -> bool:
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode == 0:
|
||||
|
|
@ -525,6 +528,7 @@ def _install_to_dir(pkg: str, target_dir: str) -> bool:
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
if result.returncode != 0:
|
||||
|
|
|
|||
|
|
@ -13,6 +13,8 @@ import urllib.error
|
|||
import urllib.request
|
||||
from typing import Callable
|
||||
|
||||
from utils.native_path_leases import child_env_without_native_path_secret
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
FLASH_ATTN_RELEASE_BASE_URL = (
|
||||
|
|
@ -59,6 +61,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non
|
|||
stderr = subprocess.PIPE,
|
||||
text = True,
|
||||
timeout = timeout,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
except subprocess.TimeoutExpired:
|
||||
return None
|
||||
|
|
@ -142,6 +145,7 @@ def install_wheel(
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
attempts.append(("uv", result))
|
||||
if result.returncode == 0:
|
||||
|
|
@ -153,6 +157,7 @@ def install_wheel(
|
|||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
env = child_env_without_native_path_secret(),
|
||||
)
|
||||
attempts.append(("pip", result))
|
||||
return attempts
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@
|
|||
"biome:fix": "biome check . --write"
|
||||
},
|
||||
"dependencies": {
|
||||
"@assistant-ui/core": "0.1.17",
|
||||
"@assistant-ui/react": "^0.12.19",
|
||||
"@assistant-ui/react-markdown": "^0.12.3",
|
||||
"@assistant-ui/react-streamdown": "^0.1.2",
|
||||
|
|
@ -42,6 +43,8 @@
|
|||
"@tanstack/react-router": "^1.159.10",
|
||||
"@tanstack/react-table": "^8.21.3",
|
||||
"@tauri-apps/api": "^2.10.1",
|
||||
"@tauri-apps/plugin-clipboard-manager": "^2.3.2",
|
||||
"@tauri-apps/plugin-notification": "^2.3.3",
|
||||
"@tauri-apps/plugin-opener": "^2.5.3",
|
||||
"@tauri-apps/plugin-process": "^2.3.1",
|
||||
"@tauri-apps/plugin-updater": "^2.10.1",
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ import {
|
|||
hasRefreshToken,
|
||||
mustChangePassword,
|
||||
refreshSession,
|
||||
tauriAutoAuth,
|
||||
} from "@/features/auth";
|
||||
|
||||
async function hasActiveSession(): Promise<boolean> {
|
||||
|
|
@ -39,7 +38,7 @@ function authRedirect(to: "/login" | "/change-password"): never {
|
|||
|
||||
export async function requireAuth(): Promise<void> {
|
||||
if (isTauri) {
|
||||
await tauriAutoAuth();
|
||||
// AppProvider owns backend startup + desktop auth; route guards run before it mounts.
|
||||
return;
|
||||
}
|
||||
|
||||
|
|
@ -59,7 +58,6 @@ export async function requireAuth(): Promise<void> {
|
|||
|
||||
export async function requireGuest(): Promise<void> {
|
||||
if (isTauri) {
|
||||
await tauriAutoAuth();
|
||||
throw redirect({ to: "/chat" });
|
||||
}
|
||||
if (!(await hasActiveSession())) return;
|
||||
|
|
@ -68,7 +66,6 @@ export async function requireGuest(): Promise<void> {
|
|||
|
||||
export async function requirePasswordChangeFlow(): Promise<void> {
|
||||
if (isTauri) {
|
||||
await tauriAutoAuth();
|
||||
throw redirect({ to: "/chat" });
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,12 +4,19 @@
|
|||
import { StartupScreen } from "@/components/tauri/startup-screen";
|
||||
import { UpdateBanner } from "@/components/tauri/update-banner";
|
||||
import { UpdateScreen } from "@/components/tauri/update-screen";
|
||||
import {
|
||||
WindowTitlebar,
|
||||
shouldUseCustomWindowTitlebar,
|
||||
} from "@/components/tauri/window-titlebar";
|
||||
import { Toaster } from "@/components/ui/sonner";
|
||||
import { useTauriBackend } from "@/hooks/use-tauri-backend";
|
||||
import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth";
|
||||
import { NativeIntentDrain } from "@/features/native-intents/native-intent-drain";
|
||||
import { useTauriBackend, type BackendStatus } from "@/hooks/use-tauri-backend";
|
||||
import { useTauriUpdate } from "@/hooks/use-tauri-update";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { useRouterState } from "@tanstack/react-router";
|
||||
import { ThemeProvider } from "next-themes";
|
||||
import { useEffect, useRef, type ReactNode } from "react";
|
||||
import { useEffect, useRef, useState, type ReactNode } from "react";
|
||||
|
||||
interface AppProviderProps {
|
||||
children: ReactNode;
|
||||
|
|
@ -19,67 +26,85 @@ interface AppProviderProps {
|
|||
// Tauri window helpers (only imported in Tauri mode)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async function showWindow(): Promise<void> {
|
||||
type TauriWindowMode = "setup" | "app";
|
||||
type WindowLayoutGuard = () => boolean;
|
||||
|
||||
async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise<void> {
|
||||
const { getCurrentWindow } = await import("@tauri-apps/api/window");
|
||||
await getCurrentWindow().show();
|
||||
}
|
||||
if (!isCurrent()) return;
|
||||
|
||||
function easeOutQuart(t: number): number {
|
||||
return 1 - (1 - t) ** 4;
|
||||
}
|
||||
|
||||
async function animateToGoldenRatio(abortRef: { current: boolean }): Promise<void> {
|
||||
const { getCurrentWindow, currentMonitor, LogicalSize } = await import("@tauri-apps/api/window");
|
||||
const win = getCurrentWindow();
|
||||
|
||||
// Ensure window is visible before resizing
|
||||
if (!isCurrent()) return;
|
||||
await win.center();
|
||||
if (!isCurrent()) return;
|
||||
await win.show();
|
||||
}
|
||||
|
||||
async function applyAppWindowLayout(isCurrent: WindowLayoutGuard): Promise<void> {
|
||||
const { getCurrentWindow, currentMonitor, LogicalSize } = await import("@tauri-apps/api/window");
|
||||
if (!isCurrent()) return;
|
||||
|
||||
const win = getCurrentWindow();
|
||||
const monitor = await currentMonitor();
|
||||
if (!monitor) return;
|
||||
if (!isCurrent()) return;
|
||||
|
||||
// Convert physical pixels to logical using scale factor
|
||||
const scale = monitor.scaleFactor;
|
||||
const screenW = monitor.size.width / scale;
|
||||
const screenH = monitor.size.height / scale;
|
||||
let finalW = 900;
|
||||
let finalH = 600;
|
||||
|
||||
// Target: 75% of screen width, golden ratio height, capped at min 900x600
|
||||
const targetW = Math.max(900, Math.round(screenW * 0.75));
|
||||
const targetH = Math.max(600, Math.round(targetW / 1.618));
|
||||
// Don't exceed screen height
|
||||
const finalH = Math.min(targetH, Math.round(screenH * 0.85));
|
||||
const finalW = targetW;
|
||||
if (monitor) {
|
||||
// Convert physical pixels to logical using scale factor
|
||||
const scale = monitor.scaleFactor;
|
||||
const screenW = monitor.size.width / scale;
|
||||
const screenH = monitor.size.height / scale;
|
||||
|
||||
// Check reduced motion preference
|
||||
const prefersReducedMotion = window.matchMedia("(prefers-reduced-motion: reduce)").matches;
|
||||
|
||||
if (prefersReducedMotion) {
|
||||
await win.setSize(new LogicalSize(finalW, finalH));
|
||||
} else {
|
||||
// Read current size instead of hardcoding — stays correct if tauri.conf.json changes
|
||||
const inner = await win.innerSize();
|
||||
const factor = await win.scaleFactor();
|
||||
const startW = Math.round(inner.width / factor);
|
||||
const startH = Math.round(inner.height / factor);
|
||||
const steps = 15;
|
||||
const stepDuration = 23; // ~350ms total
|
||||
|
||||
for (let i = 1; i <= steps; i++) {
|
||||
if (abortRef.current) return;
|
||||
const t = easeOutQuart(i / steps);
|
||||
const w = Math.round(startW + (finalW - startW) * t);
|
||||
const h = Math.round(startH + (finalH - startH) * t);
|
||||
await win.setSize(new LogicalSize(w, h));
|
||||
await new Promise((r) => setTimeout(r, stepDuration));
|
||||
}
|
||||
// Target: 75% of screen width, golden ratio height, capped at min 900x600
|
||||
finalW = Math.max(900, Math.round(screenW * 0.75));
|
||||
const targetH = Math.max(600, Math.round(finalW / 1.618));
|
||||
// Don't exceed screen height
|
||||
finalH = Math.min(targetH, Math.round(screenH * 0.85));
|
||||
}
|
||||
|
||||
if (abortRef.current) return;
|
||||
|
||||
// Apply constraints and finalize
|
||||
await win.setResizable(true);
|
||||
// Apply constraints and finalize without animating through intermediate sizes
|
||||
if (!isCurrent()) return;
|
||||
await win.setSize(new LogicalSize(finalW, finalH));
|
||||
if (!isCurrent()) return;
|
||||
await win.setSizeConstraints({ minWidth: 900, minHeight: 600 });
|
||||
if (!isCurrent()) return;
|
||||
await win.setResizable(true);
|
||||
if (!isCurrent()) return;
|
||||
await win.center();
|
||||
if (!isCurrent()) return;
|
||||
await win.show();
|
||||
}
|
||||
|
||||
async function showWindowFallback(): Promise<void> {
|
||||
const { getCurrentWindow } = await import("@tauri-apps/api/window");
|
||||
const win = getCurrentWindow();
|
||||
await win.setResizable(true);
|
||||
await win.show();
|
||||
}
|
||||
|
||||
function getTauriWindowMode(
|
||||
status: BackendStatus,
|
||||
hasEnteredAppMode: boolean,
|
||||
): TauriWindowMode | null {
|
||||
switch (status) {
|
||||
case "checking":
|
||||
return null;
|
||||
case "not-installed":
|
||||
case "installing":
|
||||
case "install-error":
|
||||
case "needs-elevation":
|
||||
case "repairing":
|
||||
case "repair-error":
|
||||
return "setup";
|
||||
case "starting":
|
||||
case "running":
|
||||
case "stopped":
|
||||
return "app";
|
||||
case "error":
|
||||
return hasEnteredAppMode ? "app" : "setup";
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
|
@ -103,6 +128,7 @@ function TauriUpdateLayer({ isExternalServer }: { isExternalServer: boolean }) {
|
|||
error={update.error}
|
||||
onRetry={update.retryUpdate}
|
||||
onSkipRestart={update.skipAndRestart}
|
||||
onCopyDiagnostics={update.copyDiagnostics}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
|
@ -112,62 +138,147 @@ function TauriUpdateLayer({ isExternalServer }: { isExternalServer: boolean }) {
|
|||
status={update.status}
|
||||
info={update.info}
|
||||
dismissed={update.dismissed}
|
||||
lastFailure={update.lastFailure}
|
||||
isExternalServer={isExternalServer}
|
||||
onInstall={update.installUpdate}
|
||||
onDismiss={update.dismiss}
|
||||
onCopyDiagnostics={update.copyDiagnostics}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
const HIDDEN_TITLEBAR_SIDEBAR_ROUTES = new Set([
|
||||
"/onboarding",
|
||||
"/login",
|
||||
"/change-password",
|
||||
"/signup",
|
||||
]);
|
||||
|
||||
function TauriWrapper({ children }: { children: ReactNode }) {
|
||||
const pathname = useRouterState({ select: (s) => s.location.pathname });
|
||||
const {
|
||||
status, logs, error, isExternalServer,
|
||||
currentStepIndex, progressDetail, elevationPackages,
|
||||
startInstall, retry, retryInstall, approveElevation,
|
||||
startInstall, retry, retryInstall, approveElevation, copyDiagnostics,
|
||||
} = useTauriBackend();
|
||||
|
||||
const hasResized = useRef(false);
|
||||
const abortRef = useRef(false);
|
||||
const appliedWindowModeRef = useRef<TauriWindowMode | null>(null);
|
||||
const hasEnteredAppModeRef = useRef(false);
|
||||
const windowLayoutGenerationRef = useRef(0);
|
||||
const [desktopAuthReady, setDesktopAuthReady] = useState(!isTauri);
|
||||
const [desktopAuthRetry, setDesktopAuthRetry] = useState(0);
|
||||
|
||||
// Show the window once the frontend mounts (for pre-running states)
|
||||
useEffect(() => {
|
||||
if (isTauri) void showWindow();
|
||||
if (!isTauri) return;
|
||||
return () => {
|
||||
windowLayoutGenerationRef.current += 1;
|
||||
appliedWindowModeRef.current = null;
|
||||
};
|
||||
}, []);
|
||||
|
||||
// Animate resize when backend becomes ready
|
||||
// Keep the Tauri window hidden during preflight, then show it centered in setup
|
||||
// mode or apply the final app layout in one instant step.
|
||||
useEffect(() => {
|
||||
if (status === "running" && !hasResized.current) {
|
||||
hasResized.current = true;
|
||||
abortRef.current = false;
|
||||
animateToGoldenRatio(abortRef).catch(async () => {
|
||||
// On failure, at minimum make the window resizable so user can fix manually
|
||||
try {
|
||||
const { getCurrentWindow } = await import("@tauri-apps/api/window");
|
||||
await getCurrentWindow().setResizable(true);
|
||||
} catch { /* swallow — window may still be functional */ }
|
||||
});
|
||||
if (!isTauri) return;
|
||||
|
||||
const nextMode = getTauriWindowMode(status, hasEnteredAppModeRef.current);
|
||||
if (!nextMode) {
|
||||
appliedWindowModeRef.current = null;
|
||||
windowLayoutGenerationRef.current += 1;
|
||||
return;
|
||||
}
|
||||
return () => { abortRef.current = true; };
|
||||
if (appliedWindowModeRef.current === nextMode) return;
|
||||
|
||||
appliedWindowModeRef.current = nextMode;
|
||||
if (nextMode === "app") hasEnteredAppModeRef.current = true;
|
||||
|
||||
const layoutGeneration = windowLayoutGenerationRef.current + 1;
|
||||
windowLayoutGenerationRef.current = layoutGeneration;
|
||||
const isCurrent = () => windowLayoutGenerationRef.current === layoutGeneration;
|
||||
const applyWindowMode = nextMode === "setup" ? showSetupWindow : applyAppWindowLayout;
|
||||
applyWindowMode(isCurrent).catch(async () => {
|
||||
if (!isCurrent()) return;
|
||||
// On failure, at minimum make the window visible and resizable so user can fix manually.
|
||||
try {
|
||||
await showWindowFallback();
|
||||
} catch { /* swallow — window may still be functional */ }
|
||||
});
|
||||
}, [status]);
|
||||
|
||||
if (!isTauri) return <>{children}</>;
|
||||
if (status === "running") return <><TauriUpdateLayer isExternalServer={isExternalServer} />{children}</>;
|
||||
useEffect(() => {
|
||||
if (!isTauri) {
|
||||
setDesktopAuthReady(true);
|
||||
return;
|
||||
}
|
||||
if (status !== "running") {
|
||||
setDesktopAuthReady(false);
|
||||
setDesktopAuthRetry(0);
|
||||
return;
|
||||
}
|
||||
|
||||
return (
|
||||
let disposed = false;
|
||||
setDesktopAuthReady(false);
|
||||
tauriAutoAuth({ force: true }).then((authenticated) => {
|
||||
if (disposed) return;
|
||||
if (authenticated) {
|
||||
setDesktopAuthReady(true);
|
||||
return;
|
||||
}
|
||||
if (!getTauriAuthFailure()) {
|
||||
window.setTimeout(() => {
|
||||
if (!disposed) setDesktopAuthRetry((value) => value + 1);
|
||||
}, 500);
|
||||
}
|
||||
});
|
||||
|
||||
return () => { disposed = true; };
|
||||
}, [status, desktopAuthRetry]);
|
||||
|
||||
if (!isTauri) return <>{children}</>;
|
||||
|
||||
const showApp = status === "running" && desktopAuthReady;
|
||||
const startupStatus = status === "running" ? "starting" : status;
|
||||
const startupProgressDetail =
|
||||
status === "running" && !desktopAuthReady
|
||||
? "Signing in to desktop session..."
|
||||
: progressDetail;
|
||||
|
||||
const content = showApp ? (
|
||||
<>
|
||||
<TauriUpdateLayer isExternalServer={isExternalServer} />
|
||||
<NativeIntentDrain />
|
||||
{children}
|
||||
</>
|
||||
) : (
|
||||
<StartupScreen
|
||||
status={status}
|
||||
status={startupStatus}
|
||||
logs={logs}
|
||||
error={error}
|
||||
currentStepIndex={currentStepIndex}
|
||||
progressDetail={progressDetail}
|
||||
progressDetail={startupProgressDetail}
|
||||
elevationPackages={elevationPackages}
|
||||
onInstall={startInstall}
|
||||
onRetry={retry}
|
||||
onRetryInstall={retryInstall}
|
||||
onApproveElevation={approveElevation}
|
||||
onStartServer={retry}
|
||||
onCopyDiagnostics={copyDiagnostics}
|
||||
/>
|
||||
);
|
||||
|
||||
if (!shouldUseCustomWindowTitlebar()) return content;
|
||||
|
||||
const showSidebarSurface =
|
||||
showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname);
|
||||
|
||||
return (
|
||||
<div className="flex h-dvh min-h-0 flex-col overflow-hidden bg-background [--studio-titlebar-height:34px]">
|
||||
<WindowTitlebar showSidebarSurface={showSidebarSurface} />
|
||||
<div className="min-h-0 flex-1 overflow-hidden">
|
||||
{content}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
export function AppProvider({ children }: AppProviderProps) {
|
||||
|
|
|
|||
|
|
@ -81,7 +81,7 @@ function RootLayout() {
|
|||
pinned={pinned}
|
||||
setPinned={setPinned}
|
||||
togglePinned={togglePinned}
|
||||
className="!min-h-0 h-dvh overflow-hidden"
|
||||
className="!min-h-0 h-[calc(100dvh-var(--studio-titlebar-height,0px))] overflow-hidden"
|
||||
>
|
||||
<AppSidebar />
|
||||
<SidebarInset className={isChatRoute ? "overflow-hidden" : "overflow-y-auto"}>
|
||||
|
|
|
|||
|
|
@ -31,19 +31,18 @@ import {
|
|||
import { useAnimatedThemeToggle } from "@/components/ui/animated-theme-toggler";
|
||||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
Book03Icon,
|
||||
ChefHatIcon,
|
||||
ColumnInsertIcon,
|
||||
CursorInfo02Icon,
|
||||
Delete02Icon,
|
||||
Download03Icon,
|
||||
GemIcon,
|
||||
MessageSearch01Icon,
|
||||
Globe02Icon,
|
||||
Search01Icon,
|
||||
NewReleasesIcon,
|
||||
PowerIcon,
|
||||
PencilEdit02Icon,
|
||||
LayoutAlignLeftIcon,
|
||||
HelpCircleIcon,
|
||||
Settings02Icon,
|
||||
ZapIcon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
|
|
@ -527,9 +526,9 @@ export function AppSidebar() {
|
|||
className="!size-8"
|
||||
/>
|
||||
</div>
|
||||
<div className="flex flex-col gap-0.5 leading-none group-data-[collapsible=icon]:hidden">
|
||||
<div className="flex flex-col gap-0.5 leading-tight group-data-[collapsible=icon]:hidden">
|
||||
<span className="truncate font-heading text-[13px] tracking-[0.02em] font-semibold text-[#383835] dark:text-[#c7c7c4]">{displayTitle}</span>
|
||||
<span className="truncate text-[11px] tracking-[0.01em] text-muted-foreground">Studio</span>
|
||||
<span className="truncate text-[11px] tracking-[0.01em] text-muted-foreground">Unsloth</span>
|
||||
</div>
|
||||
<ChevronsUpDown strokeWidth={1.25} className="ml-auto size-4 text-muted-foreground group-data-[collapsible=icon]:hidden" />
|
||||
</SidebarMenuButton>
|
||||
|
|
@ -547,6 +546,15 @@ export function AppSidebar() {
|
|||
<span>Settings</span>
|
||||
<DropdownMenuShortcut>⌘,</DropdownMenuShortcut>
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem
|
||||
onSelect={() => useSettingsDialogStore.getState().openDialog("api-keys")}
|
||||
>
|
||||
<HugeiconsIcon icon={Globe02Icon} strokeWidth={1.75} className="size-[18px]" />
|
||||
<span>API</span>
|
||||
<span className="ml-auto rounded-[6px] border border-emerald-500/25 bg-emerald-500/10 px-1.5 py-0.5 text-[10px] leading-none font-semibold text-emerald-700 dark:text-emerald-300">
|
||||
New
|
||||
</span>
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem
|
||||
ref={anchorRef as React.Ref<HTMLDivElement>}
|
||||
onSelect={(e) => { e.preventDefault(); toggleTheme(); }}
|
||||
|
|
@ -571,47 +579,12 @@ export function AppSidebar() {
|
|||
</DropdownMenuItem>
|
||||
</DropdownMenuGroup>
|
||||
<DropdownMenuSeparator className="mx-2.5! my-2.5! h-0! border-t border-border/70 bg-transparent!" />
|
||||
<DropdownMenuGroup>
|
||||
<DropdownMenuItem asChild>
|
||||
<a
|
||||
href="https://unsloth.ai/docs"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
>
|
||||
<HugeiconsIcon icon={Book03Icon} strokeWidth={1.75} className="size-[18px]" />
|
||||
<span>Learn More</span>
|
||||
</a>
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem asChild>
|
||||
<a
|
||||
href="https://unsloth.ai/docs/new/changelog"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={NewReleasesIcon}
|
||||
strokeWidth={1.75}
|
||||
className="size-[18px]"
|
||||
/>
|
||||
<span>What's New</span>
|
||||
</a>
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem asChild>
|
||||
<a
|
||||
href="https://github.com/unslothai/unsloth/issues"
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
>
|
||||
<HugeiconsIcon
|
||||
icon={MessageSearch01Icon}
|
||||
strokeWidth={1.75}
|
||||
className="size-[18px]"
|
||||
/>
|
||||
<span>Feedback</span>
|
||||
</a>
|
||||
</DropdownMenuItem>
|
||||
</DropdownMenuGroup>
|
||||
<DropdownMenuSeparator className="mx-2.5! my-2.5! h-0! border-t border-border/70 bg-transparent!" />
|
||||
<DropdownMenuItem
|
||||
onSelect={() => useSettingsDialogStore.getState().openDialog("about")}
|
||||
>
|
||||
<HugeiconsIcon icon={HelpCircleIcon} strokeWidth={1.75} className="size-[18px]" />
|
||||
<span>Help</span>
|
||||
</DropdownMenuItem>
|
||||
<DropdownMenuItem onSelect={() => setShutdownOpen(true)}>
|
||||
<HugeiconsIcon icon={PowerIcon} strokeWidth={1.75} className="size-[18px]" />
|
||||
<span>Shutdown</span>
|
||||
|
|
|
|||
|
|
@ -13,18 +13,25 @@ import { usePlatformStore } from "@/config/env";
|
|||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
ArrowDown01Icon,
|
||||
FolderSearchIcon,
|
||||
Logout01Icon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { useMemo, useState } from "react";
|
||||
import type {
|
||||
DeletedModelRef,
|
||||
LoraModelOption,
|
||||
ModelOption,
|
||||
ModelSelectorChangeMeta,
|
||||
} from "./model-selector/types";
|
||||
import { HubModelPicker, LoraModelPicker } from "./model-selector/pickers";
|
||||
|
||||
export type { LoraModelOption, ModelOption, ModelSelectorChangeMeta } from "./model-selector/types";
|
||||
export type {
|
||||
DeletedModelRef,
|
||||
LoraModelOption,
|
||||
ModelOption,
|
||||
ModelSelectorChangeMeta,
|
||||
} from "./model-selector/types";
|
||||
|
||||
interface ModelSelectorProps {
|
||||
models: ModelOption[];
|
||||
|
|
@ -35,6 +42,9 @@ interface ModelSelectorProps {
|
|||
onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
|
||||
onEject?: () => void;
|
||||
onFoldersChange?: () => void;
|
||||
onPickLocalModel?: () => void | Promise<void>;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
variant?: "outline" | "ghost" | "muted";
|
||||
size?: "sm" | "default" | "lg";
|
||||
className?: string;
|
||||
|
|
@ -66,7 +76,7 @@ function ModelSelectorTrigger({
|
|||
type="button"
|
||||
data-tour={dataTour}
|
||||
className={cn(
|
||||
"flex items-center gap-2 transition-colors",
|
||||
"flex min-w-0 items-center gap-2 transition-colors",
|
||||
variant === "outline" &&
|
||||
"rounded-[8px] border border-border/60 hover:bg-[#ececec] dark:hover:bg-[#2e3035]",
|
||||
variant === "ghost" && "rounded-[8px] hover:bg-[#ececec] dark:hover:bg-[#2e3035]",
|
||||
|
|
@ -80,17 +90,23 @@ function ModelSelectorTrigger({
|
|||
{isLoaded && (
|
||||
<span className="size-2 shrink-0 rounded-full bg-emerald-500" />
|
||||
)}
|
||||
<span className="font-heading font-medium text-[16px] text-black dark:text-white">
|
||||
{currentModel?.name ?? "Select model"}
|
||||
<span className="flex min-w-0 flex-1 items-baseline gap-2">
|
||||
<span className="min-w-0 flex-1 truncate font-heading text-[16px] font-medium leading-tight text-black dark:text-white">
|
||||
{currentModel?.name ?? "Select model"}
|
||||
</span>
|
||||
{currentModel?.description && (
|
||||
<span className="shrink-0 text-xs leading-none text-muted-foreground">
|
||||
{currentModel.description}
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
<span className="flex size-4 shrink-0 items-center justify-center">
|
||||
<HugeiconsIcon
|
||||
icon={ArrowDown01Icon}
|
||||
strokeWidth={1.75}
|
||||
className="relative top-0.5 size-3.5 text-muted-foreground"
|
||||
/>
|
||||
</span>
|
||||
{currentModel?.description && (
|
||||
<span className="text-muted-foreground text-xs">{currentModel.description}</span>
|
||||
)}
|
||||
<HugeiconsIcon
|
||||
icon={ArrowDown01Icon}
|
||||
strokeWidth={1.75}
|
||||
className="size-3.5 shrink-0 text-muted-foreground"
|
||||
/>
|
||||
</button>
|
||||
</PopoverTrigger>
|
||||
);
|
||||
|
|
@ -103,6 +119,9 @@ function ModelSelectorContent({
|
|||
onSelect,
|
||||
onEject,
|
||||
onFoldersChange,
|
||||
onPickLocalModel,
|
||||
onModelsChange,
|
||||
deleteDisabled,
|
||||
className,
|
||||
dataTour,
|
||||
}: {
|
||||
|
|
@ -112,6 +131,9 @@ function ModelSelectorContent({
|
|||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
onEject?: () => void;
|
||||
onFoldersChange?: () => void;
|
||||
onPickLocalModel?: () => void;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
className?: string;
|
||||
dataTour?: string;
|
||||
}) {
|
||||
|
|
@ -145,11 +167,26 @@ function ModelSelectorContent({
|
|||
loraModels={loraModels}
|
||||
value={value}
|
||||
onSelect={onSelect}
|
||||
onModelsChange={onModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
/>
|
||||
</TabsContent>
|
||||
</Tabs>
|
||||
)}
|
||||
|
||||
{onPickLocalModel ? (
|
||||
<div className="mt-2 border-t border-border/70 pt-2">
|
||||
<button
|
||||
type="button"
|
||||
onClick={onPickLocalModel}
|
||||
className="flex w-full items-center justify-center gap-1.5 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-muted/60"
|
||||
title="Pick a model file from disk"
|
||||
>
|
||||
<HugeiconsIcon icon={FolderSearchIcon} className="size-3.5" />
|
||||
Pick a model file from disk
|
||||
</button>
|
||||
</div>
|
||||
) : null}
|
||||
{hasSelection && onEject ? (
|
||||
<div className="mt-2 border-t border-border/70 pt-2">
|
||||
<button
|
||||
|
|
@ -176,6 +213,9 @@ export function ModelSelector({
|
|||
onValueChange,
|
||||
onEject,
|
||||
onFoldersChange,
|
||||
onPickLocalModel,
|
||||
onModelsChange,
|
||||
deleteDisabled,
|
||||
variant = "outline",
|
||||
size = "default",
|
||||
className,
|
||||
|
|
@ -253,6 +293,11 @@ export function ModelSelector({
|
|||
setOpen(false);
|
||||
}
|
||||
|
||||
function handlePickLocalModel() {
|
||||
setOpen(false);
|
||||
void onPickLocalModel?.();
|
||||
}
|
||||
|
||||
return (
|
||||
<Popover open={open} onOpenChange={setOpen}>
|
||||
<ModelSelectorTrigger
|
||||
|
|
@ -270,6 +315,9 @@ export function ModelSelector({
|
|||
onSelect={handleSelect}
|
||||
onEject={onEject ? handleEject : undefined}
|
||||
onFoldersChange={onFoldersChange}
|
||||
onPickLocalModel={onPickLocalModel ? handlePickLocalModel : undefined}
|
||||
onModelsChange={onModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
className={contentClassName}
|
||||
dataTour={contentDataTour}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,112 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
AlertDialogCancel,
|
||||
AlertDialogContent,
|
||||
AlertDialogDescription,
|
||||
AlertDialogFooter,
|
||||
AlertDialogHeader,
|
||||
AlertDialogTitle,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { Trash2Icon } from "lucide-react";
|
||||
import { useCallback, useState, type ReactNode } from "react";
|
||||
import { toast } from "sonner";
|
||||
|
||||
interface ModelDeleteActionProps {
|
||||
ariaLabel: string;
|
||||
title: string;
|
||||
description: ReactNode;
|
||||
successMessage: string;
|
||||
loadingLabel?: string;
|
||||
buttonClassName?: string;
|
||||
iconClassName?: string;
|
||||
disabled?: boolean;
|
||||
onConfirm: () => Promise<void> | void;
|
||||
onDeleted?: () => void;
|
||||
}
|
||||
|
||||
export function ModelDeleteAction({
|
||||
ariaLabel,
|
||||
title,
|
||||
description,
|
||||
successMessage,
|
||||
loadingLabel = "Deleting...",
|
||||
buttonClassName,
|
||||
iconClassName,
|
||||
disabled = false,
|
||||
onConfirm,
|
||||
onDeleted,
|
||||
}: ModelDeleteActionProps) {
|
||||
const [open, setOpen] = useState(false);
|
||||
const [deleting, setDeleting] = useState(false);
|
||||
|
||||
const handleConfirm = useCallback(async () => {
|
||||
setDeleting(true);
|
||||
try {
|
||||
await onConfirm();
|
||||
toast.success(successMessage);
|
||||
onDeleted?.();
|
||||
setOpen(false);
|
||||
} catch (err) {
|
||||
toast.error(
|
||||
err instanceof Error ? err.message : "Failed to delete model",
|
||||
);
|
||||
} finally {
|
||||
setDeleting(false);
|
||||
}
|
||||
}, [onConfirm, onDeleted, successMessage]);
|
||||
|
||||
return (
|
||||
<>
|
||||
<button
|
||||
type="button"
|
||||
onClick={(e) => {
|
||||
e.stopPropagation();
|
||||
if (disabled) return;
|
||||
setOpen(true);
|
||||
}}
|
||||
aria-label={ariaLabel}
|
||||
disabled={disabled}
|
||||
className={cn(
|
||||
"shrink-0 rounded-md p-1.5 text-muted-foreground/60 transition-colors hover:bg-destructive/10 hover:text-destructive",
|
||||
disabled && "cursor-not-allowed opacity-40 hover:bg-transparent hover:text-muted-foreground/60",
|
||||
buttonClassName,
|
||||
)}
|
||||
>
|
||||
<Trash2Icon className={cn("size-3.5", iconClassName)} />
|
||||
</button>
|
||||
|
||||
<AlertDialog
|
||||
open={open}
|
||||
onOpenChange={(nextOpen) => {
|
||||
if (!nextOpen && deleting) return;
|
||||
setOpen(nextOpen);
|
||||
}}
|
||||
>
|
||||
<AlertDialogContent size="sm">
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>{title}</AlertDialogTitle>
|
||||
<AlertDialogDescription>{description}</AlertDialogDescription>
|
||||
</AlertDialogHeader>
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel disabled={deleting}>No</AlertDialogCancel>
|
||||
<AlertDialogAction
|
||||
variant="destructive"
|
||||
disabled={deleting}
|
||||
onClick={(e) => {
|
||||
e.preventDefault();
|
||||
handleConfirm();
|
||||
}}
|
||||
>
|
||||
{deleting ? loadingLabel : "Yes"}
|
||||
</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
|
@ -1,16 +1,6 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import {
|
||||
AlertDialog,
|
||||
AlertDialogAction,
|
||||
AlertDialogCancel,
|
||||
AlertDialogContent,
|
||||
AlertDialogDescription,
|
||||
AlertDialogFooter,
|
||||
AlertDialogHeader,
|
||||
AlertDialogTitle,
|
||||
} from "@/components/ui/alert-dialog";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import {
|
||||
|
|
@ -23,6 +13,7 @@ import {
|
|||
type ScanFolderInfo,
|
||||
addScanFolder,
|
||||
deleteCachedModel,
|
||||
deleteFineTunedModel,
|
||||
listCachedGguf,
|
||||
listCachedModels,
|
||||
listGgufVariants,
|
||||
|
|
@ -50,7 +41,8 @@ import { checkVramFit, estimateLoadingVram } from "@/lib/vram";
|
|||
import { Add01Icon, Cancel01Icon, Folder02Icon, Search01Icon } from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { FolderBrowser } from "./folder-browser";
|
||||
import { ChevronDownIcon, ChevronRightIcon, DownloadIcon, StarIcon, Trash2Icon } from "lucide-react";
|
||||
import { ModelDeleteAction } from "./model-delete-action";
|
||||
import { ChevronDownIcon, ChevronRightIcon, DownloadIcon, StarIcon } from "lucide-react";
|
||||
import {
|
||||
type ReactNode,
|
||||
useCallback,
|
||||
|
|
@ -60,6 +52,7 @@ import {
|
|||
} from "react";
|
||||
import { toast } from "sonner";
|
||||
import type {
|
||||
DeletedModelRef,
|
||||
LoraModelOption,
|
||||
ModelOption,
|
||||
ModelSelectorChangeMeta,
|
||||
|
|
@ -211,12 +204,22 @@ function GgufVariantExpander({
|
|||
gpuGb,
|
||||
systemRamGb,
|
||||
onDeleteVariant,
|
||||
sourceOverride,
|
||||
deleteVariantTitle = "Delete cached model?",
|
||||
renderDeleteVariantDescription,
|
||||
getDeleteVariantSuccessMessage,
|
||||
deleteDisabled = false,
|
||||
}: {
|
||||
repoId: string;
|
||||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
gpuGb?: number;
|
||||
systemRamGb?: number;
|
||||
onDeleteVariant?: (quant: string) => void;
|
||||
onDeleteVariant?: (quant: string) => Promise<void> | void;
|
||||
sourceOverride?: ModelSelectorChangeMeta["source"];
|
||||
deleteVariantTitle?: string;
|
||||
renderDeleteVariantDescription?: (quant: string) => ReactNode;
|
||||
getDeleteVariantSuccessMessage?: (quant: string) => string;
|
||||
deleteDisabled?: boolean;
|
||||
}) {
|
||||
const [variants, setVariants] = useState<GgufVariantDetail[] | null>(null);
|
||||
const [defaultVariant, setDefaultVariant] = useState<string | null>(null);
|
||||
|
|
@ -259,14 +262,14 @@ function GgufVariantExpander({
|
|||
const handleVariantClick = useCallback(
|
||||
(quant: string, downloaded?: boolean, sizeBytes?: number) => {
|
||||
onSelect(repoId, {
|
||||
source: isLocalPath ? "local" : "hub",
|
||||
source: sourceOverride ?? (isLocalPath ? "local" : "hub"),
|
||||
isLora: false,
|
||||
ggufVariant: quant,
|
||||
isDownloaded: isLocalPath ? true : downloaded,
|
||||
expectedBytes: sizeBytes,
|
||||
});
|
||||
},
|
||||
[repoId, isLocalPath, onSelect],
|
||||
[repoId, isLocalPath, onSelect, sourceOverride],
|
||||
);
|
||||
|
||||
// GGUF fit classification matching llama-server's _select_gpus logic:
|
||||
|
|
@ -408,16 +411,29 @@ function GgufVariantExpander({
|
|||
</span>
|
||||
</button>
|
||||
{v.downloaded && onDeleteVariant && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={(e) => {
|
||||
e.stopPropagation();
|
||||
onDeleteVariant(v.quant);
|
||||
}}
|
||||
className="shrink-0 rounded-md p-1 text-muted-foreground/60 transition-colors hover:bg-destructive/10 hover:text-destructive"
|
||||
>
|
||||
<Trash2Icon className="size-3" />
|
||||
</button>
|
||||
<ModelDeleteAction
|
||||
ariaLabel={`Delete ${repoId} ${v.quant}`}
|
||||
title={deleteVariantTitle}
|
||||
description={
|
||||
renderDeleteVariantDescription?.(v.quant) ?? (
|
||||
<>
|
||||
This will remove{" "}
|
||||
<span className="font-medium text-foreground">
|
||||
{repoId} ({v.quant})
|
||||
</span>{" "}
|
||||
from disk. You can re-download it later.
|
||||
</>
|
||||
)
|
||||
}
|
||||
successMessage={
|
||||
getDeleteVariantSuccessMessage?.(v.quant) ??
|
||||
`Deleted ${repoId} ${v.quant}`
|
||||
}
|
||||
buttonClassName="p-1"
|
||||
iconClassName="size-3"
|
||||
disabled={deleteDisabled}
|
||||
onConfirm={() => onDeleteVariant(v.quant)}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
|
|
@ -512,9 +528,6 @@ export function HubModelPicker({
|
|||
// Track which GGUF repo is expanded for variant selection
|
||||
const [expandedGguf, setExpandedGguf] = useState<string | null>(null);
|
||||
|
||||
// Delete confirmation dialog state
|
||||
const [deleteTarget, setDeleteTarget] = useState<string | null>(null);
|
||||
const [deleting, setDeleting] = useState(false);
|
||||
const [downloadedCollapsed, setDownloadedCollapsed] = useState(false);
|
||||
const [customFoldersCollapsed, setCustomFoldersCollapsed] = useState(false);
|
||||
const [recommendedCollapsed, setRecommendedCollapsed] = useState(false);
|
||||
|
|
@ -675,27 +688,6 @@ export function HubModelPicker({
|
|||
.finally(check);
|
||||
}, [refreshLocalModelsList, refreshScanFolders]);
|
||||
|
||||
const handleDeleteConfirm = useCallback(async () => {
|
||||
if (!deleteTarget) return;
|
||||
setDeleting(true);
|
||||
try {
|
||||
// deleteTarget is "repo_id" or "repo_id::variant"
|
||||
const sepIdx = deleteTarget.indexOf("::");
|
||||
const repoId = sepIdx >= 0 ? deleteTarget.slice(0, sepIdx) : deleteTarget;
|
||||
const variant = sepIdx >= 0 ? deleteTarget.slice(sepIdx + 2) : undefined;
|
||||
await deleteCachedModel(repoId, variant);
|
||||
toast.success(`Deleted ${variant ? `${repoId} ${variant}` : repoId}`);
|
||||
refreshCachedLists();
|
||||
} catch (err) {
|
||||
toast.error(
|
||||
err instanceof Error ? err.message : "Failed to delete model",
|
||||
);
|
||||
} finally {
|
||||
setDeleting(false);
|
||||
setDeleteTarget(null);
|
||||
}
|
||||
}, [deleteTarget, refreshCachedLists]);
|
||||
|
||||
// Deduplicate: don't show downloaded models in the recommended list.
|
||||
// Compare case-insensitively since HF cache lowercases repo IDs.
|
||||
const downloadedSet = useMemo(() => {
|
||||
|
|
@ -952,9 +944,10 @@ export function HubModelPicker({
|
|||
systemRamGb={
|
||||
gpu.available ? gpu.systemRamAvailableGb : undefined
|
||||
}
|
||||
onDeleteVariant={(quant) =>
|
||||
setDeleteTarget(`${c.repo_id}::${quant}`)
|
||||
}
|
||||
onDeleteVariant={async (quant) => {
|
||||
await deleteCachedModel(c.repo_id, quant);
|
||||
refreshCachedLists();
|
||||
}}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
|
@ -977,16 +970,22 @@ export function HubModelPicker({
|
|||
vramStatus={null}
|
||||
/>
|
||||
</div>
|
||||
<button
|
||||
type="button"
|
||||
onClick={(e) => {
|
||||
e.stopPropagation();
|
||||
setDeleteTarget(c.repo_id);
|
||||
}}
|
||||
className="shrink-0 rounded-md p-1.5 text-muted-foreground/60 transition-colors hover:bg-destructive/10 hover:text-destructive"
|
||||
>
|
||||
<Trash2Icon className="size-3.5" />
|
||||
</button>
|
||||
<ModelDeleteAction
|
||||
ariaLabel={`Delete ${c.repo_id}`}
|
||||
title="Delete cached model?"
|
||||
description={
|
||||
<>
|
||||
This will remove{" "}
|
||||
<span className="font-medium text-foreground">
|
||||
{c.repo_id}
|
||||
</span>{" "}
|
||||
from disk. You can re-download it later.
|
||||
</>
|
||||
}
|
||||
successMessage={`Deleted ${c.repo_id}`}
|
||||
onConfirm={() => deleteCachedModel(c.repo_id)}
|
||||
onDeleted={refreshCachedLists}
|
||||
/>
|
||||
</div>
|
||||
))}
|
||||
</>
|
||||
|
|
@ -1409,40 +1408,6 @@ export function HubModelPicker({
|
|||
</div>
|
||||
</div>
|
||||
|
||||
<AlertDialog
|
||||
open={deleteTarget !== null}
|
||||
onOpenChange={(open) => {
|
||||
if (!open && !deleting) setDeleteTarget(null);
|
||||
}}
|
||||
>
|
||||
<AlertDialogContent size="sm">
|
||||
<AlertDialogHeader>
|
||||
<AlertDialogTitle>Delete cached model?</AlertDialogTitle>
|
||||
<AlertDialogDescription>
|
||||
This will remove{" "}
|
||||
<span className="font-medium text-foreground">
|
||||
{deleteTarget?.includes("::")
|
||||
? `${deleteTarget.split("::")[0]} (${deleteTarget.split("::")[1]})`
|
||||
: deleteTarget}
|
||||
</span>{" "}
|
||||
from disk. You can re-download it later.
|
||||
</AlertDialogDescription>
|
||||
</AlertDialogHeader>
|
||||
<AlertDialogFooter>
|
||||
<AlertDialogCancel disabled={deleting}>No</AlertDialogCancel>
|
||||
<AlertDialogAction
|
||||
variant="destructive"
|
||||
disabled={deleting}
|
||||
onClick={(e) => {
|
||||
e.preventDefault();
|
||||
handleDeleteConfirm();
|
||||
}}
|
||||
>
|
||||
{deleting ? "Deleting..." : "Yes"}
|
||||
</AlertDialogAction>
|
||||
</AlertDialogFooter>
|
||||
</AlertDialogContent>
|
||||
</AlertDialog>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -1451,10 +1416,14 @@ export function LoraModelPicker({
|
|||
loraModels,
|
||||
value,
|
||||
onSelect,
|
||||
onModelsChange,
|
||||
deleteDisabled = false,
|
||||
}: {
|
||||
loraModels: LoraModelOption[];
|
||||
value?: string;
|
||||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
}) {
|
||||
const [query, setQuery] = useState("");
|
||||
const [expandedGguf, setExpandedGguf] = useState<string | null>(null);
|
||||
|
|
@ -1541,6 +1510,8 @@ export function LoraModelPicker({
|
|||
const isExported = adapter.source === "exported";
|
||||
const isMerged = adapter.exportType === "merged";
|
||||
const isGguf = adapter.exportType === "gguf";
|
||||
const isExportedGguf = isExported && isGguf;
|
||||
const canDelete = (isTraining || isExported) && !isExportedGguf;
|
||||
const isTrainingFull = isTraining && isMerged;
|
||||
const isLocalGgufDir =
|
||||
isLocal &&
|
||||
|
|
@ -1569,38 +1540,69 @@ export function LoraModelPicker({
|
|||
: tag;
|
||||
return (
|
||||
<div key={adapter.id}>
|
||||
<ModelRow
|
||||
label={adapter.name}
|
||||
meta={meta}
|
||||
selected={value === adapter.id}
|
||||
onClick={() => {
|
||||
if (isLocalGgufDir) {
|
||||
setExpandedGguf((prev) =>
|
||||
prev === adapter.id ? null : adapter.id,
|
||||
);
|
||||
} else {
|
||||
onSelect(adapter.id, {
|
||||
source: isLocal
|
||||
? "local"
|
||||
: isExported
|
||||
? "exported"
|
||||
: "lora",
|
||||
isLora: !isLocal && !isMerged && !isGguf,
|
||||
isDownloaded: true,
|
||||
});
|
||||
}
|
||||
}}
|
||||
tooltipText={
|
||||
<>
|
||||
<span className="block break-words">
|
||||
{adapter.name}
|
||||
</span>
|
||||
<span className="block mt-1 text-[10px] text-muted-foreground break-all">
|
||||
{adapter.id}
|
||||
</span>
|
||||
</>
|
||||
}
|
||||
/>
|
||||
<div className="flex items-center gap-0.5">
|
||||
<div className="min-w-0 flex-1">
|
||||
<ModelRow
|
||||
label={adapter.name}
|
||||
meta={meta}
|
||||
selected={value === adapter.id}
|
||||
onClick={() => {
|
||||
if (isLocalGgufDir || isExportedGguf) {
|
||||
setExpandedGguf((prev) =>
|
||||
prev === adapter.id ? null : adapter.id,
|
||||
);
|
||||
} else {
|
||||
onSelect(adapter.id, {
|
||||
source: isLocal
|
||||
? "local"
|
||||
: isExported
|
||||
? "exported"
|
||||
: "lora",
|
||||
isLora: !isLocal && !isMerged && !isGguf,
|
||||
isDownloaded: true,
|
||||
});
|
||||
}
|
||||
}}
|
||||
tooltipText={
|
||||
<>
|
||||
<span className="block break-words">
|
||||
{adapter.name}
|
||||
</span>
|
||||
<span className="block mt-1 text-[10px] text-muted-foreground break-all">
|
||||
{adapter.id}
|
||||
</span>
|
||||
</>
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
{canDelete && (
|
||||
<ModelDeleteAction
|
||||
ariaLabel={`Delete ${adapter.name}`}
|
||||
title="Delete fine-tuned model?"
|
||||
description={
|
||||
<>
|
||||
This will remove{" "}
|
||||
<span className="font-medium text-foreground">
|
||||
{adapter.name}
|
||||
</span>{" "}
|
||||
from disk. This cannot be undone.
|
||||
</>
|
||||
}
|
||||
successMessage={`Deleted ${adapter.name}`}
|
||||
disabled={deleteDisabled}
|
||||
onConfirm={() =>
|
||||
deleteFineTunedModel({
|
||||
modelPath: adapter.id,
|
||||
source: isExported ? "exported" : "training",
|
||||
exportType: adapter.exportType,
|
||||
})
|
||||
}
|
||||
onDeleted={() =>
|
||||
onModelsChange?.({ id: adapter.id })
|
||||
}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
{expandedGguf === adapter.id && (
|
||||
<GgufVariantExpander
|
||||
repoId={adapter.id}
|
||||
|
|
@ -1609,6 +1611,37 @@ export function LoraModelPicker({
|
|||
systemRamGb={
|
||||
gpu.available ? gpu.systemRamAvailableGb : undefined
|
||||
}
|
||||
sourceOverride={isExportedGguf ? "exported" : undefined}
|
||||
deleteVariantTitle="Delete exported GGUF variant?"
|
||||
renderDeleteVariantDescription={(quant) => (
|
||||
<>
|
||||
This will remove{" "}
|
||||
<span className="font-medium text-foreground">
|
||||
{adapter.name} ({quant})
|
||||
</span>{" "}
|
||||
from disk. This cannot be undone.
|
||||
</>
|
||||
)}
|
||||
getDeleteVariantSuccessMessage={(quant) =>
|
||||
`Deleted ${adapter.name} ${quant}`
|
||||
}
|
||||
deleteDisabled={deleteDisabled}
|
||||
onDeleteVariant={
|
||||
isExportedGguf
|
||||
? async (quant) => {
|
||||
await deleteFineTunedModel({
|
||||
modelPath: adapter.id,
|
||||
source: "exported",
|
||||
exportType: "gguf",
|
||||
ggufVariant: quant,
|
||||
});
|
||||
onModelsChange?.({
|
||||
id: adapter.id,
|
||||
ggufVariant: quant,
|
||||
});
|
||||
}
|
||||
: undefined
|
||||
}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
|
@ -1619,6 +1652,7 @@ export function LoraModelPicker({
|
|||
)}
|
||||
</div>
|
||||
</div>
|
||||
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -25,3 +25,8 @@ export interface ModelSelectorChangeMeta {
|
|||
isDownloaded?: boolean;
|
||||
expectedBytes?: number;
|
||||
}
|
||||
|
||||
export interface DeletedModelRef {
|
||||
id: string;
|
||||
ggufVariant?: string;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ import {
|
|||
import { sentAudioNames } from "@/features/chat/api/chat-adapter";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { deleteThreadMessage } from "@/features/chat/utils/delete-thread-message";
|
||||
import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
|
|
@ -298,27 +299,41 @@ const Composer: FC<{ disabled?: boolean }> = ({ disabled }) => {
|
|||
[disabled],
|
||||
);
|
||||
|
||||
const composerContent = (
|
||||
<>
|
||||
<ComposerAttachments />
|
||||
<PendingAudioChip />
|
||||
<ToolStatusDisplay />
|
||||
<ComposerPrimitive.Input
|
||||
placeholder="Send a message..."
|
||||
className="aui-composer-input mb-1 min-h-12 w-full resize-none overflow-y-auto bg-transparent pl-5 pr-4 pt-2 pb-3 text-sm font-[450] outline-none placeholder:text-muted-foreground focus-visible:ring-0"
|
||||
minRows={1}
|
||||
maxRows={6}
|
||||
autoFocus={!disabled}
|
||||
disabled={disabled}
|
||||
aria-label="Message input"
|
||||
/>
|
||||
<ComposerAction disabled={disabled} />
|
||||
</>
|
||||
);
|
||||
|
||||
return (
|
||||
<ComposerPrimitive.Root
|
||||
className="aui-composer-root relative flex w-full flex-col"
|
||||
aria-disabled={disabled}
|
||||
onSubmit={handleSubmit}
|
||||
>
|
||||
<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone chat-composer-surface flex w-full flex-col rounded-3xl bg-background dark:bg-card px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:bg-accent/50">
|
||||
<ComposerAttachments />
|
||||
<PendingAudioChip />
|
||||
<ToolStatusDisplay />
|
||||
<ComposerPrimitive.Input
|
||||
placeholder="Send a message..."
|
||||
className="aui-composer-input mb-1 min-h-12 w-full resize-none overflow-y-auto bg-transparent pl-5 pr-4 pt-2 pb-3 text-sm font-[450] outline-none placeholder:text-muted-foreground focus-visible:ring-0"
|
||||
minRows={1}
|
||||
maxRows={6}
|
||||
autoFocus={!disabled}
|
||||
disabled={disabled}
|
||||
aria-label="Message input"
|
||||
/>
|
||||
<ComposerAction disabled={disabled} />
|
||||
</ComposerPrimitive.AttachmentDropzone>
|
||||
{isTauri ? (
|
||||
// Phase 1 native model drops own Tauri local-path drops. Restore browser
|
||||
// attachment drops in Tauri when Phase 1d adds attachment-token bridging.
|
||||
<div className="aui-composer-attachment-dropzone chat-composer-surface flex w-full flex-col rounded-3xl bg-background dark:bg-card px-1 pt-2 outline-none transition-shadow">
|
||||
{composerContent}
|
||||
</div>
|
||||
) : (
|
||||
<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone chat-composer-surface flex w-full flex-col rounded-3xl bg-background dark:bg-card px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:bg-accent/50">
|
||||
{composerContent}
|
||||
</ComposerPrimitive.AttachmentDropzone>
|
||||
)}
|
||||
</ComposerPrimitive.Root>
|
||||
);
|
||||
};
|
||||
|
|
@ -485,7 +500,7 @@ const PreserveThinkingToggle: FC = () => {
|
|||
: "bg-muted text-muted-foreground hover:bg-muted-foreground/15",
|
||||
)}
|
||||
aria-label={
|
||||
preserveThinking ? "Disable preserve thinking" : "Enable preserve thinking"
|
||||
preserveThinking ? "Disable preserve think" : "Enable preserve think"
|
||||
}
|
||||
>
|
||||
{preserveThinking && !disabled ? (
|
||||
|
|
@ -493,7 +508,7 @@ const PreserveThinkingToggle: FC = () => {
|
|||
) : (
|
||||
<LightbulbOffIcon className="size-3.5" />
|
||||
)}
|
||||
<span>Preserve Thinking</span>
|
||||
<span>Preserve Think</span>
|
||||
</button>
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@
|
|||
|
||||
import { ShimmerButton } from "@/components/ui/shimmer-button";
|
||||
import type { BackendStatus } from "@/hooks/use-tauri-backend";
|
||||
import type { CopySupportDiagnosticsResult } from "@/lib/tauri-diagnostics";
|
||||
import { AnimatePresence, motion } from "motion/react";
|
||||
import { useState } from "react";
|
||||
|
||||
interface StartupScreenProps {
|
||||
status: BackendStatus;
|
||||
|
|
@ -17,6 +19,63 @@ interface StartupScreenProps {
|
|||
onRetryInstall: () => void;
|
||||
onApproveElevation: () => void;
|
||||
onStartServer: () => void;
|
||||
onCopyDiagnostics: () => Promise<CopySupportDiagnosticsResult>;
|
||||
}
|
||||
|
||||
function DiagnosticsCopyActions({
|
||||
onCopyDiagnostics,
|
||||
children,
|
||||
}: {
|
||||
onCopyDiagnostics: () => Promise<CopySupportDiagnosticsResult>;
|
||||
children: React.ReactNode;
|
||||
}) {
|
||||
const [copying, setCopying] = useState(false);
|
||||
const [manualReport, setManualReport] = useState<string | null>(null);
|
||||
const [manualMessage, setManualMessage] = useState<string | null>(null);
|
||||
|
||||
async function handleCopyDiagnostics() {
|
||||
setCopying(true);
|
||||
try {
|
||||
const result = await onCopyDiagnostics();
|
||||
if (result.ok) {
|
||||
setManualReport(null);
|
||||
setManualMessage(null);
|
||||
} else {
|
||||
setManualReport(result.report);
|
||||
setManualMessage(result.error ?? "Clipboard copy failed. Select and copy the diagnostics below.");
|
||||
}
|
||||
} catch (error) {
|
||||
setManualReport(null);
|
||||
setManualMessage(`Diagnostics copy failed: ${String(error)}`);
|
||||
} finally {
|
||||
setCopying(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="mt-4 flex w-full flex-col items-center gap-3">
|
||||
<div className="flex gap-3">
|
||||
<ActionButton
|
||||
variant="secondary"
|
||||
onClick={() => void handleCopyDiagnostics()}
|
||||
>
|
||||
{copying ? "Copying..." : "Copy Diagnostics"}
|
||||
</ActionButton>
|
||||
{children}
|
||||
</div>
|
||||
{manualMessage && (
|
||||
<p className="max-w-md text-center text-xs text-destructive">{manualMessage}</p>
|
||||
)}
|
||||
{manualReport && (
|
||||
<textarea
|
||||
readOnly
|
||||
value={manualReport}
|
||||
onFocus={(event) => event.currentTarget.select()}
|
||||
className="h-32 w-full max-w-md resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[10px] text-muted-foreground"
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
|
@ -180,12 +239,12 @@ function RepairingContent({
|
|||
|
||||
function InstallErrorContent({
|
||||
error,
|
||||
logs,
|
||||
onRetryInstall,
|
||||
onCopyDiagnostics,
|
||||
}: {
|
||||
error: string | null;
|
||||
logs: string[];
|
||||
onRetryInstall: () => void;
|
||||
onCopyDiagnostics: () => Promise<CopySupportDiagnosticsResult>;
|
||||
}) {
|
||||
return (
|
||||
<>
|
||||
|
|
@ -195,15 +254,9 @@ function InstallErrorContent({
|
|||
{error && (
|
||||
<p className="max-w-xs text-center text-xs text-muted-foreground">{error}</p>
|
||||
)}
|
||||
<div className="mt-4 flex gap-3">
|
||||
<ActionButton
|
||||
variant="secondary"
|
||||
onClick={() => void navigator.clipboard.writeText(logs.join("\n"))}
|
||||
>
|
||||
Copy Logs
|
||||
</ActionButton>
|
||||
<DiagnosticsCopyActions onCopyDiagnostics={onCopyDiagnostics}>
|
||||
<ActionButton onClick={onRetryInstall}>Try Again</ActionButton>
|
||||
</div>
|
||||
</DiagnosticsCopyActions>
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
|
|
@ -211,12 +264,12 @@ function InstallErrorContent({
|
|||
|
||||
function RepairErrorContent({
|
||||
error,
|
||||
logs,
|
||||
onRetry,
|
||||
onCopyDiagnostics,
|
||||
}: {
|
||||
error: string | null;
|
||||
logs: string[];
|
||||
onRetry: () => void;
|
||||
onCopyDiagnostics: () => Promise<CopySupportDiagnosticsResult>;
|
||||
}) {
|
||||
return (
|
||||
<>
|
||||
|
|
@ -226,15 +279,9 @@ function RepairErrorContent({
|
|||
{error && (
|
||||
<p className="max-w-md text-center text-xs text-muted-foreground">{error}</p>
|
||||
)}
|
||||
<div className="mt-4 flex gap-3">
|
||||
<ActionButton
|
||||
variant="secondary"
|
||||
onClick={() => void navigator.clipboard.writeText(logs.join("\n"))}
|
||||
>
|
||||
Copy Logs
|
||||
</ActionButton>
|
||||
<DiagnosticsCopyActions onCopyDiagnostics={onCopyDiagnostics}>
|
||||
<ActionButton onClick={onRetry}>Retry</ActionButton>
|
||||
</div>
|
||||
</DiagnosticsCopyActions>
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
|
|
@ -301,12 +348,12 @@ function StoppedContent({ onStartServer }: { onStartServer: () => void }) {
|
|||
|
||||
function ErrorContent({
|
||||
error,
|
||||
logs,
|
||||
onRetry,
|
||||
onCopyDiagnostics,
|
||||
}: {
|
||||
error: string | null;
|
||||
logs: string[];
|
||||
onRetry: () => void;
|
||||
onCopyDiagnostics: () => Promise<CopySupportDiagnosticsResult>;
|
||||
}) {
|
||||
return (
|
||||
<>
|
||||
|
|
@ -316,15 +363,9 @@ function ErrorContent({
|
|||
{error && (
|
||||
<p className="max-w-md text-center text-xs text-muted-foreground">{error}</p>
|
||||
)}
|
||||
<div className="mt-4 flex gap-3">
|
||||
<ActionButton
|
||||
variant="secondary"
|
||||
onClick={() => void navigator.clipboard.writeText(logs.join("\n"))}
|
||||
>
|
||||
Copy Logs
|
||||
</ActionButton>
|
||||
<DiagnosticsCopyActions onCopyDiagnostics={onCopyDiagnostics}>
|
||||
<ActionButton onClick={onRetry}>Retry</ActionButton>
|
||||
</div>
|
||||
</DiagnosticsCopyActions>
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
|
|
@ -346,6 +387,7 @@ export function StartupScreen({
|
|||
onRetryInstall,
|
||||
onApproveElevation,
|
||||
onStartServer,
|
||||
onCopyDiagnostics,
|
||||
}: StartupScreenProps) {
|
||||
function renderContent() {
|
||||
switch (status) {
|
||||
|
|
@ -356,11 +398,23 @@ export function StartupScreen({
|
|||
case "installing":
|
||||
return <InstallingContent currentStepIndex={currentStepIndex} progressDetail={progressDetail} />;
|
||||
case "install-error":
|
||||
return <InstallErrorContent error={error} logs={logs} onRetryInstall={onRetryInstall} />;
|
||||
return (
|
||||
<InstallErrorContent
|
||||
error={error}
|
||||
onRetryInstall={onRetryInstall}
|
||||
onCopyDiagnostics={onCopyDiagnostics}
|
||||
/>
|
||||
);
|
||||
case "repairing":
|
||||
return <RepairingContent logs={logs} progressDetail={progressDetail} />;
|
||||
case "repair-error":
|
||||
return <RepairErrorContent error={error} logs={logs} onRetry={onRetry} />;
|
||||
return (
|
||||
<RepairErrorContent
|
||||
error={error}
|
||||
onRetry={onRetry}
|
||||
onCopyDiagnostics={onCopyDiagnostics}
|
||||
/>
|
||||
);
|
||||
case "needs-elevation":
|
||||
return (
|
||||
<NeedsElevationContent
|
||||
|
|
@ -376,12 +430,18 @@ export function StartupScreen({
|
|||
case "stopped":
|
||||
return <StoppedContent onStartServer={onStartServer} />;
|
||||
case "error":
|
||||
return <ErrorContent error={error} logs={logs} onRetry={onRetry} />;
|
||||
return (
|
||||
<ErrorContent
|
||||
error={error}
|
||||
onRetry={onRetry}
|
||||
onCopyDiagnostics={onCopyDiagnostics}
|
||||
/>
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="flex h-screen w-full flex-col items-center bg-background">
|
||||
<div className="flex h-full w-full flex-col items-center bg-background">
|
||||
<div className="flex flex-1 w-full max-w-md items-center justify-center px-6">
|
||||
<AnimatePresence mode="wait">
|
||||
<motion.div
|
||||
|
|
|
|||
|
|
@ -2,16 +2,20 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import type { UpdateInfo, UpdateStatus } from "@/hooks/use-tauri-update";
|
||||
import type { RetainedUpdateFailure, UpdateInfo, UpdateStatus } from "@/hooks/use-tauri-update";
|
||||
import type { CopySupportDiagnosticsResult } from "@/lib/tauri-diagnostics";
|
||||
import { AnimatePresence, motion } from "motion/react";
|
||||
import { useState } from "react";
|
||||
|
||||
interface UpdateBannerProps {
|
||||
status: UpdateStatus;
|
||||
info: UpdateInfo | null;
|
||||
dismissed: boolean;
|
||||
lastFailure: RetainedUpdateFailure | null;
|
||||
isExternalServer?: boolean;
|
||||
onInstall: () => void;
|
||||
onDismiss: () => void;
|
||||
onCopyDiagnostics: () => Promise<CopySupportDiagnosticsResult>;
|
||||
}
|
||||
|
||||
const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1];
|
||||
|
|
@ -20,16 +24,41 @@ export function UpdateBanner({
|
|||
status,
|
||||
info,
|
||||
dismissed,
|
||||
lastFailure,
|
||||
isExternalServer = false,
|
||||
onInstall,
|
||||
onDismiss,
|
||||
onCopyDiagnostics,
|
||||
}: UpdateBannerProps) {
|
||||
const visible = status === "available";
|
||||
const show = visible && !dismissed;
|
||||
const [copying, setCopying] = useState(false);
|
||||
const [manualReport, setManualReport] = useState<string | null>(null);
|
||||
const [manualMessage, setManualMessage] = useState<string | null>(null);
|
||||
const showFailure = Boolean(lastFailure) && !dismissed;
|
||||
const showAvailable = status === "available" && !dismissed && !showFailure;
|
||||
const show = showFailure || (showAvailable && Boolean(info));
|
||||
|
||||
async function handleCopyDiagnostics() {
|
||||
setCopying(true);
|
||||
try {
|
||||
const result = await onCopyDiagnostics();
|
||||
if (result.ok) {
|
||||
setManualReport(null);
|
||||
setManualMessage(null);
|
||||
} else {
|
||||
setManualReport(result.report);
|
||||
setManualMessage(result.error ?? "Clipboard copy failed. Select and copy the diagnostics below.");
|
||||
}
|
||||
} catch (error) {
|
||||
setManualReport(null);
|
||||
setManualMessage(`Diagnostics copy failed: ${String(error)}`);
|
||||
} finally {
|
||||
setCopying(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<AnimatePresence>
|
||||
{show && info && (
|
||||
{show && (
|
||||
<motion.div
|
||||
initial={{ opacity: 0, y: -12, scale: 0.96 }}
|
||||
animate={{ opacity: 1, y: 0, scale: 1 }}
|
||||
|
|
@ -54,28 +83,61 @@ export function UpdateBanner({
|
|||
<span className="text-lg">🦥</span>
|
||||
<div>
|
||||
<p className="text-sm font-semibold text-foreground">
|
||||
New version: v{info.version}
|
||||
{showFailure ? "App update failed" : `New version: v${info?.version}`}
|
||||
</p>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{isExternalServer
|
||||
? "Run `unsloth studio update` from your terminal"
|
||||
: "A new app update is available"}
|
||||
{showFailure
|
||||
? "Backend recovered. Diagnostics are still available."
|
||||
: isExternalServer
|
||||
? "Run `unsloth studio update` from your terminal"
|
||||
: "A new app update is available"}
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Retained failure */}
|
||||
{showFailure && lastFailure && (
|
||||
<p className="mt-3 line-clamp-2 text-xs text-destructive">
|
||||
{lastFailure.error}
|
||||
</p>
|
||||
)}
|
||||
|
||||
{/* Actions */}
|
||||
<div className="mt-3 flex items-center gap-2">
|
||||
<Button size="sm" className="corner-squircle" onClick={onInstall} disabled={isExternalServer}>
|
||||
Update Now
|
||||
</Button>
|
||||
<Button size="sm" variant="outline" className="corner-squircle" disabled>
|
||||
Release Notes
|
||||
</Button>
|
||||
{showFailure ? (
|
||||
<>
|
||||
<Button size="sm" variant="outline" className="corner-squircle" onClick={() => void handleCopyDiagnostics()}>
|
||||
{copying ? "Copying..." : "Copy Diagnostics"}
|
||||
</Button>
|
||||
<Button size="sm" className="corner-squircle" onClick={onInstall} disabled={isExternalServer}>
|
||||
Retry Update
|
||||
</Button>
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
<Button size="sm" className="corner-squircle" onClick={onInstall} disabled={isExternalServer}>
|
||||
Update Now
|
||||
</Button>
|
||||
<Button size="sm" variant="outline" className="corner-squircle" disabled>
|
||||
Release Notes
|
||||
</Button>
|
||||
</>
|
||||
)}
|
||||
<Button size="sm" variant="ghost" className="corner-squircle" onClick={onDismiss}>
|
||||
Later
|
||||
</Button>
|
||||
</div>
|
||||
{manualMessage && (
|
||||
<p className="mt-3 text-xs text-destructive">{manualMessage}</p>
|
||||
)}
|
||||
{manualReport && (
|
||||
<textarea
|
||||
readOnly
|
||||
value={manualReport}
|
||||
onFocus={(event) => event.currentTarget.select()}
|
||||
className="mt-2 h-28 w-full resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[10px] text-muted-foreground"
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
</motion.div>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -2,8 +2,9 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import type { UpdateStatus } from "@/hooks/use-tauri-update";
|
||||
import type { CopySupportDiagnosticsResult } from "@/lib/tauri-diagnostics";
|
||||
import { AnimatePresence, motion } from "motion/react";
|
||||
import { useEffect, useRef } from "react";
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
|
||||
interface UpdateScreenProps {
|
||||
status: UpdateStatus;
|
||||
|
|
@ -12,6 +13,7 @@ interface UpdateScreenProps {
|
|||
error: string | null;
|
||||
onRetry: () => void;
|
||||
onSkipRestart: () => void;
|
||||
onCopyDiagnostics: () => Promise<CopySupportDiagnosticsResult>;
|
||||
}
|
||||
|
||||
const EASE_OUT_QUART: [number, number, number, number] = [0.165, 0.84, 0.44, 1];
|
||||
|
|
@ -96,11 +98,34 @@ export function UpdateScreen({
|
|||
error,
|
||||
onRetry,
|
||||
onSkipRestart,
|
||||
onCopyDiagnostics,
|
||||
}: UpdateScreenProps) {
|
||||
const isError = status === "error";
|
||||
const [copying, setCopying] = useState(false);
|
||||
const [manualReport, setManualReport] = useState<string | null>(null);
|
||||
const [manualMessage, setManualMessage] = useState<string | null>(null);
|
||||
|
||||
async function handleCopyDiagnostics() {
|
||||
setCopying(true);
|
||||
try {
|
||||
const result = await onCopyDiagnostics();
|
||||
if (result.ok) {
|
||||
setManualReport(null);
|
||||
setManualMessage(null);
|
||||
} else {
|
||||
setManualReport(result.report);
|
||||
setManualMessage(result.error ?? "Clipboard copy failed. Select and copy the diagnostics below.");
|
||||
}
|
||||
} catch (copyError) {
|
||||
setManualReport(null);
|
||||
setManualMessage(`Diagnostics copy failed: ${String(copyError)}`);
|
||||
} finally {
|
||||
setCopying(false);
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="flex h-screen w-full items-center justify-center bg-background">
|
||||
<div className="flex h-full w-full items-center justify-center bg-background">
|
||||
<motion.div
|
||||
initial={{ opacity: 0, y: 8 }}
|
||||
animate={{ opacity: 1, y: 0 }}
|
||||
|
|
@ -148,6 +173,13 @@ export function UpdateScreen({
|
|||
{/* Error actions */}
|
||||
{isError && (
|
||||
<div className="mt-4 flex items-center gap-2">
|
||||
<button
|
||||
type="button"
|
||||
className="rounded-lg bg-muted px-5 py-2.5 text-sm font-medium text-foreground transition-colors hover:bg-muted/80"
|
||||
onClick={() => void handleCopyDiagnostics()}
|
||||
>
|
||||
{copying ? "Copying..." : "Copy Diagnostics"}
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className="rounded-lg bg-primary px-5 py-2.5 text-sm font-medium text-primary-foreground transition-colors hover:bg-primary/80"
|
||||
|
|
@ -165,6 +197,18 @@ export function UpdateScreen({
|
|||
</div>
|
||||
)}
|
||||
|
||||
{manualMessage && (
|
||||
<p className="mt-3 max-w-xl text-center text-xs text-destructive">{manualMessage}</p>
|
||||
)}
|
||||
{manualReport && (
|
||||
<textarea
|
||||
readOnly
|
||||
value={manualReport}
|
||||
onFocus={(event) => event.currentTarget.select()}
|
||||
className="mt-2 h-32 w-full max-w-xl resize-none rounded-lg border border-border/50 bg-muted/30 p-2 font-mono text-[10px] text-muted-foreground"
|
||||
/>
|
||||
)}
|
||||
|
||||
{/* Log viewer */}
|
||||
<LogViewer logs={logs} />
|
||||
</motion.div>
|
||||
|
|
|
|||
329
studio/frontend/src/components/tauri/window-titlebar.tsx
Normal file
329
studio/frontend/src/components/tauri/window-titlebar.tsx
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { useSidebarPin } from "@/hooks/use-sidebar-pin";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { cn } from "@/lib/utils";
|
||||
import type { Window as TauriWindow } from "@tauri-apps/api/window";
|
||||
import {
|
||||
type MouseEvent,
|
||||
type ReactElement,
|
||||
type ReactNode,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useState,
|
||||
} from "react";
|
||||
|
||||
const CUSTOM_TITLEBAR_PLATFORMS = ["win", "linux", "x11"] as const;
|
||||
|
||||
type WindowResizeDirection =
|
||||
| "East"
|
||||
| "North"
|
||||
| "NorthEast"
|
||||
| "NorthWest"
|
||||
| "South"
|
||||
| "SouthEast"
|
||||
| "SouthWest"
|
||||
| "West";
|
||||
|
||||
type NavigatorWithUserAgentData = Navigator & {
|
||||
userAgentData?: {
|
||||
platform?: string;
|
||||
};
|
||||
};
|
||||
|
||||
function getClientPlatform(): string {
|
||||
if (typeof navigator === "undefined") {
|
||||
return "";
|
||||
}
|
||||
const nav = navigator as NavigatorWithUserAgentData;
|
||||
return (
|
||||
nav.userAgentData?.platform ??
|
||||
navigator.platform ??
|
||||
navigator.userAgent
|
||||
).toLowerCase();
|
||||
}
|
||||
|
||||
export function shouldUseCustomWindowTitlebar(): boolean {
|
||||
if (!isTauri) {
|
||||
return false;
|
||||
}
|
||||
const platform = getClientPlatform();
|
||||
if (!platform || platform.includes("mac")) {
|
||||
return false;
|
||||
}
|
||||
return CUSTOM_TITLEBAR_PLATFORMS.some((token) => platform.includes(token));
|
||||
}
|
||||
|
||||
async function getAppWindow(): Promise<TauriWindow> {
|
||||
const { getCurrentWindow } = await import("@tauri-apps/api/window");
|
||||
return getCurrentWindow();
|
||||
}
|
||||
|
||||
function WindowControlButton({
|
||||
label,
|
||||
className,
|
||||
onClick,
|
||||
children,
|
||||
}: {
|
||||
label: string;
|
||||
className?: string;
|
||||
onClick: () => void;
|
||||
children: ReactNode;
|
||||
}): ReactElement {
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
aria-label={label}
|
||||
title={label}
|
||||
onClick={onClick}
|
||||
className={cn(
|
||||
"relative z-[80] inline-flex size-8 items-center justify-center rounded-md text-muted-foreground transition-colors hover:bg-muted/80 hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
|
||||
className,
|
||||
)}
|
||||
>
|
||||
{children}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
function MinimizeGlyph(): ReactElement {
|
||||
return (
|
||||
<span aria-hidden="true" className="h-px w-3.5 rounded-full bg-current" />
|
||||
);
|
||||
}
|
||||
|
||||
function MaximizeGlyph(): ReactElement {
|
||||
return (
|
||||
<span
|
||||
aria-hidden="true"
|
||||
className="size-3 rounded-[2px] border border-current"
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
function RestoreGlyph(): ReactElement {
|
||||
return (
|
||||
<span aria-hidden="true" className="relative size-3.5">
|
||||
<span className="absolute left-0.5 top-0 size-2.5 rounded-[2px] border border-current" />
|
||||
<span className="absolute bottom-0 right-0 size-2.5 rounded-[2px] border border-current bg-muted" />
|
||||
</span>
|
||||
);
|
||||
}
|
||||
|
||||
function CloseGlyph(): ReactElement {
|
||||
return (
|
||||
<span aria-hidden="true" className="relative size-3.5">
|
||||
<span className="absolute left-1/2 top-0 h-3.5 w-px -translate-x-1/2 rotate-45 rounded-full bg-current" />
|
||||
<span className="absolute left-1/2 top-0 h-3.5 w-px -translate-x-1/2 -rotate-45 rounded-full bg-current" />
|
||||
</span>
|
||||
);
|
||||
}
|
||||
|
||||
export function WindowTitlebar({
|
||||
showSidebarSurface = false,
|
||||
}: {
|
||||
showSidebarSurface?: boolean;
|
||||
}): ReactElement | null {
|
||||
const [enabled] = useState(shouldUseCustomWindowTitlebar);
|
||||
const [maximized, setMaximized] = useState(false);
|
||||
const { pinned } = useSidebarPin();
|
||||
|
||||
const refreshMaximized = useCallback(async () => {
|
||||
if (!enabled) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const appWindow = await getAppWindow();
|
||||
setMaximized(await appWindow.isMaximized());
|
||||
} catch {
|
||||
// If a window permission is not ready yet, keep the previous visual state.
|
||||
}
|
||||
}, [enabled]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!enabled) {
|
||||
return;
|
||||
}
|
||||
let mounted = true;
|
||||
let unlistenResize: (() => void) | undefined;
|
||||
let unlistenFocus: (() => void) | undefined;
|
||||
|
||||
const setupWindowListeners = async () => {
|
||||
try {
|
||||
const appWindow = await getAppWindow();
|
||||
if (!mounted) {
|
||||
return;
|
||||
}
|
||||
setMaximized(await appWindow.isMaximized());
|
||||
unlistenResize = await appWindow.onResized(() => {
|
||||
refreshMaximized().catch(() => undefined);
|
||||
});
|
||||
unlistenFocus = await appWindow.onFocusChanged(() => {
|
||||
refreshMaximized().catch(() => undefined);
|
||||
});
|
||||
} catch {
|
||||
// Missing capabilities should not break the rest of the app shell.
|
||||
}
|
||||
};
|
||||
|
||||
setupWindowListeners().catch(() => undefined);
|
||||
|
||||
return () => {
|
||||
mounted = false;
|
||||
unlistenResize?.();
|
||||
unlistenFocus?.();
|
||||
};
|
||||
}, [enabled, refreshMaximized]);
|
||||
|
||||
const runWindowAction = useCallback(
|
||||
(action: (appWindow: TauriWindow) => Promise<void>) => {
|
||||
const runAction = async () => {
|
||||
try {
|
||||
const appWindow = await getAppWindow();
|
||||
await action(appWindow);
|
||||
await refreshMaximized();
|
||||
} catch {
|
||||
// Keep custom chrome inert rather than throwing into React on denied commands.
|
||||
}
|
||||
};
|
||||
|
||||
runAction().catch(() => undefined);
|
||||
},
|
||||
[refreshMaximized],
|
||||
);
|
||||
|
||||
const handleDragMouseDown = useCallback(
|
||||
(event: MouseEvent<HTMLDivElement>) => {
|
||||
if (event.button !== 0 || event.detail > 1) {
|
||||
return;
|
||||
}
|
||||
runWindowAction((appWindow) => appWindow.startDragging());
|
||||
},
|
||||
[runWindowAction],
|
||||
);
|
||||
|
||||
const handleDragDoubleClick = useCallback(
|
||||
(event: MouseEvent<HTMLDivElement>) => {
|
||||
if (event.button !== 0) {
|
||||
return;
|
||||
}
|
||||
runWindowAction((appWindow) => appWindow.toggleMaximize());
|
||||
},
|
||||
[runWindowAction],
|
||||
);
|
||||
|
||||
const handleResizeMouseDown = useCallback(
|
||||
(direction: WindowResizeDirection) =>
|
||||
(event: MouseEvent<HTMLDivElement>) => {
|
||||
if (event.button !== 0) {
|
||||
return;
|
||||
}
|
||||
event.preventDefault();
|
||||
event.stopPropagation();
|
||||
runWindowAction(async (appWindow) => {
|
||||
if (!(await appWindow.isResizable())) {
|
||||
return;
|
||||
}
|
||||
await appWindow.startResizeDragging(direction);
|
||||
});
|
||||
},
|
||||
[runWindowAction],
|
||||
);
|
||||
|
||||
if (!enabled) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<header
|
||||
className="relative z-[60] flex h-[var(--studio-titlebar-height)] shrink-0 select-none items-center text-foreground"
|
||||
aria-label="Window titlebar"
|
||||
>
|
||||
{showSidebarSurface && (
|
||||
<div
|
||||
className="h-full shrink-0 border-r border-sidebar-border bg-sidebar"
|
||||
style={{ width: pinned ? "16rem" : "3rem" }}
|
||||
onMouseDown={handleDragMouseDown}
|
||||
onDoubleClick={handleDragDoubleClick}
|
||||
aria-hidden="true"
|
||||
/>
|
||||
)}
|
||||
<div
|
||||
className="h-full min-w-0 flex-1 border-b border-border/35 bg-muted/35"
|
||||
onMouseDown={handleDragMouseDown}
|
||||
onDoubleClick={handleDragDoubleClick}
|
||||
aria-hidden="true"
|
||||
/>
|
||||
<div
|
||||
className="flex h-full shrink-0 items-center gap-0.5 border-b border-border/35 bg-muted/35 px-1"
|
||||
role="toolbar"
|
||||
aria-label="Window controls"
|
||||
>
|
||||
<WindowControlButton
|
||||
label="Minimize window"
|
||||
onClick={() => runWindowAction((appWindow) => appWindow.minimize())}
|
||||
>
|
||||
<MinimizeGlyph />
|
||||
</WindowControlButton>
|
||||
<WindowControlButton
|
||||
label={maximized ? "Restore window" : "Maximize window"}
|
||||
onClick={() =>
|
||||
runWindowAction((appWindow) => appWindow.toggleMaximize())
|
||||
}
|
||||
>
|
||||
{maximized ? <RestoreGlyph /> : <MaximizeGlyph />}
|
||||
</WindowControlButton>
|
||||
<WindowControlButton
|
||||
label="Close window"
|
||||
onClick={() => runWindowAction((appWindow) => appWindow.close())}
|
||||
className="hover:bg-destructive hover:text-destructive-foreground focus-visible:ring-destructive/70"
|
||||
>
|
||||
<CloseGlyph />
|
||||
</WindowControlButton>
|
||||
</div>
|
||||
</header>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed inset-x-2 top-0 z-[70] h-1 cursor-n-resize"
|
||||
onMouseDown={handleResizeMouseDown("North")}
|
||||
/>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed inset-x-2 bottom-0 z-[70] h-1 cursor-s-resize"
|
||||
onMouseDown={handleResizeMouseDown("South")}
|
||||
/>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed inset-y-2 left-0 z-[70] w-1 cursor-w-resize"
|
||||
onMouseDown={handleResizeMouseDown("West")}
|
||||
/>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed inset-y-2 right-0 z-[70] w-1 cursor-e-resize"
|
||||
onMouseDown={handleResizeMouseDown("East")}
|
||||
/>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed left-0 top-0 z-[70] size-3 cursor-nw-resize"
|
||||
onMouseDown={handleResizeMouseDown("NorthWest")}
|
||||
/>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed right-0 top-0 z-[70] size-3 cursor-ne-resize"
|
||||
onMouseDown={handleResizeMouseDown("NorthEast")}
|
||||
/>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed bottom-0 left-0 z-[70] size-3 cursor-sw-resize"
|
||||
onMouseDown={handleResizeMouseDown("SouthWest")}
|
||||
/>
|
||||
<div
|
||||
aria-hidden="true"
|
||||
className="fixed bottom-0 right-0 z-[70] size-3 cursor-se-resize"
|
||||
onMouseDown={handleResizeMouseDown("SouthEast")}
|
||||
/>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
|
@ -18,6 +18,32 @@ type RefreshResponse = {
|
|||
|
||||
let isRedirecting = false;
|
||||
|
||||
const TAURI_FETCH_RETRY_DELAYS_MS = [250, 750, 1500] as const;
|
||||
|
||||
function wait(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
async function fetchWithTauriNetworkRetry(
|
||||
input: RequestInfo | URL,
|
||||
init?: RequestInit,
|
||||
): Promise<Response> {
|
||||
for (let attempt = 0; ; attempt++) {
|
||||
try {
|
||||
return await fetch(input, init);
|
||||
} catch (error) {
|
||||
if (
|
||||
!isTauri ||
|
||||
!(error instanceof TypeError) ||
|
||||
attempt >= TAURI_FETCH_RETRY_DELAYS_MS.length
|
||||
) {
|
||||
throw error;
|
||||
}
|
||||
await wait(TAURI_FETCH_RETRY_DELAYS_MS[attempt]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async function isPasswordChangeRequiredResponse(response: Response): Promise<boolean> {
|
||||
if (response.status !== 403) return false;
|
||||
|
||||
|
|
@ -54,7 +80,7 @@ async function retryWithCurrentToken(
|
|||
const retryHeaders = new Headers(init?.headers);
|
||||
const token = getAuthToken();
|
||||
if (token) retryHeaders.set("Authorization", `Bearer ${token}`);
|
||||
return fetch(input, { ...init, headers: retryHeaders });
|
||||
return fetchWithTauriNetworkRetry(input, { ...init, headers: retryHeaders });
|
||||
}
|
||||
|
||||
async function retryWithTauriAutoAuth(
|
||||
|
|
@ -72,11 +98,14 @@ export async function refreshSession(): Promise<boolean> {
|
|||
if (!refreshToken) return false;
|
||||
|
||||
try {
|
||||
const response = await fetch(apiUrl("/api/auth/refresh"), {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ refresh_token: refreshToken }),
|
||||
});
|
||||
const response = await fetchWithTauriNetworkRetry(
|
||||
apiUrl("/api/auth/refresh"),
|
||||
{
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ refresh_token: refreshToken }),
|
||||
},
|
||||
);
|
||||
|
||||
if (!response.ok) {
|
||||
clearAuthTokens();
|
||||
|
|
@ -108,7 +137,10 @@ export async function authFetch(
|
|||
|
||||
let response: Response;
|
||||
try {
|
||||
response = await fetch(resolvedInput, { ...init, headers });
|
||||
response = await fetchWithTauriNetworkRetry(resolvedInput, {
|
||||
...init,
|
||||
headers,
|
||||
});
|
||||
} catch (err) {
|
||||
if (err instanceof TypeError) {
|
||||
throw new Error("Studio isn't running -- please relaunch it.");
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import { AuthForm } from "./components/auth-form";
|
|||
|
||||
export function ChangePasswordPage() {
|
||||
return (
|
||||
<div className="relative flex min-h-screen items-center justify-center overflow-hidden bg-background px-4 py-8 sm:px-6 sm:py-10 md:px-10">
|
||||
<div className="relative flex min-h-[calc(100dvh-var(--studio-titlebar-height,0px))] items-center justify-center overflow-hidden bg-background px-4 py-8 sm:px-6 sm:py-10 md:px-10">
|
||||
<LightRays
|
||||
count={6}
|
||||
color="rgba(34, 197, 94, 0.25)"
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import { AuthForm } from "./components/auth-form";
|
|||
|
||||
export function LoginPage() {
|
||||
return (
|
||||
<div className="relative flex min-h-screen items-center justify-center overflow-hidden bg-background px-4 py-8 sm:px-6 sm:py-10 md:px-10">
|
||||
<div className="relative flex min-h-[calc(100dvh-var(--studio-titlebar-height,0px))] items-center justify-center overflow-hidden bg-background px-4 py-8 sm:px-6 sm:py-10 md:px-10">
|
||||
<LightRays
|
||||
count={6}
|
||||
color="rgba(34, 197, 94, 0.25)"
|
||||
|
|
|
|||
|
|
@ -15,9 +15,13 @@ type DesktopAuthResponse = {
|
|||
refresh_token: string;
|
||||
};
|
||||
|
||||
type TauriAutoAuthOptions = {
|
||||
force?: boolean;
|
||||
};
|
||||
|
||||
// Concurrency guard: multiple route guards can call tauriAutoAuth simultaneously.
|
||||
// Without this, the first-launch password-change could race with itself.
|
||||
let pending: Promise<boolean> | null = null;
|
||||
let pending: { promise: Promise<boolean>; force: boolean } | null = null;
|
||||
let lastTauriAuthFailure: string | null = null;
|
||||
|
||||
const TAURI_AUTH_FAILURE_FALLBACK =
|
||||
|
|
@ -49,15 +53,15 @@ function isBackendNotReady(error: unknown): boolean {
|
|||
return authFailureMessage(error).includes(BACKEND_NOT_READY_MESSAGE);
|
||||
}
|
||||
|
||||
async function doTauriAutoAuth(): Promise<boolean> {
|
||||
async function doTauriAutoAuth(options: TauriAutoAuthOptions): Promise<boolean> {
|
||||
// Desktop must handle password-change state internally in Rust.
|
||||
if (hasAuthToken() && !mustChangePassword()) {
|
||||
if (!options.force && hasAuthToken() && !mustChangePassword()) {
|
||||
clearTauriAuthFailure();
|
||||
return true;
|
||||
}
|
||||
|
||||
// Try refreshing existing session
|
||||
if (hasRefreshToken()) {
|
||||
if (!options.force && hasRefreshToken()) {
|
||||
const refreshed = await refreshSession();
|
||||
if (refreshed && hasAuthToken() && !mustChangePassword()) {
|
||||
clearTauriAuthFailure();
|
||||
|
|
@ -86,10 +90,17 @@ async function doTauriAutoAuth(): Promise<boolean> {
|
|||
* Returns true if authentication succeeded.
|
||||
* Concurrent calls are coalesced into a single in-flight attempt.
|
||||
*/
|
||||
export function tauriAutoAuth(): Promise<boolean> {
|
||||
export function tauriAutoAuth(
|
||||
options: TauriAutoAuthOptions = {},
|
||||
): Promise<boolean> {
|
||||
if (!isTauri) return Promise.resolve(false);
|
||||
if (!pending) {
|
||||
pending = doTauriAutoAuth().finally(() => { pending = null; });
|
||||
const force = options.force === true;
|
||||
if (!pending || (force && !pending.force)) {
|
||||
let promise: Promise<boolean>;
|
||||
promise = doTauriAutoAuth({ force }).finally(() => {
|
||||
if (pending?.promise === promise) pending = null;
|
||||
});
|
||||
pending = { promise, force };
|
||||
}
|
||||
return pending;
|
||||
return pending.promise;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@
|
|||
import type { ChatModelAdapter } from "@assistant-ui/react";
|
||||
import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core";
|
||||
import { toast } from "sonner";
|
||||
import { getAuthToken } from "@/features/auth/session";
|
||||
import { apiUrl } from "@/lib/api-base";
|
||||
import {
|
||||
generateAudio,
|
||||
listCachedGguf,
|
||||
|
|
@ -707,7 +709,42 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
const toolCallParts: ToolCallMessagePart[] = [];
|
||||
let serverMetadata: { usage?: ServerUsage; timings?: ServerTimings } | null = null;
|
||||
|
||||
// Per-run cancellation token so a delayed stop POST cannot match
|
||||
// the next run on the same thread.
|
||||
const cancelId =
|
||||
typeof crypto !== "undefined" && "randomUUID" in crypto
|
||||
? crypto.randomUUID()
|
||||
: `${Date.now()}-${Math.random().toString(36).slice(2)}`;
|
||||
|
||||
// Colab-style proxies can swallow fetch aborts, so also POST
|
||||
// /inference/cancel explicitly on abort.
|
||||
const onAbortCancel = () => {
|
||||
const body: Record<string, string> = { cancel_id: cancelId };
|
||||
if (resolvedThreadId) body.session_id = resolvedThreadId;
|
||||
// Plain fetch, not authFetch: authFetch redirects to login on
|
||||
// 401, which would kick the user out mid-stop.
|
||||
const token = getAuthToken();
|
||||
// Use apiUrl so the cancel POST reaches the right origin in
|
||||
// Tauri production builds (where the webview origin is not the
|
||||
// backend at 127.0.0.1:<port>). Browser/dev builds get the empty
|
||||
// base, so the path is unchanged there.
|
||||
void fetch(apiUrl("/api/inference/cancel"), {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
...(token ? { Authorization: `Bearer ${token}` } : {}),
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
keepalive: true,
|
||||
}).catch(() => {});
|
||||
};
|
||||
try {
|
||||
if (abortSignal.aborted) {
|
||||
onAbortCancel();
|
||||
} else {
|
||||
abortSignal.addEventListener("abort", onAbortCancel, { once: true });
|
||||
}
|
||||
|
||||
const {
|
||||
supportsReasoning,
|
||||
reasoningEnabled,
|
||||
|
|
@ -730,6 +767,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
presence_penalty: params.presencePenalty,
|
||||
image_base64: imageBase64,
|
||||
audio_base64: audioBase64,
|
||||
cancel_id: cancelId,
|
||||
...(resolvedThreadId ? { session_id: resolvedThreadId } : {}),
|
||||
...(useAdapter === undefined ? {} : { use_adapter: useAdapter }),
|
||||
...(supportsReasoning
|
||||
? reasoningStyle === "reasoning_effort"
|
||||
|
|
@ -750,7 +789,6 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
const mins = useChatRuntimeStore.getState().toolCallTimeout;
|
||||
return mins >= 9999 ? 9999 : mins * 60;
|
||||
})(),
|
||||
session_id: resolvedThreadId,
|
||||
}
|
||||
: {}),
|
||||
},
|
||||
|
|
@ -948,6 +986,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
}
|
||||
throw err;
|
||||
} finally {
|
||||
abortSignal.removeEventListener("abort", onAbortCancel);
|
||||
runtime.setGeneratingStatus(null);
|
||||
runtime.setToolStatus(null);
|
||||
clearTimeout(warmupTimer);
|
||||
|
|
|
|||
|
|
@ -68,7 +68,11 @@ export async function loadModel(
|
|||
const response = await authFetch("/api/inference/load", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
body: JSON.stringify({
|
||||
...payload,
|
||||
native_path_lease: payload.nativePathLease ?? null,
|
||||
nativePathLease: undefined,
|
||||
}),
|
||||
});
|
||||
return parseJsonOrThrow<LoadModelResponse>(response);
|
||||
}
|
||||
|
|
@ -81,6 +85,7 @@ export async function validateModel(
|
|||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
model_path: payload.model_path,
|
||||
native_path_lease: payload.nativePathLease ?? null,
|
||||
hf_token: payload.hf_token,
|
||||
gguf_variant: payload.gguf_variant ?? null,
|
||||
}),
|
||||
|
|
@ -219,6 +224,25 @@ export async function deleteCachedModel(repoId: string, variant?: string): Promi
|
|||
await parseJsonOrThrow<unknown>(response);
|
||||
}
|
||||
|
||||
export async function deleteFineTunedModel(args: {
|
||||
modelPath: string;
|
||||
source: "training" | "exported";
|
||||
exportType?: "lora" | "merged" | "gguf";
|
||||
ggufVariant?: string;
|
||||
}): Promise<void> {
|
||||
const response = await authFetch("/api/models/delete-finetuned", {
|
||||
method: "DELETE",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
model_path: args.modelPath,
|
||||
source: args.source,
|
||||
export_type: args.exportType ?? null,
|
||||
gguf_variant: args.ggufVariant ?? null,
|
||||
}),
|
||||
});
|
||||
await parseJsonOrThrow<unknown>(response);
|
||||
}
|
||||
|
||||
export interface ScanFolderInfo {
|
||||
id: number;
|
||||
path: string;
|
||||
|
|
|
|||
|
|
@ -2,11 +2,20 @@
|
|||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import {
|
||||
type DeletedModelRef,
|
||||
type LoraModelOption,
|
||||
type ModelOption,
|
||||
ModelSelector,
|
||||
} from "@/components/assistant-ui/model-selector";
|
||||
import { Thread } from "@/components/assistant-ui/thread";
|
||||
import { NativeModelChip } from "@/features/native-intents/components/native-model-chip";
|
||||
import { NativeModelDropOverlay } from "@/features/native-intents/components/native-model-drop-overlay";
|
||||
import { useChooseNativeModel } from "@/features/native-intents/use-native-dialogs";
|
||||
import { useNativeModelDrop } from "@/features/native-intents/use-native-drop";
|
||||
import { useNativePathLeasesSupported } from "@/features/native-intents/use-native-readiness";
|
||||
import { useNativeIntentStore } from "@/features/native-intents/store";
|
||||
import type { NativeIntent } from "@/features/native-intents/types";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { GuidedTour, useGuidedTourController } from "@/features/tour";
|
||||
import { useSidebar } from "@/components/ui/sidebar";
|
||||
|
|
@ -129,6 +138,17 @@ type CompareModelSelection = {
|
|||
ggufVariant?: string;
|
||||
};
|
||||
|
||||
function modelMatchesDeleted(
|
||||
model: { id: string; ggufVariant?: string | null },
|
||||
deletedModel?: DeletedModelRef,
|
||||
): boolean {
|
||||
if (!deletedModel || model.id !== deletedModel.id) return false;
|
||||
return (
|
||||
deletedModel.ggufVariant == null ||
|
||||
(model.ggufVariant ?? null) === deletedModel.ggufVariant
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Detect if this is a LoRA base-vs-fine-tuned compare.
|
||||
* Returns true when the loaded checkpoint is a LoRA — in that case
|
||||
|
|
@ -147,11 +167,15 @@ const CompareContent = memo(function CompareContent({
|
|||
models,
|
||||
loraModels,
|
||||
onFoldersChange,
|
||||
onModelsChange,
|
||||
deleteDisabled,
|
||||
}: {
|
||||
pairId: string;
|
||||
models: ModelOption[];
|
||||
loraModels: LoraModelOption[];
|
||||
onFoldersChange?: () => void;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
}): ReactElement {
|
||||
const isLoraCompare = useIsLoraCompare();
|
||||
|
||||
|
|
@ -163,6 +187,8 @@ const CompareContent = memo(function CompareContent({
|
|||
models={models}
|
||||
loraModels={loraModels}
|
||||
onFoldersChange={onFoldersChange}
|
||||
onModelsChange={onModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
/>
|
||||
);
|
||||
});
|
||||
|
|
@ -331,6 +357,8 @@ function GeneralCompareHeader({
|
|||
value,
|
||||
onValueChange,
|
||||
onFoldersChange,
|
||||
onModelsChange,
|
||||
deleteDisabled,
|
||||
side,
|
||||
}: {
|
||||
models: ModelOption[];
|
||||
|
|
@ -341,6 +369,8 @@ function GeneralCompareHeader({
|
|||
meta: { isLora: boolean; ggufVariant?: string },
|
||||
) => void;
|
||||
onFoldersChange?: () => void;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
side: "left" | "right";
|
||||
}): ReactElement {
|
||||
return (
|
||||
|
|
@ -356,6 +386,8 @@ function GeneralCompareHeader({
|
|||
value={value}
|
||||
onValueChange={onValueChange}
|
||||
onFoldersChange={onFoldersChange}
|
||||
onModelsChange={onModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
variant="ghost"
|
||||
className="max-w-[80%] !h-[34px]"
|
||||
/>
|
||||
|
|
@ -369,11 +401,15 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
|
|||
models,
|
||||
loraModels,
|
||||
onFoldersChange,
|
||||
onModelsChange,
|
||||
deleteDisabled,
|
||||
}: {
|
||||
pairId: string;
|
||||
models: ModelOption[];
|
||||
loraModels: LoraModelOption[];
|
||||
onFoldersChange?: () => void;
|
||||
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
|
||||
deleteDisabled?: boolean;
|
||||
}): ReactElement {
|
||||
const handlesRef = useRef<Record<string, CompareHandle>>({});
|
||||
const [model1ThreadId, setModel1ThreadId] = useState<string>();
|
||||
|
|
@ -391,6 +427,19 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
|
|||
isLora: false,
|
||||
});
|
||||
|
||||
const handleModelsChange = useCallback(
|
||||
(deletedModel?: DeletedModelRef) => {
|
||||
if (modelMatchesDeleted(model1, deletedModel)) {
|
||||
setModel1({ id: "", isLora: false });
|
||||
}
|
||||
if (modelMatchesDeleted(model2, deletedModel)) {
|
||||
setModel2({ id: "", isLora: false });
|
||||
}
|
||||
onModelsChange?.(deletedModel);
|
||||
},
|
||||
[model1, model2, onModelsChange],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
let isActive = true;
|
||||
db.threads
|
||||
|
|
@ -446,6 +495,8 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
|
|||
})
|
||||
}
|
||||
onFoldersChange={onFoldersChange}
|
||||
onModelsChange={handleModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
/>
|
||||
}
|
||||
/>
|
||||
|
|
@ -469,6 +520,8 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
|
|||
})
|
||||
}
|
||||
onFoldersChange={onFoldersChange}
|
||||
onModelsChange={handleModelsChange}
|
||||
deleteDisabled={deleteDisabled}
|
||||
/>
|
||||
}
|
||||
/>
|
||||
|
|
@ -484,10 +537,6 @@ export function ChatPage(): ReactElement {
|
|||
const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen);
|
||||
const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen);
|
||||
|
||||
useEffect(() => {
|
||||
return () => setSettingsOpen(false);
|
||||
}, [setSettingsOpen]);
|
||||
|
||||
useEffect(() => {
|
||||
const threadId = search.thread;
|
||||
if (!threadId) return;
|
||||
|
|
@ -533,7 +582,11 @@ export function ChatPage(): ReactElement {
|
|||
const modelsFromStore = useChatRuntimeStore((state) => state.models);
|
||||
const lorasFromStore = useChatRuntimeStore((state) => state.loras);
|
||||
const modelsError = useChatRuntimeStore((state) => state.modelsError);
|
||||
const modelLoading = useChatRuntimeStore((state) => state.modelLoading);
|
||||
const activeThreadId = useChatRuntimeStore((state) => state.activeThreadId);
|
||||
const modelOperationInProgress = useChatRuntimeStore(
|
||||
(state) => state.modelLoading,
|
||||
);
|
||||
const {
|
||||
refresh,
|
||||
selectModel,
|
||||
|
|
@ -543,6 +596,8 @@ export function ChatPage(): ReactElement {
|
|||
loadProgress,
|
||||
loadToastDismissed,
|
||||
} = useChatModelRuntime();
|
||||
const pendingNativeModelIntent = useNativeIntentStore((state) => state.pendingModelIntent);
|
||||
const nativePathLeasesSupported = useNativePathLeasesSupported();
|
||||
const refreshRef = useRef(refresh);
|
||||
const selectModelRef = useRef(selectModel);
|
||||
|
||||
|
|
@ -574,6 +629,60 @@ export function ChatPage(): ReactElement {
|
|||
return { mode: "single" };
|
||||
}, [search.thread, search.compare, search.new, activeThreadId]);
|
||||
|
||||
const hasActiveModel = Boolean(inferenceParams.checkpoint);
|
||||
const loadNativeModelIntent = useCallback(
|
||||
async (intent: NativeIntent, loadingDescription: string) => {
|
||||
const label = intent.path.displayLabel || intent.displayLabel || "Local GGUF model";
|
||||
await selectModel({
|
||||
id: label,
|
||||
nativePathToken: intent.path.token,
|
||||
isDownloaded: true,
|
||||
loadingDescription,
|
||||
forceReload: true,
|
||||
throwOnError: true,
|
||||
});
|
||||
useNativeIntentStore.getState().clearModelIntent(intent.id);
|
||||
},
|
||||
[selectModel],
|
||||
);
|
||||
const handleNativeModelDropAutoLoad = useCallback(
|
||||
(intent: NativeIntent) =>
|
||||
loadNativeModelIntent(
|
||||
intent,
|
||||
hasActiveModel
|
||||
? "Replacing with dropped local GGUF model."
|
||||
: "Loading dropped local GGUF model.",
|
||||
),
|
||||
[hasActiveModel, loadNativeModelIntent],
|
||||
);
|
||||
const handleNativeModelPickerAutoLoad = useCallback(
|
||||
(intent: NativeIntent) =>
|
||||
loadNativeModelIntent(intent, "Loading chosen local GGUF model."),
|
||||
[loadNativeModelIntent],
|
||||
);
|
||||
const canAutoLoadPickedNativeModel = useCallback(() => {
|
||||
const store = useChatRuntimeStore.getState();
|
||||
return (
|
||||
view.mode === "single" &&
|
||||
nativePathLeasesSupported &&
|
||||
!loadingModel &&
|
||||
!modelLoading &&
|
||||
!store.modelLoading &&
|
||||
!store.params.checkpoint
|
||||
);
|
||||
}, [loadingModel, modelLoading, nativePathLeasesSupported, view.mode]);
|
||||
const chooseNativeModel = useChooseNativeModel({
|
||||
shouldAutoLoad: canAutoLoadPickedNativeModel,
|
||||
onAutoLoad: handleNativeModelPickerAutoLoad,
|
||||
});
|
||||
const nativeModelDropState = useNativeModelDrop({
|
||||
enabled: view.mode === "single",
|
||||
nativePathLeasesSupported,
|
||||
hasActiveModel,
|
||||
isModelLoading: Boolean(loadingModel) || modelLoading,
|
||||
onAutoLoad: handleNativeModelDropAutoLoad,
|
||||
});
|
||||
|
||||
const handleCheckpointChange = useCallback(
|
||||
(
|
||||
value: string,
|
||||
|
|
@ -738,6 +847,21 @@ export function ChatPage(): ReactElement {
|
|||
.catch(() => {});
|
||||
}, [navigate]);
|
||||
|
||||
const refreshModelLists = useCallback((deletedModel?: DeletedModelRef) => {
|
||||
const { checkpoint } = useChatRuntimeStore.getState().params;
|
||||
const activeGgufVariant = useChatRuntimeStore.getState().activeGgufVariant;
|
||||
if (
|
||||
modelMatchesDeleted(
|
||||
{ id: checkpoint, ggufVariant: activeGgufVariant },
|
||||
deletedModel,
|
||||
)
|
||||
) {
|
||||
useChatRuntimeStore.getState().clearCheckpoint();
|
||||
}
|
||||
void refresh();
|
||||
refreshLocalModels();
|
||||
}, [refresh, refreshLocalModels]);
|
||||
|
||||
const loraModels = useMemo<LoraModelOption[]>(() => {
|
||||
const fromLoras = lorasFromStore.map((lora) => ({
|
||||
id: lora.id,
|
||||
|
|
@ -863,6 +987,7 @@ export function ChatPage(): ReactElement {
|
|||
<div className="flex min-h-0 min-w-0 flex-1 basis-0 bg-background overflow-hidden">
|
||||
<GuidedTour {...tour.tourProps} />
|
||||
<div className="relative flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-hidden">
|
||||
<NativeModelDropOverlay state={nativeModelDropState} />
|
||||
<div
|
||||
className={cn(
|
||||
"absolute top-0 left-0 right-[10px] z-30 flex h-[48px] shrink-0 items-start pt-[11px] pr-2 bg-background",
|
||||
|
|
@ -881,14 +1006,24 @@ export function ChatPage(): ReactElement {
|
|||
onValueChange={handleCheckpointChange}
|
||||
onEject={handleEject}
|
||||
onFoldersChange={refreshLocalModels}
|
||||
onPickLocalModel={isTauri ? chooseNativeModel : undefined}
|
||||
onModelsChange={refreshModelLists}
|
||||
deleteDisabled={modelOperationInProgress}
|
||||
variant="ghost"
|
||||
open={modelSelectorOpen}
|
||||
onOpenChange={handleModelSelectorOpenChange}
|
||||
triggerDataTour="chat-model-selector"
|
||||
contentDataTour="chat-model-selector-popover"
|
||||
className="max-w-[62vw] sm:max-w-none !h-[34px]"
|
||||
className="max-w-[62vw] !pr-3 sm:max-w-none !h-[34px]"
|
||||
/>
|
||||
)}
|
||||
{pendingNativeModelIntent && view.mode !== "compare" ? (
|
||||
<NativeModelChip
|
||||
intent={pendingNativeModelIntent}
|
||||
nativeReadsDisabled={!nativePathLeasesSupported}
|
||||
onLoad={(selection) => selectModel(selection)}
|
||||
/>
|
||||
) : null}
|
||||
{loadingModel && loadToastDismissed ? (
|
||||
<ModelLoadInlineStatus
|
||||
label={
|
||||
|
|
@ -910,12 +1045,17 @@ export function ChatPage(): ReactElement {
|
|||
onStop={cancelLoading}
|
||||
/>
|
||||
) : null}
|
||||
{!loadingModel && modelsError ? (
|
||||
<div
|
||||
className="relative top-0.5 max-w-[28rem] truncate pl-0.5 text-xs text-destructive"
|
||||
title={modelsError}
|
||||
role="status"
|
||||
aria-live="polite"
|
||||
>
|
||||
{modelsError}
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
{modelsError && (
|
||||
<div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">
|
||||
{modelsError}
|
||||
</div>
|
||||
)}
|
||||
<div className="ml-auto flex items-center gap-2">
|
||||
{view.mode === "single" && ggufContextLength && contextUsage ? (
|
||||
<ContextUsageBar
|
||||
|
|
@ -961,6 +1101,8 @@ export function ChatPage(): ReactElement {
|
|||
models={models}
|
||||
loraModels={loraModels}
|
||||
onFoldersChange={refreshLocalModels}
|
||||
onModelsChange={refreshModelLists}
|
||||
deleteDisabled={modelOperationInProgress}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ import {
|
|||
DropdownMenu,
|
||||
DropdownMenuContent,
|
||||
DropdownMenuItem,
|
||||
DropdownMenuSeparator,
|
||||
DropdownMenuTrigger,
|
||||
} from "@/components/ui/dropdown-menu";
|
||||
import { Input } from "@/components/ui/input";
|
||||
|
|
@ -64,54 +65,36 @@ import {
|
|||
} from "@/components/ui/tooltip";
|
||||
import { Tooltip as TooltipPrimitive } from "radix-ui";
|
||||
import { AnimatePresence, motion } from "motion/react";
|
||||
import type { ReactNode } from "react";
|
||||
import { Fragment, type ReactNode } from "react";
|
||||
import { useEffect, useLayoutEffect, useMemo, useRef, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
|
||||
import {
|
||||
DEFAULT_INFERENCE_PARAMS,
|
||||
type InferenceParams,
|
||||
} from "./types/runtime";
|
||||
applyPresetParams,
|
||||
BUILTIN_PRESET_NAMES,
|
||||
BUILTIN_PRESETS,
|
||||
defaultInferenceParams,
|
||||
getBuiltinVariantName,
|
||||
getOrderedPresets,
|
||||
getPresetOwnedConfigKey,
|
||||
getPresetSaveState,
|
||||
getPresetSource,
|
||||
getUniquePresetName,
|
||||
isSamePresetConfig,
|
||||
normalizeCustomPresets,
|
||||
toPresetParams,
|
||||
type Preset,
|
||||
} from "./presets/preset-policy";
|
||||
import type { InferenceParams } from "./types/runtime";
|
||||
|
||||
export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS;
|
||||
export { defaultInferenceParams, type Preset } from "./presets/preset-policy";
|
||||
export type { InferenceParams } from "./types/runtime";
|
||||
|
||||
export interface Preset {
|
||||
name: string;
|
||||
params: InferenceParams;
|
||||
}
|
||||
|
||||
interface LegacySystemPromptTemplate {
|
||||
name: string;
|
||||
content: string;
|
||||
}
|
||||
|
||||
const BUILTIN_PRESETS: Preset[] = [
|
||||
{ name: "Default", params: { ...defaultInferenceParams } },
|
||||
{
|
||||
name: "Creative",
|
||||
params: {
|
||||
...defaultInferenceParams,
|
||||
temperature: 1.5,
|
||||
topP: 1.0,
|
||||
topK: 0,
|
||||
minP: 0.1,
|
||||
repetitionPenalty: 1.0,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "Precise",
|
||||
params: {
|
||||
...defaultInferenceParams,
|
||||
temperature: 0.1,
|
||||
topP: 0.95,
|
||||
topK: 80,
|
||||
minP: 0.01,
|
||||
repetitionPenalty: 1.0,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
const CHAT_PRESETS_KEY = "unsloth_chat_custom_presets";
|
||||
const CHAT_ACTIVE_PRESET_KEY = "unsloth_chat_active_preset";
|
||||
const LEGACY_CHAT_SYSTEM_PROMPTS_KEY = "unsloth_chat_system_prompts";
|
||||
|
|
@ -122,16 +105,13 @@ function canUseStorage(): boolean {
|
|||
return typeof window !== "undefined";
|
||||
}
|
||||
|
||||
function getUniquePresetName(baseName: string, usedNames: Set<string>): string {
|
||||
const normalizedBase = baseName.trim() || "Imported Prompt";
|
||||
let nextName = normalizedBase;
|
||||
let suffix = 2;
|
||||
while (usedNames.has(nextName)) {
|
||||
nextName = `${normalizedBase} ${suffix}`;
|
||||
suffix += 1;
|
||||
function saveCustomPresets(presets: Preset[]): void {
|
||||
if (!canUseStorage()) return;
|
||||
try {
|
||||
localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(presets));
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
usedNames.add(nextName);
|
||||
return nextName;
|
||||
}
|
||||
|
||||
function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] {
|
||||
|
|
@ -161,18 +141,7 @@ function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] {
|
|||
]);
|
||||
const seenImportedConfigKeys = new Set(
|
||||
[...BUILTIN_PRESETS, ...presets].map((preset) =>
|
||||
JSON.stringify({
|
||||
temperature: preset.params.temperature,
|
||||
topP: preset.params.topP,
|
||||
topK: preset.params.topK,
|
||||
minP: preset.params.minP,
|
||||
repetitionPenalty: preset.params.repetitionPenalty,
|
||||
presencePenalty: preset.params.presencePenalty,
|
||||
maxSeqLength: preset.params.maxSeqLength,
|
||||
maxTokens: preset.params.maxTokens,
|
||||
systemPrompt: preset.params.systemPrompt,
|
||||
trustRemoteCode: preset.params.trustRemoteCode ?? false,
|
||||
}),
|
||||
getPresetOwnedConfigKey(preset.params),
|
||||
),
|
||||
);
|
||||
const importedPresets = parsed
|
||||
|
|
@ -191,18 +160,7 @@ function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] {
|
|||
},
|
||||
}))
|
||||
.filter(({ importedParams }) => {
|
||||
const configKey = JSON.stringify({
|
||||
temperature: importedParams.temperature,
|
||||
topP: importedParams.topP,
|
||||
topK: importedParams.topK,
|
||||
minP: importedParams.minP,
|
||||
repetitionPenalty: importedParams.repetitionPenalty,
|
||||
presencePenalty: importedParams.presencePenalty,
|
||||
maxSeqLength: importedParams.maxSeqLength,
|
||||
maxTokens: importedParams.maxTokens,
|
||||
systemPrompt: importedParams.systemPrompt,
|
||||
trustRemoteCode: importedParams.trustRemoteCode ?? false,
|
||||
});
|
||||
const configKey = getPresetOwnedConfigKey(importedParams);
|
||||
if (seenImportedConfigKeys.has(configKey)) return false;
|
||||
seenImportedConfigKeys.add(configKey);
|
||||
return true;
|
||||
|
|
@ -216,8 +174,8 @@ function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] {
|
|||
localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw);
|
||||
return presets;
|
||||
}
|
||||
const mergedPresets = [...presets, ...importedPresets];
|
||||
localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(mergedPresets));
|
||||
const mergedPresets = normalizeCustomPresets([...presets, ...importedPresets]);
|
||||
saveCustomPresets(mergedPresets);
|
||||
try {
|
||||
localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw);
|
||||
localStorage.removeItem(LEGACY_CHAT_SYSTEM_PROMPTS_KEY);
|
||||
|
|
@ -255,7 +213,11 @@ function loadSavedCustomPresets(): Preset[] {
|
|||
},
|
||||
}))
|
||||
.filter((preset) => preset.name.length > 0);
|
||||
return migrateLegacySystemPromptTemplates(presets);
|
||||
const normalized = normalizeCustomPresets(presets);
|
||||
if (JSON.stringify(normalized) !== JSON.stringify(presets)) {
|
||||
saveCustomPresets(normalized);
|
||||
}
|
||||
return migrateLegacySystemPromptTemplates(normalized);
|
||||
} catch {
|
||||
return migrateLegacySystemPromptTemplates([]);
|
||||
}
|
||||
|
|
@ -270,82 +232,6 @@ function loadSavedActivePreset(): string {
|
|||
}
|
||||
}
|
||||
|
||||
type PresetSaveMode =
|
||||
| "disabled"
|
||||
| "overwrite-active"
|
||||
| "overwrite-other"
|
||||
| "create";
|
||||
|
||||
interface PresetSaveState {
|
||||
mode: PresetSaveMode;
|
||||
canSubmit: boolean;
|
||||
isSaveReady: boolean;
|
||||
buttonLabel: string;
|
||||
title: string;
|
||||
}
|
||||
|
||||
function isSamePresetConfig(a: InferenceParams, b: InferenceParams): boolean {
|
||||
return (
|
||||
a.temperature === b.temperature &&
|
||||
a.topP === b.topP &&
|
||||
a.topK === b.topK &&
|
||||
a.minP === b.minP &&
|
||||
a.repetitionPenalty === b.repetitionPenalty &&
|
||||
a.presencePenalty === b.presencePenalty &&
|
||||
a.maxSeqLength === b.maxSeqLength &&
|
||||
a.maxTokens === b.maxTokens &&
|
||||
a.systemPrompt === b.systemPrompt &&
|
||||
(a.trustRemoteCode ?? false) === (b.trustRemoteCode ?? false)
|
||||
);
|
||||
}
|
||||
|
||||
function getPresetSaveState({
|
||||
rawName,
|
||||
activePreset,
|
||||
presets,
|
||||
activePresetDirty,
|
||||
}: {
|
||||
rawName: string;
|
||||
activePreset: string;
|
||||
presets: Preset[];
|
||||
activePresetDirty: boolean;
|
||||
}): PresetSaveState {
|
||||
const trimmedName = rawName.trim();
|
||||
if (!trimmedName) {
|
||||
return {
|
||||
mode: "disabled",
|
||||
canSubmit: false,
|
||||
isSaveReady: false,
|
||||
buttonLabel: "Save",
|
||||
title: "Enter a preset name",
|
||||
};
|
||||
}
|
||||
|
||||
const matchingPreset = presets.find((preset) => preset.name === trimmedName);
|
||||
if (matchingPreset) {
|
||||
const isActiveMatch = matchingPreset.name === activePreset;
|
||||
return {
|
||||
mode: isActiveMatch ? "overwrite-active" : "overwrite-other",
|
||||
canSubmit: !isActiveMatch || activePresetDirty,
|
||||
isSaveReady: !isActiveMatch || activePresetDirty,
|
||||
buttonLabel: isActiveMatch && !activePresetDirty ? "Saved" : "Overwrite",
|
||||
title: isActiveMatch
|
||||
? activePresetDirty
|
||||
? "Save current settings to this preset"
|
||||
: "No unsaved changes"
|
||||
: `Overwrite preset "${trimmedName}"`,
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
mode: "create",
|
||||
canSubmit: true,
|
||||
isSaveReady: true,
|
||||
buttonLabel: "Save as New",
|
||||
title: `Save current settings as "${trimmedName}"`,
|
||||
};
|
||||
}
|
||||
|
||||
function ParamSlider({
|
||||
label,
|
||||
value,
|
||||
|
|
@ -518,6 +404,10 @@ export function ChatSettingsPanel({
|
|||
const setCustomContextLength = useChatRuntimeStore(
|
||||
(s) => s.setCustomContextLength,
|
||||
);
|
||||
const setActivePresetSource = useChatRuntimeStore(
|
||||
(s) => s.setActivePresetSource,
|
||||
);
|
||||
const activePresetSource = useChatRuntimeStore((s) => s.activePresetSource);
|
||||
|
||||
const ctxDisplayValue = customContextLength ?? ggufContextLength ?? "";
|
||||
const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null;
|
||||
|
|
@ -540,12 +430,9 @@ export function ChatSettingsPanel({
|
|||
>(undefined);
|
||||
const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false);
|
||||
const [systemPromptDraft, setSystemPromptDraft] = useState("");
|
||||
const [activePresetBaseline, setActivePresetBaseline] = useState(params);
|
||||
const presets = useMemo(() => {
|
||||
const overrides = new Set(customPresets.map((preset) => preset.name));
|
||||
return [
|
||||
...BUILTIN_PRESETS.filter((preset) => !overrides.has(preset.name)),
|
||||
...customPresets,
|
||||
];
|
||||
return getOrderedPresets(customPresets);
|
||||
}, [customPresets]);
|
||||
const activePresetDefinition = useMemo(
|
||||
() => presets.find((preset) => preset.name === activePreset) ?? null,
|
||||
|
|
@ -555,17 +442,23 @@ export function ChatSettingsPanel({
|
|||
() => customPresets.find((preset) => preset.name === activePreset) ?? null,
|
||||
[activePreset, customPresets],
|
||||
);
|
||||
const activeBuiltinPreset = useMemo(
|
||||
() =>
|
||||
BUILTIN_PRESETS.find((preset) => preset.name === activePreset) ?? null,
|
||||
[activePreset],
|
||||
);
|
||||
const activePresetDirty = useMemo(
|
||||
() =>
|
||||
activePresetDefinition == null
|
||||
? false
|
||||
: !isSamePresetConfig(activePresetDefinition.params, params),
|
||||
[activePresetDefinition, params],
|
||||
const hasUnsavedPresetChanges = useMemo(
|
||||
() => {
|
||||
if (activePresetDefinition == null) {
|
||||
return false;
|
||||
}
|
||||
if (BUILTIN_PRESET_NAMES.has(activePresetDefinition.name)) {
|
||||
if (activePresetDefinition.name === "Default") {
|
||||
return activePresetSource === "modified";
|
||||
}
|
||||
return (
|
||||
activePresetSource === "modified" ||
|
||||
!isSamePresetConfig(activePresetDefinition.params, params)
|
||||
);
|
||||
}
|
||||
return !isSamePresetConfig(activePresetDefinition.params, params);
|
||||
},
|
||||
[activePresetDefinition, activePresetSource, params],
|
||||
);
|
||||
const presetSaveState = useMemo(
|
||||
() =>
|
||||
|
|
@ -573,9 +466,9 @@ export function ChatSettingsPanel({
|
|||
rawName: presetNameInput,
|
||||
activePreset,
|
||||
presets,
|
||||
activePresetDirty,
|
||||
hasUnsavedPresetChanges,
|
||||
}),
|
||||
[activePreset, activePresetDirty, presetNameInput, presets],
|
||||
[activePreset, hasUnsavedPresetChanges, presetNameInput, presets],
|
||||
);
|
||||
const systemPromptEditorDirty = systemPromptDraft !== params.systemPrompt;
|
||||
const trustRemoteCodeMissing =
|
||||
|
|
@ -584,27 +477,24 @@ export function ChatSettingsPanel({
|
|||
!(params.trustRemoteCode ?? false);
|
||||
|
||||
function set<K extends keyof InferenceParams>(key: K) {
|
||||
return (v: InferenceParams[K]) => onParamsChange({ ...params, [key]: v });
|
||||
return (v: InferenceParams[K]) => {
|
||||
const nextParams = { ...params, [key]: v };
|
||||
const nextSource = isSamePresetConfig(activePresetBaseline, nextParams)
|
||||
? getPresetSource(activePreset)
|
||||
: "modified";
|
||||
setActivePresetSource(nextSource);
|
||||
onParamsChange(nextParams);
|
||||
};
|
||||
}
|
||||
|
||||
function applyPreset(name: string) {
|
||||
const p = presets.find((pr) => pr.name === name);
|
||||
if (p) {
|
||||
if (
|
||||
modelRequiresTrustRemoteCode &&
|
||||
!(p.params.trustRemoteCode ?? false)
|
||||
) {
|
||||
toast.warning("This configuration turns custom code off", {
|
||||
description:
|
||||
"The current model needs custom code enabled to load. Keep it on for this model.",
|
||||
});
|
||||
return;
|
||||
}
|
||||
onParamsChange({
|
||||
...p.params,
|
||||
checkpoint: params.checkpoint,
|
||||
...applyPresetParams(params, p.params),
|
||||
});
|
||||
setActivePreset(name);
|
||||
setActivePresetSource(getPresetSource(name));
|
||||
if (canUseStorage()) {
|
||||
try {
|
||||
localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, name);
|
||||
|
|
@ -621,27 +511,29 @@ export function ChatSettingsPanel({
|
|||
toast.error("Enter a preset name");
|
||||
return;
|
||||
}
|
||||
const usedNames = new Set([
|
||||
...BUILTIN_PRESET_NAMES,
|
||||
...customPresets.map((preset) => preset.name),
|
||||
]);
|
||||
const saveName = BUILTIN_PRESET_NAMES.has(trimmed)
|
||||
? getBuiltinVariantName(trimmed, usedNames)
|
||||
: trimmed;
|
||||
setCustomPresets((prev) => {
|
||||
const next = prev.filter((p) => p.name !== trimmed);
|
||||
const merged = [...next, { name: trimmed, params: { ...params } }];
|
||||
if (canUseStorage()) {
|
||||
try {
|
||||
localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(merged));
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
const next = prev.filter((p) => p.name !== saveName);
|
||||
const merged = [...next, { name: saveName, params: toPresetParams(params) }];
|
||||
saveCustomPresets(merged);
|
||||
return merged;
|
||||
});
|
||||
if (canUseStorage()) {
|
||||
try {
|
||||
localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, trimmed);
|
||||
localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, saveName);
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
setActivePreset(trimmed);
|
||||
setPresetNameInput(trimmed);
|
||||
setActivePreset(saveName);
|
||||
setActivePresetSource("custom");
|
||||
setPresetNameInput(saveName);
|
||||
}
|
||||
|
||||
function deletePreset(name: string) {
|
||||
|
|
@ -651,41 +543,20 @@ export function ChatSettingsPanel({
|
|||
if (!hasCustomPreset) {
|
||||
return;
|
||||
}
|
||||
const builtinPreset = BUILTIN_PRESETS.find((preset) => preset.name === name);
|
||||
const fallbackPreset =
|
||||
builtinPreset ??
|
||||
BUILTIN_PRESETS.find((preset) => preset.name === "Default") ??
|
||||
null;
|
||||
if (
|
||||
activePreset === name &&
|
||||
fallbackPreset &&
|
||||
modelRequiresTrustRemoteCode &&
|
||||
!(fallbackPreset.params.trustRemoteCode ?? false)
|
||||
) {
|
||||
toast.warning("Reset would turn custom code off", {
|
||||
description:
|
||||
"The current model needs custom code enabled to load. Keep it on for this model.",
|
||||
});
|
||||
return;
|
||||
}
|
||||
BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? null;
|
||||
setCustomPresets((prev) => {
|
||||
const next = prev.filter((preset) => preset.name !== name);
|
||||
if (canUseStorage()) {
|
||||
try {
|
||||
localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(next));
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
}
|
||||
saveCustomPresets(next);
|
||||
return next;
|
||||
});
|
||||
if (activePreset === name) {
|
||||
if (fallbackPreset) {
|
||||
onParamsChange({
|
||||
...fallbackPreset.params,
|
||||
checkpoint: params.checkpoint,
|
||||
...applyPresetParams(params, fallbackPreset.params),
|
||||
});
|
||||
setActivePreset(fallbackPreset.name);
|
||||
setActivePresetSource("builtin-default");
|
||||
if (canUseStorage()) {
|
||||
try {
|
||||
localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, fallbackPreset.name);
|
||||
|
|
@ -708,8 +579,46 @@ export function ChatSettingsPanel({
|
|||
}
|
||||
|
||||
useEffect(() => {
|
||||
if (presets.some((preset) => preset.name === activePreset)) return;
|
||||
if (activePresetSource !== "modified") {
|
||||
setActivePresetBaseline(params);
|
||||
}
|
||||
}, [activePresetSource, params]);
|
||||
|
||||
useEffect(() => {
|
||||
if (presets.some((preset) => preset.name === activePreset)) {
|
||||
const expectedSource = getPresetSource(activePreset);
|
||||
if (activePresetDefinition != null) {
|
||||
if (BUILTIN_PRESET_NAMES.has(activePresetDefinition.name)) {
|
||||
if (activePresetDefinition.name === "Default") {
|
||||
if (
|
||||
activePresetSource !== "modified" &&
|
||||
activePresetSource !== expectedSource
|
||||
) {
|
||||
setActivePresetSource(expectedSource);
|
||||
}
|
||||
return;
|
||||
}
|
||||
const matchesActivePreset = isSamePresetConfig(
|
||||
activePresetDefinition.params,
|
||||
params,
|
||||
);
|
||||
const nextSource = matchesActivePreset ? expectedSource : "modified";
|
||||
if (activePresetSource !== nextSource) {
|
||||
setActivePresetSource(nextSource);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
if (
|
||||
activePresetSource !== "modified" &&
|
||||
activePresetSource !== expectedSource
|
||||
) {
|
||||
setActivePresetSource(expectedSource);
|
||||
}
|
||||
return;
|
||||
}
|
||||
setActivePreset("Default");
|
||||
setActivePresetSource("builtin-default");
|
||||
if (canUseStorage()) {
|
||||
try {
|
||||
localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, "Default");
|
||||
|
|
@ -717,7 +626,14 @@ export function ChatSettingsPanel({
|
|||
// ignore
|
||||
}
|
||||
}
|
||||
}, [activePreset, presets]);
|
||||
}, [
|
||||
activePreset,
|
||||
activePresetDefinition,
|
||||
activePresetSource,
|
||||
params,
|
||||
presets,
|
||||
setActivePresetSource,
|
||||
]);
|
||||
|
||||
useEffect(() => {
|
||||
setPresetNameInput(activePreset);
|
||||
|
|
@ -741,6 +657,187 @@ export function ChatSettingsPanel({
|
|||
return () => ro.disconnect();
|
||||
}, [open]);
|
||||
|
||||
const modelSection = (
|
||||
<CollapsibleSection
|
||||
icon={Settings02Icon}
|
||||
label="Model"
|
||||
defaultOpen={true}
|
||||
>
|
||||
<div className="flex flex-col gap-3 py-1">
|
||||
{isGguf && (
|
||||
<>
|
||||
<div className="space-y-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium">Context Length</span>
|
||||
<Input
|
||||
type="number"
|
||||
value={
|
||||
typeof ctxDisplayValue === "number"
|
||||
? ctxDisplayValue
|
||||
: (ggufContextLength ?? "")
|
||||
}
|
||||
placeholder="..."
|
||||
min={128}
|
||||
max={ctxMaxValue ?? undefined}
|
||||
step={1024}
|
||||
className="h-6 w-[100px] text-right text-xs tabular-nums"
|
||||
onChange={(e) => {
|
||||
const raw = e.target.value;
|
||||
if (raw === "") {
|
||||
setCustomContextLength(null);
|
||||
return;
|
||||
}
|
||||
const v = Number.parseInt(raw, 10);
|
||||
if (!Number.isNaN(v) && v >= 0) {
|
||||
const maxCtx = ctxMaxValue ?? Number.POSITIVE_INFINITY;
|
||||
const clamped = Math.min(v, maxCtx);
|
||||
setCustomContextLength(
|
||||
clamped === (ggufContextLength ?? 0) ? null : clamped,
|
||||
);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
<Slider
|
||||
min={1024}
|
||||
max={ctxMaxValue ?? 4096}
|
||||
step={1024}
|
||||
value={[
|
||||
Math.min(
|
||||
typeof ctxDisplayValue === "number"
|
||||
? ctxDisplayValue
|
||||
: (ggufContextLength ?? 4096),
|
||||
ctxMaxValue ?? 4096,
|
||||
),
|
||||
]}
|
||||
onValueChange={([v]) => {
|
||||
setCustomContextLength(
|
||||
v === (ggufContextLength ?? 0) ? null : v,
|
||||
);
|
||||
}}
|
||||
/>
|
||||
{ggufMaxContextLength != null &&
|
||||
typeof ctxDisplayValue === "number" &&
|
||||
ctxDisplayValue > ggufMaxContextLength && (
|
||||
<p className="text-[11px] text-amber-500">
|
||||
Exceeds estimated VRAM capacity (
|
||||
{ggufMaxContextLength.toLocaleString()} tokens). The model
|
||||
may use system RAM.
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
<div className="grid grid-cols-[minmax(0,1fr)_65px] items-center gap-x-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">KV Cache Dtype</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Quantize KV cache to reduce VRAM.
|
||||
</div>
|
||||
</div>
|
||||
<div className="w-full min-w-0">
|
||||
<Select
|
||||
value={kvCacheDtype ?? "f16"}
|
||||
onValueChange={(v) => {
|
||||
setKvCacheDtype(v === "f16" ? null : v);
|
||||
}}
|
||||
>
|
||||
<SelectTrigger className="grid h-7 w-full min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 px-2 py-0 text-xs [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="f16">f16</SelectItem>
|
||||
<SelectItem value="bf16">bf16</SelectItem>
|
||||
<SelectItem value="q8_0">q8_0</SelectItem>
|
||||
<SelectItem value="q5_1">q5_1</SelectItem>
|
||||
<SelectItem value="q4_1">q4_1</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</div>
|
||||
{!currentModelIsVision && (
|
||||
<div className="grid grid-cols-[minmax(0,1fr)_65px] items-center gap-x-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">
|
||||
Speculative Decoding
|
||||
</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Speed up generation with no VRAM cost.
|
||||
</div>
|
||||
</div>
|
||||
<div className="w-full min-w-0">
|
||||
<Select
|
||||
value={speculativeType ?? "off"}
|
||||
onValueChange={(v) => {
|
||||
setSpeculativeType(v === "off" ? null : v);
|
||||
}}
|
||||
>
|
||||
<SelectTrigger className="grid h-7 w-full min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 px-2 py-0 text-xs [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="default">On</SelectItem>
|
||||
<SelectItem value="off">Off</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
{modelSettingsDirty && (
|
||||
<div className="flex flex-wrap gap-1.5 pt-1">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => onReloadModel?.()}
|
||||
className="rounded-md bg-primary px-2.5 py-1 text-[11px] font-medium text-primary-foreground transition-colors hover:bg-primary/90"
|
||||
>
|
||||
Apply
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setCustomContextLength(null);
|
||||
setKvCacheDtype(loadedKvCacheDtype);
|
||||
setSpeculativeType(loadedSpeculativeType);
|
||||
}}
|
||||
className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
|
||||
>
|
||||
Reset
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
{!isGguf && params.checkpoint && (
|
||||
<>
|
||||
<div className="flex items-center justify-between gap-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">Enable custom code</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Allow models with custom code (e.g. Nemotron). Only enable if
|
||||
sure.
|
||||
</div>
|
||||
</div>
|
||||
<Switch
|
||||
checked={params.trustRemoteCode ?? false}
|
||||
onCheckedChange={set("trustRemoteCode")}
|
||||
/>
|
||||
</div>
|
||||
{trustRemoteCodeMissing && (
|
||||
<Alert className="border-amber-200/70 bg-amber-50/70 px-3 py-2 text-amber-950 dark:border-amber-900/70 dark:bg-amber-950/35 dark:text-amber-100">
|
||||
<AlertTitle className="text-[11px] font-medium">
|
||||
Keep custom code enabled for this model
|
||||
</AlertTitle>
|
||||
<AlertDescription className="text-[11px] text-amber-800 dark:text-amber-200">
|
||||
This model requires custom code to load. You can edit the
|
||||
toggle, but loading will stay blocked until it is turned back
|
||||
on.
|
||||
</AlertDescription>
|
||||
</Alert>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</CollapsibleSection>
|
||||
);
|
||||
|
||||
const settingsContent = (
|
||||
<>
|
||||
<div className="aui-thread-viewport relative h-full overflow-y-auto">
|
||||
|
|
@ -834,13 +931,16 @@ export function ChatSettingsPanel({
|
|||
: undefined
|
||||
}
|
||||
>
|
||||
{presets.map((p) => (
|
||||
<DropdownMenuItem
|
||||
key={p.name}
|
||||
onSelect={() => applyPreset(p.name)}
|
||||
>
|
||||
{p.name}
|
||||
</DropdownMenuItem>
|
||||
{presets.map((p, index) => (
|
||||
<Fragment key={p.name}>
|
||||
<DropdownMenuItem onSelect={() => applyPreset(p.name)}>
|
||||
{p.name}
|
||||
</DropdownMenuItem>
|
||||
{index === BUILTIN_PRESETS.length - 1 &&
|
||||
presets.length > BUILTIN_PRESETS.length && (
|
||||
<DropdownMenuSeparator className="mx-2.5! my-1.5! h-0! border-t border-border/70 bg-transparent!" />
|
||||
)}
|
||||
</Fragment>
|
||||
))}
|
||||
</DropdownMenuContent>
|
||||
</DropdownMenu>
|
||||
|
|
@ -874,9 +974,7 @@ export function ChatSettingsPanel({
|
|||
className="h-8 w-full text-xs text-muted-foreground"
|
||||
title={
|
||||
activeCustomPreset
|
||||
? activeBuiltinPreset
|
||||
? "Reset selected preset to built-in defaults"
|
||||
: "Delete selected preset"
|
||||
? "Delete selected preset"
|
||||
: "No saved override to delete"
|
||||
}
|
||||
>
|
||||
|
|
@ -918,188 +1016,6 @@ export function ChatSettingsPanel({
|
|||
/>
|
||||
</div>
|
||||
|
||||
<CollapsibleSection
|
||||
icon={Settings02Icon}
|
||||
label="Model"
|
||||
defaultOpen={true}
|
||||
>
|
||||
<div className="flex flex-col gap-3 py-1">
|
||||
{isGguf && (
|
||||
<>
|
||||
<div className="space-y-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="text-xs font-medium">Context Length</span>
|
||||
<Input
|
||||
type="number"
|
||||
value={
|
||||
typeof ctxDisplayValue === "number"
|
||||
? ctxDisplayValue
|
||||
: (ggufContextLength ?? "")
|
||||
}
|
||||
placeholder="..."
|
||||
min={128}
|
||||
max={ctxMaxValue ?? undefined}
|
||||
step={1024}
|
||||
className="h-6 w-[100px] text-right text-xs tabular-nums"
|
||||
onChange={(e) => {
|
||||
const raw = e.target.value;
|
||||
if (raw === "") {
|
||||
setCustomContextLength(null);
|
||||
return;
|
||||
}
|
||||
const v = Number.parseInt(raw, 10);
|
||||
if (!Number.isNaN(v) && v >= 0) {
|
||||
const maxCtx =
|
||||
ctxMaxValue ?? Number.POSITIVE_INFINITY;
|
||||
const clamped = Math.min(v, maxCtx);
|
||||
setCustomContextLength(
|
||||
clamped === (ggufContextLength ?? 0)
|
||||
? null
|
||||
: clamped,
|
||||
);
|
||||
}
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
<Slider
|
||||
min={1024}
|
||||
max={ctxMaxValue ?? 4096}
|
||||
step={1024}
|
||||
value={[
|
||||
Math.min(
|
||||
typeof ctxDisplayValue === "number"
|
||||
? ctxDisplayValue
|
||||
: (ggufContextLength ?? 4096),
|
||||
ctxMaxValue ?? 4096,
|
||||
),
|
||||
]}
|
||||
onValueChange={([v]) => {
|
||||
setCustomContextLength(
|
||||
v === (ggufContextLength ?? 0) ? null : v,
|
||||
);
|
||||
}}
|
||||
/>
|
||||
{ggufMaxContextLength != null &&
|
||||
typeof ctxDisplayValue === "number" &&
|
||||
ctxDisplayValue > ggufMaxContextLength && (
|
||||
<p className="text-[11px] text-amber-500">
|
||||
Exceeds estimated VRAM capacity (
|
||||
{ggufMaxContextLength.toLocaleString()} tokens). The
|
||||
model may use system RAM.
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
<div className="grid grid-cols-[minmax(0,1fr)_65px] items-center gap-x-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">KV Cache Dtype</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Quantize KV cache to reduce VRAM.
|
||||
</div>
|
||||
</div>
|
||||
<div className="w-full min-w-0">
|
||||
<Select
|
||||
value={kvCacheDtype ?? "f16"}
|
||||
onValueChange={(v) => {
|
||||
setKvCacheDtype(v === "f16" ? null : v);
|
||||
}}
|
||||
>
|
||||
<SelectTrigger className="grid h-7 w-full min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 px-2 py-0 text-xs [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="f16">f16</SelectItem>
|
||||
<SelectItem value="bf16">bf16</SelectItem>
|
||||
<SelectItem value="q8_0">q8_0</SelectItem>
|
||||
<SelectItem value="q5_1">q5_1</SelectItem>
|
||||
<SelectItem value="q4_1">q4_1</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</div>
|
||||
{!currentModelIsVision && (
|
||||
<div className="grid grid-cols-[minmax(0,1fr)_65px] items-center gap-x-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">
|
||||
Speculative Decoding
|
||||
</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Speed up generation with no VRAM cost.
|
||||
</div>
|
||||
</div>
|
||||
<div className="w-full min-w-0">
|
||||
<Select
|
||||
value={speculativeType ?? "off"}
|
||||
onValueChange={(v) => {
|
||||
setSpeculativeType(v === "off" ? null : v);
|
||||
}}
|
||||
>
|
||||
<SelectTrigger className="grid h-7 w-full min-w-0 grid-cols-[minmax(0,1fr)_auto] items-center gap-1 px-2 py-0 text-xs [&_[data-slot=select-value]]:min-w-0 [&_[data-slot=select-value]]:truncate [&>svg]:shrink-0">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="ngram-mod">On</SelectItem>
|
||||
<SelectItem value="off">Off</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
{modelSettingsDirty && (
|
||||
<div className="flex flex-wrap gap-1.5 pt-1">
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => onReloadModel?.()}
|
||||
className="rounded-md bg-primary px-2.5 py-1 text-[11px] font-medium text-primary-foreground transition-colors hover:bg-primary/90"
|
||||
>
|
||||
Apply
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setCustomContextLength(null);
|
||||
setKvCacheDtype(loadedKvCacheDtype);
|
||||
setSpeculativeType(loadedSpeculativeType);
|
||||
}}
|
||||
className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
|
||||
>
|
||||
Reset
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
{!isGguf && params.checkpoint && (
|
||||
<>
|
||||
<div className="flex items-center justify-between gap-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">Enable custom code</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Allow models with custom code (e.g. Nemotron). Only
|
||||
enable if sure.
|
||||
</div>
|
||||
</div>
|
||||
<Switch
|
||||
checked={params.trustRemoteCode ?? false}
|
||||
onCheckedChange={set("trustRemoteCode")}
|
||||
/>
|
||||
</div>
|
||||
{trustRemoteCodeMissing && (
|
||||
<Alert className="border-amber-200/70 bg-amber-50/70 px-3 py-2 text-amber-950 dark:border-amber-900/70 dark:bg-amber-950/35 dark:text-amber-100">
|
||||
<AlertTitle className="text-[11px] font-medium">
|
||||
Keep custom code enabled for this model
|
||||
</AlertTitle>
|
||||
<AlertDescription className="text-[11px] text-amber-800 dark:text-amber-200">
|
||||
This model requires custom code to load. You can edit the
|
||||
toggle, but loading will stay blocked until it is turned
|
||||
back on.
|
||||
</AlertDescription>
|
||||
</Alert>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</CollapsibleSection>
|
||||
|
||||
<CollapsibleSection
|
||||
icon={SlidersHorizontalIcon}
|
||||
label="Sampling"
|
||||
|
|
@ -1186,6 +1102,8 @@ export function ChatSettingsPanel({
|
|||
</div>
|
||||
</CollapsibleSection>
|
||||
|
||||
{modelSection}
|
||||
|
||||
<CollapsibleSection icon={Wrench01Icon} label="Tools">
|
||||
<div className="flex flex-col gap-3 py-1">
|
||||
<AutoHealToolCallsToggle />
|
||||
|
|
|
|||
|
|
@ -3,6 +3,13 @@
|
|||
|
||||
import { createElement, useCallback, useRef, useState } from "react";
|
||||
import { toast } from "sonner";
|
||||
import { consumeNativePathToken } from "@/features/native-intents/api";
|
||||
import {
|
||||
notifyNative,
|
||||
primeNativeNotificationPermission,
|
||||
safeNotificationLabel,
|
||||
sanitizeNotificationBody,
|
||||
} from "@/lib/native-notifications";
|
||||
import { ModelLoadDescription } from "../components/model-load-status";
|
||||
import {
|
||||
getDownloadProgress,
|
||||
|
|
@ -17,13 +24,28 @@ import {
|
|||
} from "../api/chat-api";
|
||||
import { formatEta, formatRate } from "../utils/format-transfer";
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
import type { InferenceStatusResponse, LoadModelResponse } from "../types/api";
|
||||
import {
|
||||
mergeBackendRecommendedInference,
|
||||
resolveLoadMaxSeqLength,
|
||||
} from "../presets/preset-policy";
|
||||
import type {
|
||||
ChatLoraSummary,
|
||||
ChatModelSummary,
|
||||
InferenceParams,
|
||||
} from "../types/runtime";
|
||||
|
||||
// The simplified Speculative Decoding control surfaces "default" (which
|
||||
// maps to llama.cpp's --spec-default) and "off". A backend status / load
|
||||
// response can still report the older manual modes (ngram-mod,
|
||||
// ngram-simple) when a model is loaded via the API or carried over from an
|
||||
// older Studio version. The Select would render an empty trigger for those
|
||||
// values, so coerce them to "default" -- llama.cpp's own --spec-default
|
||||
// picks an equivalent strategy and keeps the dropdown coherent.
|
||||
function normalizeSpeculativeType(v: string | null | undefined): string | null {
|
||||
if (v == null) return null;
|
||||
if (v === "default" || v === "off") return v;
|
||||
return "default";
|
||||
}
|
||||
|
||||
type SelectedModelInput = {
|
||||
id: string;
|
||||
isLora?: boolean;
|
||||
|
|
@ -32,6 +54,8 @@ type SelectedModelInput = {
|
|||
isDownloaded?: boolean;
|
||||
expectedBytes?: number;
|
||||
forceReload?: boolean;
|
||||
nativePathToken?: string;
|
||||
throwOnError?: boolean;
|
||||
};
|
||||
|
||||
const MODEL_LOAD_TOAST_CLASSNAMES = {
|
||||
|
|
@ -119,46 +143,10 @@ function toLoraSummary(lora: {
|
|||
};
|
||||
}
|
||||
|
||||
function toFiniteNumber(value: unknown): number | undefined {
|
||||
if (typeof value !== "number" || !Number.isFinite(value)) {
|
||||
return undefined;
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
function getTrustRemoteCodeRequiredMessage(modelName: string): string {
|
||||
return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`;
|
||||
}
|
||||
|
||||
function mergeRecommendedInference(
|
||||
current: InferenceParams,
|
||||
response: LoadModelResponse | InferenceStatusResponse,
|
||||
modelId: string,
|
||||
): InferenceParams {
|
||||
const inference = response.inference;
|
||||
// GGUF: use actual context length from GGUF metadata, fallback to 131072
|
||||
// Non-GGUF: 4096
|
||||
const defaultMaxTokens = response.is_gguf
|
||||
? (response.context_length ?? 131072)
|
||||
: 4096;
|
||||
return {
|
||||
...current,
|
||||
checkpoint: modelId,
|
||||
maxTokens: defaultMaxTokens,
|
||||
temperature:
|
||||
toFiniteNumber(inference?.temperature) ?? current.temperature,
|
||||
topP: toFiniteNumber(inference?.top_p) ?? current.topP,
|
||||
topK: toFiniteNumber(inference?.top_k) ?? current.topK,
|
||||
minP: toFiniteNumber(inference?.min_p) ?? current.minP,
|
||||
presencePenalty:
|
||||
toFiniteNumber(inference?.presence_penalty) ?? current.presencePenalty,
|
||||
trustRemoteCode:
|
||||
typeof inference?.trust_remote_code === "boolean"
|
||||
? inference.trust_remote_code
|
||||
: current.trustRemoteCode,
|
||||
};
|
||||
}
|
||||
|
||||
export function useChatModelRuntime() {
|
||||
const params = useChatRuntimeStore((state) => state.params);
|
||||
const models = useChatRuntimeStore((state) => state.models);
|
||||
|
|
@ -175,6 +163,7 @@ export function useChatModelRuntime() {
|
|||
displayName: string;
|
||||
isDownloaded?: boolean;
|
||||
isCachedLora?: boolean;
|
||||
nativePathToken?: string | null;
|
||||
} | null>(null);
|
||||
const [loadToastDismissed, setLoadToastDismissed] = useState(false);
|
||||
const [loadProgress, setLoadProgress] = useState<{
|
||||
|
|
@ -185,7 +174,9 @@ export function useChatModelRuntime() {
|
|||
const loadAbortRef = useRef<AbortController | null>(null);
|
||||
const loadingModelRef = useRef<typeof loadingModel>(null);
|
||||
const loadToastIdRef = useRef<string | number | null>(null);
|
||||
const loadAttemptRef = useRef(0);
|
||||
const loadToastDismissedRef = useRef(false);
|
||||
const cancelUnloadPendingRef = useRef(false);
|
||||
|
||||
const setLoadToastDismissedState = useCallback((dismissed: boolean) => {
|
||||
loadToastDismissedRef.current = dismissed;
|
||||
|
|
@ -199,7 +190,9 @@ export function useChatModelRuntime() {
|
|||
loadAbortRef.current = null;
|
||||
loadToastIdRef.current = null;
|
||||
setLoadToastDismissedState(false);
|
||||
useChatRuntimeStore.getState().setModelLoading(false);
|
||||
if (!cancelUnloadPendingRef.current) {
|
||||
useChatRuntimeStore.getState().setModelLoading(false);
|
||||
}
|
||||
}, [setLoadToastDismissedState]);
|
||||
|
||||
const renderLoadDescription = useCallback(
|
||||
|
|
@ -238,29 +231,13 @@ export function useChatModelRuntime() {
|
|||
// Apply inference defaults on reconnect (page refresh with model already loaded)
|
||||
if (statusRes.inference) {
|
||||
const currentParams = useChatRuntimeStore.getState().params;
|
||||
const reconnectResponse: LoadModelResponse = {
|
||||
status: "already_loaded",
|
||||
model: statusRes.active_model,
|
||||
display_name: statusRes.active_model,
|
||||
is_vision: statusRes.is_vision,
|
||||
is_lora: false,
|
||||
is_gguf: statusRes.is_gguf,
|
||||
is_audio: statusRes.is_audio,
|
||||
audio_type: statusRes.audio_type,
|
||||
has_audio_input: statusRes.has_audio_input,
|
||||
inference: statusRes.inference,
|
||||
context_length: statusRes.context_length,
|
||||
max_context_length: statusRes.max_context_length,
|
||||
native_context_length: statusRes.native_context_length,
|
||||
supports_reasoning: statusRes.supports_reasoning,
|
||||
reasoning_style: statusRes.reasoning_style,
|
||||
reasoning_always_on: statusRes.reasoning_always_on,
|
||||
supports_preserve_thinking: statusRes.supports_preserve_thinking,
|
||||
supports_tools: statusRes.supports_tools,
|
||||
speculative_type: statusRes.speculative_type,
|
||||
};
|
||||
setParams(
|
||||
mergeRecommendedInference(currentParams, statusRes, statusRes.active_model),
|
||||
mergeBackendRecommendedInference({
|
||||
current: currentParams,
|
||||
response: statusRes,
|
||||
modelId: statusRes.active_model,
|
||||
presetSource: useChatRuntimeStore.getState().activePresetSource,
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -279,7 +256,11 @@ export function useChatModelRuntime() {
|
|||
const ggufNativeContextLength = statusRes.is_gguf
|
||||
? (statusRes.native_context_length ?? null)
|
||||
: null;
|
||||
const currentSpecType = statusRes.speculative_type ?? null;
|
||||
const currentSpecType = normalizeSpeculativeType(statusRes.speculative_type);
|
||||
const nextDefaultChatTemplate =
|
||||
statusRes.chat_template === undefined
|
||||
? useChatRuntimeStore.getState().defaultChatTemplate
|
||||
: statusRes.chat_template;
|
||||
useChatRuntimeStore.setState({
|
||||
supportsReasoning,
|
||||
reasoningAlwaysOn,
|
||||
|
|
@ -296,6 +277,7 @@ export function useChatModelRuntime() {
|
|||
ggufNativeContextLength,
|
||||
modelRequiresTrustRemoteCode:
|
||||
statusRes.requires_trust_remote_code ?? false,
|
||||
defaultChatTemplate: nextDefaultChatTemplate,
|
||||
speculativeType: currentSpecType,
|
||||
loadedSpeculativeType: currentSpecType,
|
||||
});
|
||||
|
|
@ -346,8 +328,18 @@ export function useChatModelRuntime() {
|
|||
? undefined
|
||||
: "The current download may still finish in the background.",
|
||||
});
|
||||
// Fire-and-forget: tell backend to stop, don't block UI
|
||||
unloadModel({ model_path: model.id }).catch(() => {});
|
||||
cancelUnloadPendingRef.current = true;
|
||||
useChatRuntimeStore.getState().setModelLoading(true);
|
||||
void (async () => {
|
||||
try {
|
||||
await unloadModel({ model_path: model.id }).catch(() => {});
|
||||
} finally {
|
||||
cancelUnloadPendingRef.current = false;
|
||||
if (!loadingModelRef.current) {
|
||||
useChatRuntimeStore.getState().setModelLoading(false);
|
||||
}
|
||||
}
|
||||
})();
|
||||
}, [clearCheckpoint, setLoadToastDismissedState]);
|
||||
|
||||
const selectModel = useCallback(
|
||||
|
|
@ -357,12 +349,20 @@ export function useChatModelRuntime() {
|
|||
typeof selection === "string" ? undefined : selection.ggufVariant;
|
||||
const forceReload =
|
||||
typeof selection === "string" ? false : selection.forceReload ?? false;
|
||||
const nativePathToken =
|
||||
typeof selection === "string" ? undefined : selection.nativePathToken;
|
||||
const throwOnError =
|
||||
typeof selection === "string" ? false : selection.throwOnError ?? false;
|
||||
const currentVariant = useChatRuntimeStore.getState().activeGgufVariant;
|
||||
if (!forceReload && (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null)))) {
|
||||
return;
|
||||
}
|
||||
// Prevent duplicate loads if already loading this model
|
||||
if (loadingModelRef.current?.id === modelId) return;
|
||||
if (
|
||||
loadingModelRef.current?.id === modelId &&
|
||||
(loadingModelRef.current?.nativePathToken ?? null) === (nativePathToken ?? null)
|
||||
)
|
||||
return;
|
||||
|
||||
const explicitIsLora =
|
||||
typeof selection === "string" ? undefined : selection.isLora;
|
||||
|
|
@ -376,6 +376,10 @@ export function useChatModelRuntime() {
|
|||
const isLora =
|
||||
explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false;
|
||||
const displayName = model?.name || lora?.name || modelId;
|
||||
const loadAttemptId = ++loadAttemptRef.current;
|
||||
primeNativeNotificationPermission().catch(() => undefined);
|
||||
const notificationModelKey = `${modelId}:${ggufVariant ?? ""}:${loadAttemptId}`;
|
||||
const safeModelName = safeNotificationLabel(displayName, "The model");
|
||||
const currentCheckpoint =
|
||||
useChatRuntimeStore.getState().params.checkpoint;
|
||||
const previousCheckpoint = currentCheckpoint;
|
||||
|
|
@ -402,7 +406,13 @@ export function useChatModelRuntime() {
|
|||
.join(" ");
|
||||
setModelsError(null);
|
||||
setLoadToastDismissedState(false);
|
||||
const loadInfo = { id: modelId, displayName, isDownloaded, isCachedLora };
|
||||
const loadInfo = {
|
||||
id: modelId,
|
||||
displayName,
|
||||
isDownloaded,
|
||||
isCachedLora,
|
||||
nativePathToken: nativePathToken ?? null,
|
||||
};
|
||||
setLoadingModel(loadInfo);
|
||||
useChatRuntimeStore.getState().setModelLoading(true);
|
||||
setLoadProgress(
|
||||
|
|
@ -419,17 +429,30 @@ export function useChatModelRuntime() {
|
|||
let previousWasUnloaded = false;
|
||||
const currentCheckpoint =
|
||||
useChatRuntimeStore.getState().params.checkpoint;
|
||||
const paramsBeforeLoad = useChatRuntimeStore.getState().params;
|
||||
const trustRemoteCode = paramsBeforeLoad.trustRemoteCode ?? false;
|
||||
const maxSeqLength = paramsBeforeLoad.maxSeqLength;
|
||||
const hfToken = useChatRuntimeStore.getState().hfToken || null;
|
||||
const stateBeforeUnload = useChatRuntimeStore.getState();
|
||||
const trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false;
|
||||
const maxSeqLength = stateBeforeUnload.params.maxSeqLength;
|
||||
const previousIsGguf =
|
||||
previousModel?.isGguf === true
|
||||
|| previousVariant != null
|
||||
|| (previousCheckpoint?.toLowerCase().endsWith(".gguf") ?? false);
|
||||
const rollbackMaxSeqLength = previousIsGguf
|
||||
? (stateBeforeUnload.ggufContextLength ?? 0)
|
||||
: maxSeqLength;
|
||||
const hfToken = stateBeforeUnload.hfToken || null;
|
||||
const previousModelRequiresTrustRemoteCode =
|
||||
useChatRuntimeStore.getState().modelRequiresTrustRemoteCode;
|
||||
stateBeforeUnload.modelRequiresTrustRemoteCode;
|
||||
const previousActiveNativePathToken =
|
||||
stateBeforeUnload.activeNativePathToken;
|
||||
try {
|
||||
// Lightweight pre-flight validation: avoid unloading a working model
|
||||
// if the new identifier is clearly invalid (e.g. bad HF id / path).
|
||||
const validateNativePathLease = nativePathToken
|
||||
? (await consumeNativePathToken(nativePathToken, "validate-model")).nativePathLease
|
||||
: undefined;
|
||||
const validation = await validateModel({
|
||||
model_path: modelId,
|
||||
nativePathLease: validateNativePathLease,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: maxSeqLength,
|
||||
load_in_4bit: true,
|
||||
|
|
@ -439,21 +462,39 @@ export function useChatModelRuntime() {
|
|||
if (validation.requires_trust_remote_code && !trustRemoteCode) {
|
||||
throw new Error(getTrustRemoteCodeRequiredMessage(displayName));
|
||||
}
|
||||
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
|
||||
const loadNativePathLease = nativePathToken
|
||||
? (await consumeNativePathToken(nativePathToken, "load-model")).nativePathLease
|
||||
: undefined;
|
||||
|
||||
if (currentCheckpoint) {
|
||||
await unloadModel({ model_path: currentCheckpoint });
|
||||
previousWasUnloaded = true;
|
||||
}
|
||||
if (abortCtrl.signal.aborted) throw new Error("Cancelled");
|
||||
|
||||
const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength, speculativeType } = useChatRuntimeStore.getState();
|
||||
// GGUF: use custom context length, or 0 = model's native context
|
||||
// Non-GGUF: use the Max Seq Length slider value
|
||||
const isDirectGgufFile = modelId.toLowerCase().endsWith(".gguf");
|
||||
const effectiveMaxSeqLength = customContextLength != null
|
||||
? customContextLength
|
||||
: (ggufVariant != null || isDirectGgufFile) ? (ggufContextLength ?? 0) : maxSeqLength;
|
||||
const {
|
||||
chatTemplateOverride,
|
||||
kvCacheDtype,
|
||||
customContextLength,
|
||||
ggufContextLength,
|
||||
speculativeType,
|
||||
activePresetSource,
|
||||
activeGgufVariant,
|
||||
} = useChatRuntimeStore.getState();
|
||||
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
|
||||
modelId,
|
||||
ggufVariant,
|
||||
customContextLength,
|
||||
ggufContextLength,
|
||||
currentCheckpoint,
|
||||
activeGgufVariant,
|
||||
maxSeqLength,
|
||||
presetSource: activePresetSource,
|
||||
});
|
||||
const loadResponse = await loadModel({
|
||||
model_path: modelId,
|
||||
nativePathLease: loadNativePathLease,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: effectiveMaxSeqLength,
|
||||
load_in_4bit: true,
|
||||
|
|
@ -471,7 +512,12 @@ export function useChatModelRuntime() {
|
|||
|
||||
const currentParams = useChatRuntimeStore.getState().params;
|
||||
setParams(
|
||||
mergeRecommendedInference(currentParams, loadResponse, modelId),
|
||||
mergeBackendRecommendedInference({
|
||||
current: currentParams,
|
||||
response: loadResponse,
|
||||
modelId,
|
||||
presetSource: useChatRuntimeStore.getState().activePresetSource,
|
||||
}),
|
||||
);
|
||||
// Qwen3.5/3.6 small models (0.8B, 2B, 4B, 9B) disable thinking by default
|
||||
let reasoningDefault = loadResponse.supports_reasoning ?? false;
|
||||
|
|
@ -485,7 +531,7 @@ export function useChatModelRuntime() {
|
|||
}
|
||||
}
|
||||
const loadedKv = loadResponse.cache_type_kv ?? null;
|
||||
const loadedSpec = loadResponse.speculative_type ?? null;
|
||||
const loadedSpec = normalizeSpeculativeType(loadResponse.speculative_type);
|
||||
const nativeCtx = loadResponse.is_gguf
|
||||
? (loadResponse.context_length ?? 131072)
|
||||
: null;
|
||||
|
|
@ -521,16 +567,36 @@ export function useChatModelRuntime() {
|
|||
customContextLength: keepCustomCtx,
|
||||
defaultChatTemplate: loadResponse.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
activeNativePathToken: nativePathToken ?? null,
|
||||
});
|
||||
// Qwen3/3.5/3.6: apply thinking-mode-specific params after load
|
||||
if (modelId.toLowerCase().includes("qwen3") && (loadResponse.supports_reasoning ?? false)) {
|
||||
const store = useChatRuntimeStore.getState();
|
||||
const mid = modelId.toLowerCase();
|
||||
const needsPresencePenalty = mid.includes("qwen3.5") || mid.includes("qwen3.6");
|
||||
const p = reasoningDefault
|
||||
? { temperature: 0.6, topP: 0.95, topK: 20, minP: 0.0, ...(needsPresencePenalty ? { presencePenalty: 1.5 } : {}) }
|
||||
: { temperature: 0.7, topP: 0.8, topK: 20, minP: 0.0, ...(needsPresencePenalty ? { presencePenalty: 1.5 } : {}) };
|
||||
store.setParams({ ...store.params, ...p });
|
||||
if (store.activePresetSource === "builtin-default") {
|
||||
const mid = modelId.toLowerCase();
|
||||
const needsPresencePenalty =
|
||||
mid.includes("qwen3.5") || mid.includes("qwen3.6");
|
||||
const p = reasoningDefault
|
||||
? {
|
||||
temperature: 0.6,
|
||||
topP: 0.95,
|
||||
topK: 20,
|
||||
minP: 0.0,
|
||||
...(needsPresencePenalty
|
||||
? { presencePenalty: 1.5 }
|
||||
: {}),
|
||||
}
|
||||
: {
|
||||
temperature: 0.7,
|
||||
topP: 0.8,
|
||||
topK: 20,
|
||||
minP: 0.0,
|
||||
...(needsPresencePenalty
|
||||
? { presencePenalty: 1.5 }
|
||||
: {}),
|
||||
};
|
||||
store.setParams({ ...store.params, ...p });
|
||||
}
|
||||
}
|
||||
await refresh();
|
||||
} catch (error) {
|
||||
|
|
@ -538,20 +604,36 @@ export function useChatModelRuntime() {
|
|||
if (abortCtrl.signal.aborted) throw error;
|
||||
// If we unloaded a previous model and the new load failed, attempt a rollback.
|
||||
if (previousWasUnloaded && previousCheckpoint) {
|
||||
let rollbackNativePathLease: string | undefined;
|
||||
if (previousActiveNativePathToken) {
|
||||
try {
|
||||
rollbackNativePathLease = (
|
||||
await consumeNativePathToken(previousActiveNativePathToken, "load-model")
|
||||
).nativePathLease;
|
||||
} catch {
|
||||
throw new Error(
|
||||
"Could not reload the previous local model: please re-select the file.",
|
||||
);
|
||||
}
|
||||
}
|
||||
try {
|
||||
await loadModel({
|
||||
model_path: previousCheckpoint,
|
||||
nativePathLease: rollbackNativePathLease,
|
||||
hf_token: hfToken,
|
||||
max_seq_length: maxSeqLength,
|
||||
max_seq_length: rollbackMaxSeqLength,
|
||||
load_in_4bit: true,
|
||||
is_lora: previousIsLora,
|
||||
gguf_variant: previousVariant,
|
||||
trust_remote_code:
|
||||
previousModelRequiresTrustRemoteCode || trustRemoteCode,
|
||||
});
|
||||
useChatRuntimeStore.setState({
|
||||
activeNativePathToken: previousActiveNativePathToken ?? null,
|
||||
});
|
||||
await refresh();
|
||||
} catch {
|
||||
// If rollback also fails, surface the original error.
|
||||
// Rollback also failed; surface the original load error below.
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
|
|
@ -742,6 +824,12 @@ export function useChatModelRuntime() {
|
|||
},
|
||||
});
|
||||
}
|
||||
notifyNative({
|
||||
key: `model-downloaded:${notificationModelKey}`,
|
||||
title: "Model downloaded",
|
||||
body: `${safeModelName} finished downloading and is loading into memory.`,
|
||||
requestPermission: false,
|
||||
}).catch(() => undefined);
|
||||
// Keep polling: the mmap branch below takes over from here.
|
||||
}
|
||||
} catch {
|
||||
|
|
@ -827,6 +915,12 @@ export function useChatModelRuntime() {
|
|||
duration: 2000,
|
||||
});
|
||||
}
|
||||
notifyNative({
|
||||
key: `model-loaded:${notificationModelKey}`,
|
||||
title: "Model ready",
|
||||
body: `${safeModelName} is loaded and ready to chat.`,
|
||||
requestPermission: false,
|
||||
}).catch(() => undefined);
|
||||
} catch (err) {
|
||||
if (!abortCtrl.signal.aborted) {
|
||||
const message =
|
||||
|
|
@ -841,6 +935,12 @@ export function useChatModelRuntime() {
|
|||
duration: 5000,
|
||||
});
|
||||
}
|
||||
notifyNative({
|
||||
key: `model-load-failed:${notificationModelKey}`,
|
||||
title: "Model failed to load",
|
||||
body: sanitizeNotificationBody(message, "The model failed to load."),
|
||||
requestPermission: false,
|
||||
}).catch(() => undefined);
|
||||
}
|
||||
throw err;
|
||||
} finally {
|
||||
|
|
@ -853,6 +953,9 @@ export function useChatModelRuntime() {
|
|||
const message =
|
||||
error instanceof Error ? error.message : "Failed to load model";
|
||||
setModelsError(message);
|
||||
if (throwOnError) {
|
||||
throw error instanceof Error ? error : new Error(message);
|
||||
}
|
||||
}
|
||||
},
|
||||
[
|
||||
|
|
|
|||
|
|
@ -105,6 +105,9 @@ async function buildIndex(): Promise<ChatSearchItem[]> {
|
|||
const arr = byThreadId.get(tid);
|
||||
if (arr) merged.push(...arr);
|
||||
}
|
||||
if (merged.length === 0) {
|
||||
continue;
|
||||
}
|
||||
merged.sort((a, b) => b.createdAt - a.createdAt);
|
||||
|
||||
let preview = "";
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue