diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 3c7cea919c..9b85b20177 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -6,12 +6,12 @@ # Local Agent Guides CI. All failures from here are failure class (c) # "guide drift": the server preflight already passed and the agent CLI # already installed, so a failure here means the documented recipe in -# unsloth_cli/commands/connect.py no longer produces a working flow. +# unsloth_cli/commands/start.py no longer produces a working flow. # -# Self-updating: for the 5 agents with a connect.py recipe we obtain the -# exact env + command from `unsloth connect --no-launch` and run -# THAT, so a recipe change is exercised automatically. Pi (no connect.py -# command at HEAD) is driven by a hand-written recipe. +# Self-updating: for all six agents (claude, codex, hermes, openclaw, +# opencode, pi) we obtain the exact env + command from +# `unsloth start --no-launch` and run THAT, so a recipe change is +# exercised automatically. # # Every agent invocation is wrapped in `timeout` so a headless-TTY prompt # can never hang the runner -- a timeout is reported as guide drift with a @@ -53,14 +53,14 @@ REDACTED_DIR="$REPO_ROOT/redacted-configs" WORKDIR_BASE="$REPO_ROOT/agent-workdir" CACHE_HELPER="$SCRIPT_DIR/assert-prompt-cache.sh" mkdir -p "$LOGS_DIR" "$REDACTED_DIR" -CONNECT_REF="unsloth_cli/commands/connect.py" +CONNECT_REF="unsloth_cli/commands/start.py" # Prefill-shrinking flags for Claude Code. The heavyweight agents send # multi-thousand-token system prompts + full tool schemas, which on a CPU-only # runner is minutes of prefill per model round-trip (~16 tok/s for a 4B model). # Replacing the ~5.7k default system prompt with a tiny one (--system-prompt-file) # and restricting tools cuts the prefill to a few hundred tokens so it completes -# quickly on CPU. These only shape the request size; the connect.py recipe +# quickly on CPU. These only shape the request size; the start.py recipe # (endpoint, auth, model) is still exercised end to end. # # The bulk of Claude Code's prompt is the built-in tool JSON schemas: measured @@ -105,6 +105,13 @@ redact() { done } +# Print a file to the log with the key scrubbed, without mutating it (the raw file is +# still needed to parse the real env). Use this instead of `cat` for any transcript that +# carries an `export UNSLOTH_API_KEY=...` line, so a live key never reaches Actions logs. +cat_redacted() { + sed "s#${UNSLOTH_API_KEY}##g" "$1" +} + # A reply must be non-empty and free of connection/auth errors. assert_reply() { local out="$1" @@ -131,45 +138,29 @@ run_timed() { # $1=outfile, rest=command return "$rc" } -# ── Pi: no connect.py command at HEAD -> hand-written recipe ────────────── -write_pi_config() { - if unsloth connect pi --help >/dev/null 2>&1; then - # Tripwire: once a real recipe exists, the hand-written config would mask any - # drift in it, defeating the point of this CI. Fail hard so the cell is - # migrated to the self-updating `unsloth connect pi --no-launch` path. - guide_fail "connect.py now ships a 'pi' command -- migrate this CI cell to the 'unsloth connect pi --no-launch' path so the documented recipe is exercised (the hand-written Pi config no longer reflects it)" - fi - mkdir -p "$HOME/.pi/agent" - python3 - "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" "$UNSLOTH_MODEL_ID" <<'PY' -import json, os, sys -base, key, model = sys.argv[1], sys.argv[2], sys.argv[3] -cfg = {"providers": {"unsloth": { - "api": "openai-completions", - "baseUrl": f"{base}/v1", - "apiKey": key, - "models": [{"id": model}], -}}} -path = os.path.expanduser("~/.pi/agent/models.json") -with open(path, "w") as fh: - json.dump(cfg, fh, indent=2) -PY - cp "$HOME/.pi/agent/models.json" "$REDACTED_DIR/pi-models.json" 2>/dev/null || true - redact "$REDACTED_DIR/pi-models.json" +# Read a value from an `export VAR=...` line in the connect --no-launch output. +# `unsloth start` writes each agent's session config off the user's ~ and points +# at it through a relocation env var (CODEX_HOME / OPENCODE_CONFIG / +# OPENCLAW_CONFIG_PATH), so the contract checks read the path from here. +raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) + local raw="$LOGS_DIR/connect-${AGENT}.txt" + local v; v="$(sed -n "s/^export $1=//p" "$raw" | tail -1)" + v="${v#\'}"; v="${v%\'}"; printf '%s' "$v" } -# ── 5-agent connect.py path: parse env + command from --no-launch ───────── +# ── 5-agent start.py path: parse env + command from --no-launch ───────── # Populates globals CONNECT_ENV (export/unset lines) and CONNECT_CMD (the -# launch command on the last printed line), and runs connect.py's config -# writers as a side effect (it writes ~/.codex, ~/.claude, etc.). +# launch command on the last printed line), and runs start.py's config +# writers as a side effect (it writes each agent's relocated session config). parse_connect() { local raw="$LOGS_DIR/connect-${AGENT}.txt" - if ! unsloth connect "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then - cat "$raw" - guide_fail "'unsloth connect ${AGENT} --no-launch' exited non-zero" + if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + cat_redacted "$raw" + guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" fi - echo "[$AGENT] connect --no-launch printed:"; cat "$raw" + echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw" CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)" - # The launch command is the last non-export, non-status line. connect.py + # The launch command is the last non-export, non-status line. start.py # prints "Studio · model " and "Updated ..." status lines first. CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \ | grep -E '[^[:space:]]' | tail -1)" @@ -177,45 +168,63 @@ parse_connect() { redact "$raw" } -# Cross-check the documented contract knobs so silent connect.py changes +# Cross-check the documented contract knobs so silent start.py changes # (env-var rename, wire_api flip, attribution setting drop) also fail/flag. crosscheck_contract() { local raw="$LOGS_DIR/connect-${AGENT}.txt" + local cfg home case "$AGENT" in codex) grep -q 'UNSLOTH_STUDIO_AUTH_TOKEN' "$raw" \ - || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (connect.py _CODEX_ENV_KEY)" - if [ -f "$HOME/.codex/config.toml" ]; then - grep -q 'wire_api = "responses"' "$HOME/.codex/config.toml" \ - || guide_fail "Codex wire_api is no longer \"responses\" in ~/.codex/config.toml" - cp "$HOME/.codex/config.toml" "$REDACTED_DIR/codex-config.toml" + || guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (start.py _CODEX_ENV_KEY)" + home="$(raw_env CODEX_HOME)" + # An empty relocation var would make cfg "/config.toml" and silently + # skip the [ -f ] contract check below; fail loudly instead. + [ -n "$home" ] || guide_fail "CODEX_HOME missing from connect output (start.py codex())" + cfg="$home/config.toml" + if [ -f "$cfg" ]; then + grep -q 'wire_api = "responses"' "$cfg" \ + || guide_fail "Codex wire_api is no longer \"responses\" in \$CODEX_HOME/config.toml" + cp "$cfg" "$REDACTED_DIR/codex-config.toml" fi grep -q 'codex --oss --profile unsloth_api' "$raw" \ || echo "::warning::Codex launch command changed from 'codex --oss --profile unsloth_api'" ;; claude) grep -q 'ANTHROPIC_AUTH_TOKEN' "$raw" \ - || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (connect.py claude())" - if [ -f "$HOME/.claude/settings.json" ]; then - grep -q '"CLAUDE_CODE_ATTRIBUTION_HEADER"' "$HOME/.claude/settings.json" \ - || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER not written to ~/.claude/settings.json (ensure_claude_attribution_header)" - cp "$HOME/.claude/settings.json" "$REDACTED_DIR/claude-settings.json" - fi + || guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (start.py claude())" + grep -q 'CLAUDE_CODE_ATTRIBUTION_HEADER' "$raw" \ + || echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER no longer set for the session (start.py claude())" ;; hermes) grep -q 'UNSLOTH_API_KEY' "$raw" \ - || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (connect.py _HERMES_ENV_KEY)" - [ -f "$HOME/.hermes/config.yaml" ] && cp "$HOME/.hermes/config.yaml" "$REDACTED_DIR/hermes-config.yaml" + || guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (start.py _HERMES_ENV_KEY)" + home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "HERMES_HOME missing from connect output (start.py hermes())" + cfg="$home/config.yaml" + [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/hermes-config.yaml" ;; openclaw) - if [ -f "$HOME/.openclaw/openclaw.json" ]; then - grep -q '"openai-completions"' "$HOME/.openclaw/openclaw.json" \ + cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + if [ -n "$cfg" ] && [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ || echo "::warning::OpenClaw provider api is no longer 'openai-completions' (write_openclaw_config)" - cp "$HOME/.openclaw/openclaw.json" "$REDACTED_DIR/openclaw.json" + cp "$cfg" "$REDACTED_DIR/openclaw.json" fi ;; opencode) - [ -f "$HOME/.config/opencode/opencode.json" ] && cp "$HOME/.config/opencode/opencode.json" "$REDACTED_DIR/opencode.json" + cfg="$(raw_env OPENCODE_CONFIG)" + [ -n "$cfg" ] && [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/opencode.json" + ;; + pi) + # Pi has no config-dir env var; the session is HOME-relocated, and the + # provider config lives at $HOME/.pi/agent/models.json. + cfg="$(raw_env HOME)/.pi/agent/models.json" + if [ -f "$cfg" ]; then + grep -q '"openai-completions"' "$cfg" \ + || echo "::warning::Pi provider api is no longer 'openai-completions' (write_pi_config)" + cp "$cfg" "$REDACTED_DIR/pi-models.json" + fi ;; esac redact "$REDACTED_DIR"/* 2>/dev/null || true @@ -229,16 +238,23 @@ crosscheck_contract() { # Hermes: an explicit empty cli toolset disables all tools (and drops the # tool-gated guidance blocks), so -z sends ~300 tokens instead of thousands. -# hermes ships a DEFAULT config.yaml that already has a populated -# platform_toolsets, and `unsloth connect` merges into it, so we must override -# cli (not just append). That needs a YAML parser, and the runner's bare -# python3 has no PyYAML -- but the venv that ships `unsloth` does (connect.py -# imports yaml), so run the patch with that interpreter. +# Hermes enables its default cli toolset when the session config does not pin one, +# so we must set platform_toolsets.cli explicitly to [] (not just append) to get +# zero tools. That needs a YAML parser, and the runner's bare python3 has no +# PyYAML -- but the venv that ships `unsloth` does (start.py imports yaml), so run +# the patch with that interpreter. We patch the relocated $HERMES_HOME/config.yaml +# that `unsloth start` printed, not the user's ~/.hermes. # (-z reads platform_toolsets.cli; --ignore-rules is a no-op under -z.) patch_hermes_tools() { # $1 = none|default + # Check the raw var BEFORE appending /config.yaml: the joined path is never + # empty, so the old guard could not fire and the patcher would die on + # "/config.yaml" with a bare traceback instead of this clear failure. + local home; home="$(raw_env HERMES_HOME)" + [ -n "$home" ] || guide_fail "Hermes HERMES_HOME missing from connect output (start.py hermes())" + local cfg; cfg="$home/config.yaml" # Find a python that can import yaml. The runner's bare python3 cannot, but the # interpreter in the `unsloth` console-script shebang provably can (it runs - # connect.py's write_hermes_config, which imports yaml). Try that first, then + # start.py's write_hermes_config, which imports yaml). Try that first, then # any python on PATH, then the venv sibling, picking the first with PyYAML. local cand py="" shebang shebang="$(head -1 "$(command -v unsloth)" 2>/dev/null | sed -n 's/^#![[:space:]]*//p' | awk '{print $1}')" @@ -247,13 +263,13 @@ patch_hermes_tools() { # $1 = none|default { [ -x "$cand" ] || command -v "$cand" >/dev/null 2>&1; } || continue if "$cand" -c 'import yaml' 2>/dev/null; then py="$cand"; break; fi done - [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch ~/.hermes/config.yaml" - echo "[hermes] patching config with $py" - "$py" - "$1" <<'PY' + [ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch the hermes session config" + echo "[hermes] patching $cfg with $py" + "$py" - "$1" "$cfg" <<'PY' import os, sys import yaml mode = sys.argv[1] -p = os.path.expanduser("~/.hermes/config.yaml") +p = sys.argv[2] cfg = (yaml.safe_load(open(p)) or {}) if os.path.exists(p) else {} ts = cfg.get("platform_toolsets") if not isinstance(ts, dict): @@ -274,10 +290,14 @@ PY # drop the auto-injected AGENTS.md/SOUL.md bootstrap (the bulk of the prompt) for # both modes. --agent must reference a defined agent, so write it before invoking. patch_openclaw_agent() { # $1 = notools|tools - python3 - "$1" <<'PY' + # OpenClaw reads its config from the relocated OPENCLAW_CONFIG_PATH that + # `unsloth start` printed, so patch THAT file (not the user's ~/.openclaw). + local cfg; cfg="$(raw_env OPENCLAW_CONFIG_PATH)" + [ -n "$cfg" ] || guide_fail "OpenClaw OPENCLAW_CONFIG_PATH missing from connect output (start.py openclaw())" + python3 - "$1" "$cfg" <<'PY' import os, sys, json mode = sys.argv[1] -p = os.path.expanduser("~/.openclaw/openclaw.json") +p = sys.argv[2] cfg = json.load(open(p)) if os.path.exists(p) else {} agents = cfg.setdefault("agents", {}) agents.setdefault("defaults", {})["skipBootstrap"] = True @@ -293,20 +313,24 @@ print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}") PY } -# Build an invoke script that applies connect.py's env then runs the launch +# Build an invoke script that applies start.py's env then runs the launch # command (with extra args appended) under bash. We do NOT eval connect's env # into this shell; we write it into a one-shot script so the export/unset -# semantics are exactly what connect.py printed. The script path is absolute +# semantics are exactly what start.py printed. The script path is absolute # so it is valid even when the caller has cd'd into a scratch work dir. invoke_via_connect() { # $1=outfile, rest=extra args appended to the command local out="$1"; shift local script="$LOGS_DIR/invoke-${AGENT}.sh" local real; real="$(mktemp)" + # CONNECT_ENV_EXTRA / CONNECT_CMD_OVERRIDE let a caller (attribution-ab) flip a + # session knob without editing the user's config; empty -> use what start.py emitted. + local cmd="${CONNECT_CMD_OVERRIDE:-$CONNECT_CMD}" { echo "set -uo pipefail" echo "$CONNECT_ENV" + [ -n "${CONNECT_ENV_EXTRA:-}" ] && echo "$CONNECT_ENV_EXTRA" # Append extra args (the prompt / flags) to the launch command verbatim. - printf '%s' "$CONNECT_CMD" + printf '%s' "$cmd" local a for a in "$@"; do printf ' %q' "$a"; done printf '\n' @@ -318,7 +342,9 @@ invoke_via_connect() { # $1=outfile, rest=extra args appended to the command # Writing the redacted copy up front keeps the key out of the artifact even if # the run times out (run_timed exits before returning here). cp "$real" "$script"; redact "$script" - echo "[$AGENT] invoking (timeout ${TIMEOUT}s): $CONNECT_CMD $*" + # The connect one-liner now carries the key as an inline env assignment; scrub it on + # the way to the log (the executed $real keeps the live value). + echo "[$AGENT] invoking (timeout ${TIMEOUT}s): ${cmd//${UNSLOTH_API_KEY}/} $*" run_timed "$out" bash "$real" local rc=$? rm -f "$real" @@ -332,27 +358,23 @@ case "$MODE" in connection) PROMPT='Reply with exactly the single word: pong' OUT="$LOGS_DIR/${AGENT}-connection.txt" - if [ "$AGENT" = "pi" ]; then - write_pi_config - run_timed "$OUT" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$PROMPT" - else - parse_connect - crosscheck_contract - # claude/codex run in print mode via the flags connect.py emits - # (claude -p / codex exec). For agents whose default subcommand prints - # to stdout we pass the prompt through ctx.args. - case "$AGENT" in - claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; - codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; - opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; - hermes) patch_hermes_tools none - invoke_via_connect "$OUT" -z "$PROMPT" ;; - openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ - --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; - *) invoke_via_connect "$OUT" "$PROMPT" ;; - esac - fi + parse_connect + crosscheck_contract + # claude/codex run in print mode via the flags start.py emits + # (claude -p / codex exec). For agents whose default subcommand prints + # to stdout we pass the prompt through ctx.args. + case "$AGENT" in + claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;; + codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;; + opencode) invoke_via_connect "$OUT" run "$PROMPT" ;; + pi) invoke_via_connect "$OUT" -p "$PROMPT" ;; + hermes) patch_hermes_tools none + invoke_via_connect "$OUT" -z "$PROMPT" ;; + openclaw) patch_openclaw_agent notools + invoke_via_connect "$OUT" agent --local --agent ci \ + --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; + *) invoke_via_connect "$OUT" "$PROMPT" ;; + esac # A non-zero exit from the documented launch command is drift even if it # printed something: a benign-looking "command not found" / usage dump would # otherwise slip past assert_reply (which only flags empty/error-keyword text). @@ -371,22 +393,18 @@ case "$MODE" in T1='Create a file named hello.py in the current directory whose entire contents are a single line: print("Hello"). Do not run it.' T2='Run hello.py with python and show me the exact output.' - # The connect.py recipe writers + crosscheck must see the repo; run them + # The start.py recipe writers + crosscheck must see the repo; run them # from the repo root BEFORE cd-ing into the scratch work dir. - if [ "$AGENT" != "pi" ]; then - parse_connect - crosscheck_contract - # File-edit needs real tools, so we cannot zero them as in connection. - # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md - # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work - # dir is empty, so no project context files are auto-loaded either. - case "$AGENT" in - hermes) patch_hermes_tools default ;; - openclaw) patch_openclaw_agent tools ;; - esac - else - write_pi_config - fi + parse_connect + crosscheck_contract + # File-edit needs real tools, so we cannot zero them as in connection. + # hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md + # bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work + # dir is empty, so no project context files are auto-loaded either. + case "$AGENT" in + hermes) patch_hermes_tools default ;; + openclaw) patch_openclaw_agent tools ;; + esac # Drive from inside the work dir so the agent edits files there. All log # writes use absolute $LOGS_DIR, so cwd does not matter for them. @@ -395,7 +413,14 @@ case "$MODE" in invoke_turn() { # $1=outfile $2=continue? $3=prompt local out="$1" cont="$2" prompt="$3" case "$AGENT" in - pi) run_timed "$out" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$prompt" ;; + pi) + # Pi continues the previous session with -c; provider/model come from + # the parsed `unsloth start pi` recipe (CONNECT_CMD), not hardcoded here. + if [ "$cont" = "continue" ]; then + invoke_via_connect "$out" -p --continue "$prompt" + else + invoke_via_connect "$out" -p "$prompt" + fi ;; claude) # --dangerously-skip-permissions lets headless claude actually use the # Write/Bash tools (otherwise it blocks on an approval prompt and emits @@ -466,33 +491,32 @@ case "$MODE" in # right before the measured turn, so an earlier turn's reuse can't leak in. LLAMA_LOG_DIR="${UNSLOTH_LLAMA_LOG_DIR:-$HOME/.unsloth/studio/logs/llama-server}" export LLAMA_LOG_DIR - parse_connect # writes ~/.claude/settings.json (header=0) + env + parse_connect # prints session env + suppression flags (no ~/.claude write) crosscheck_contract PROMPT='Reply with exactly the single word: pong' - # Phase A: header DISABLED (=0, the documented setting) -> expect a HIT on - # the continued turn. connect.py's ensure_claude_attribution_header() set 0. + # Phase A: the suppression start.py ships (CLAUDE_CODE_ATTRIBUTION_HEADER=0 + + # --exclude-dynamic-system-prompt-sections + --settings overlay) -> expect a + # HIT on the continued turn, since the system-prompt prefix is stable. invoke_via_connect "$LOGS_DIR/claude-ab-hit-1.txt" -p "$PROMPT" # turn 1 primes FROM_HIT="$(bash "$CACHE_HELPER" mark)" # offset before turn 2 invoke_via_connect "$LOGS_DIR/claude-ab-hit-2.txt" -p --continue "$PROMPT again" CACHE_LOG_FROM="$FROM_HIT" bash "$CACHE_HELPER" log HIT - # Phase B: header ENABLED -> expect a MISS. The header prepends a - # per-request-changing attribution line to the system prompt, so the shared - # prefix changes every turn and the KV cache is invalidated (~90% slower); - # this is exactly what the guide flag prevents. - python3 - <<'PY' -import json, os -p = os.path.expanduser("~/.claude/settings.json") -s = json.load(open(p)) if os.path.exists(p) else {} -s.setdefault("env", {})["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "1" -json.dump(s, open(p, "w"), indent=2) -PY + # Phase B: vanilla Claude with the header ENABLED -> expect a MISS. We flip + # the env var to 1 and strip the suppression flags from the launch command + # (without them the dynamic attribution line is included and changes every + # turn, so the shared prefix moves and the KV cache is invalidated, ~90% + # slower). This is session-only: nothing is written to ~/.claude. + CONNECT_ENV_EXTRA='export CLAUDE_CODE_ATTRIBUTION_HEADER=1' + CONNECT_CMD_OVERRIDE="$(printf '%s' "$CONNECT_CMD" \ + | sed -E "s/ --exclude-dynamic-system-prompt-sections//; s/ --settings '[^']*'//")" invoke_via_connect "$LOGS_DIR/claude-ab-miss-1.txt" -p "$PROMPT" FROM_MISS="$(bash "$CACHE_HELPER" mark)" invoke_via_connect "$LOGS_DIR/claude-ab-miss-2.txt" -p --continue "$PROMPT again" CACHE_LOG_FROM="$FROM_MISS" bash "$CACHE_HELPER" log MISS - echo "[claude] attribution A/B OK (header=0 HIT, header=1 MISS)" + unset CONNECT_ENV_EXTRA CONNECT_CMD_OVERRIDE + echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; *) diff --git a/.github/scripts/agent-guides-install.sh b/.github/scripts/agent-guides-install.sh index dfab8aec80..daf4bacd3e 100755 --- a/.github/scripts/agent-guides-install.sh +++ b/.github/scripts/agent-guides-install.sh @@ -7,7 +7,7 @@ # is the single biggest source of false reds, so installs retry with # backoff and the only ::error:: this script can emit is class (b). The # install recipes mirror the install_hint strings in -# unsloth_cli/commands/connect.py at HEAD. +# unsloth_cli/commands/start.py at HEAD. # # Usage: agent-guides-install.sh # agent in: claude codex hermes openclaw opencode pi @@ -25,13 +25,14 @@ install_fail() { } # npm registry flakiness is common in CI; retry 3x with linear backoff. +# Extra npm flags may precede the package (e.g. npm_retry --ignore-scripts pkg). npm_retry() { - local pkg="$1" i + local i for i in 1 2 3; do - if npm install -g "$pkg" >> "$LOG" 2>&1; then + if npm install -g "$@" >> "$LOG" 2>&1; then return 0 fi - echo "[install] npm install -g $pkg attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" + echo "[install] npm install -g $* attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG" sleep "$((i * 10))" done return 1 @@ -60,30 +61,30 @@ curl_bash() { echo "[install] agent=$AGENT (log=$LOG)" case "$AGENT" in claude) - # connect.py install_hint: curl -fsSL https://claude.ai/install.sh | bash + # start.py install_hint: curl -fsSL https://claude.ai/install.sh | bash curl_bash "https://claude.ai/install.sh" || install_fail "claude installer failed" # The installer drops the binary under ~/.local/bin. echo "$HOME/.local/bin" >> "$GITHUB_PATH" ;; codex) - # connect.py install_hint: npm install -g @openai/codex + # start.py install_hint: npm install -g @openai/codex npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed" ;; opencode) - # connect.py install_hint: npm install -g opencode-ai + # start.py install_hint: npm install -g opencode-ai npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed" ;; openclaw) - # connect.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash + # start.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash # npm is the more deterministic path in CI and matches the agent's docs; - # fall back to the connect.py curl installer if the npm tag is missing. + # fall back to the start.py curl installer if the npm tag is missing. if ! npm_retry "openclaw@latest"; then curl_bash "https://openclaw.ai/install.sh" || install_fail "openclaw install failed (npm + curl)" echo "$HOME/.local/bin" >> "$GITHUB_PATH" fi ;; hermes) - # connect.py install_hint: + # start.py install_hint: # curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash curl_bash "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh" \ --non-interactive --skip-setup --skip-browser --no-skills \ @@ -91,11 +92,13 @@ case "$AGENT" in echo "$HOME/.local/bin" >> "$GITHUB_PATH" ;; pi) - # No connect.py recipe; the agent's documented package name. The CLI moved - # from the now-deprecated @mariozechner scope to @earendil-works (the old - # scope is frozen, so installing it would test a stale Pi against the API). - npm_retry "@earendil-works/pi-coding-agent" \ - || install_fail "npm install -g @earendil-works/pi-coding-agent failed" + # start.py install_hint: npm install -g --ignore-scripts @earendil-works/pi-coding-agent + # (--ignore-scripts matches Pi's documented recipe; exercising the exact hint + # catches guide drift). The CLI moved from the now-deprecated @mariozechner + # scope to @earendil-works (the old scope is frozen, so installing it would + # test a stale Pi against the API). + npm_retry --ignore-scripts "@earendil-works/pi-coding-agent" \ + || install_fail "npm install -g --ignore-scripts @earendil-works/pi-coding-agent failed" ;; *) install_fail "unknown agent '$AGENT'" diff --git a/.github/scripts/serve-unsloth-run.sh b/.github/scripts/serve-unsloth-run.sh index 34b8b962c6..6ac98ded7c 100755 --- a/.github/scripts/serve-unsloth-run.sh +++ b/.github/scripts/serve-unsloth-run.sh @@ -27,7 +27,7 @@ # # Outputs written to $GITHUB_ENV (and echoed): # UNSLOTH_API_KEY the sk-unsloth-* key minted on the banner -# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth connect` +# UNSLOTH_STUDIO_URL http://127.0.0.1: (so `unsloth start` # finds THIS server, not the hardcoded :8888) # UNSLOTH_BASE_URL same as UNSLOTH_STUDIO_URL (alias for clarity) # UNSLOTH_MODEL_ID the canonical id reported by /v1/models diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 299ee3f18b..47f75dc1ba 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -6,29 +6,27 @@ # Detects when our local-agent setup recipes drift out of sync with # `unsloth run`. Boots a real `unsloth run --disable-tools` server and # drives the coding agents end to end through the *exact* recipes defined -# in unsloth_cli/commands/connect.py (the in-repo source of truth -- there -# is no docs/ tree). Wherever connect.py has a recipe we drive the agent -# via `unsloth connect --no-launch` and execute what it prints, so -# the test self-updates against connect.py and catches silent recipe drift. +# in unsloth_cli/commands/start.py (the in-repo source of truth -- there +# is no docs/ tree). Wherever start.py has a recipe we drive the agent +# via `unsloth start --no-launch` and execute what it prints, so +# the test self-updates against start.py and catches silent recipe drift. # # Source-of-truth files this workflow guards: -# unsloth_cli/commands/connect.py the `unsloth connect ` recipes +# unsloth_cli/commands/start.py the `unsloth start ` recipes # unsloth_cli/commands/studio.py the `unsloth run` banner (API Key line) # # Failure taxonomy (each surfaced with a distinct ::error:: + the agent name -# + the connect.py location, so a red X is immediately triageable): +# + the start.py location, so a red X is immediately triageable): # (a) Unsloth server/API regression -- the dialect HTTP preflight fails # BEFORE the agent runs (or the server never becomes healthy). # (b) Agent package install failed -- npm/curl install of the CLI failed. # (c) Guide drift -- preflight passed + install ok, but -# the documented `unsloth connect` flow produced no/garbled output. +# the documented `unsloth start` flow produced no/garbled output. # # Agents covered (6): claude, codex, hermes, openclaw, opencode, pi. -# - claude/codex/hermes/openclaw/opencode have a connect.py recipe. -# - pi has NO `unsloth connect pi` command in connect.py at HEAD; it is -# driven by a hand-written recipe and the matrix cell asserts that the -# missing connect recipe is the (known) reason, so the day connect.py -# grows a `pi` command this cell flips to the self-updating path. +# - All six have a `unsloth start ` recipe, so each cell obtains its +# env + command from `unsloth start --no-launch` and runs THAT +# (self-updating: a recipe change is exercised automatically). name: Local Agent Guides CI @@ -83,7 +81,7 @@ jobs: # ═════════════════════════════════════════════════════════════════════ # Job 1: connection # Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect, - # install the agent, run `unsloth connect --no-launch`, execute + # install the agent, run `unsloth start --no-launch`, execute # the emitted recipe with a trivial prompt, assert a non-empty reply. # Runs on PR + weekly + dispatch. Each matrix cell is its own runner so # it serves exactly one model on its own port. @@ -103,7 +101,9 @@ jobs: env: # gemma-4-E4B (128K context, capable enough to drive every agent for a # trivial reply; the 270m model produced empty/failed responses for - # codex/openclaw and is below hermes' 64K context floor). Served as a flat + # codex/openclaw). Hermes' 64K context floor no longer constrains the model + # choice: write_hermes_config claims the floor for smaller windows and + # scales compaction back to the real window. Served as a flat # GGUF file (the -MTP- repo ships no separate draft, so this is plain 4B). GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf @@ -209,7 +209,7 @@ jobs: ;; *) # OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw). - # OpenClaw's connect.py recipe writes an "openai-completions" + # OpenClaw's start.py recipe writes an "openai-completions" # provider (write_openclaw_config), so it uses this path, not # /v1/messages. code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ @@ -227,13 +227,13 @@ jobs: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-install.sh "$AGENT" - # ── (c) drive the agent via connect.py and assert a reply ────────── - # For the 5 agents with a connect.py recipe we run - # `unsloth connect --no-launch`, eval its env/unset exports, + # ── (c) drive the agent via start.py and assert a reply ────────── + # For the 5 agents with a start.py recipe we run + # `unsloth start --no-launch`, eval its env/unset exports, # then run the printed command with a hard timeout (no headless-TTY # hang). Pi has no connect recipe, so it is driven by hand and the # cell asserts that absence is the (known) reason. - - name: Drive ${{ matrix.agent }} via unsloth connect (class-c isolation) + - name: Drive ${{ matrix.agent }} via unsloth start (class-c isolation) env: AGENT: ${{ matrix.agent }} run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT" @@ -248,8 +248,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi @@ -438,8 +440,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi @@ -582,8 +586,10 @@ jobs: # `API Key: `) into logs/unsloth-run-.log, and the upload # step publishes all of logs/, so scrubbing only studio-logs would leak # the bearer token in the retained artifact. + # Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and + # agent-workdir/ are published by the same upload step. if [ -n "${UNSLOTH_API_KEY:-}" ]; then - grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true done fi diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index bce355458a..3022127a2b 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -68,9 +68,10 @@ jobs: pip install -r studio/backend/requirements/studio.txt # Extras that studio.txt does not list but the import chain needs # (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography - # for the auth DB, yaml/jinja2 for utils.models.model_config, etc.): + # for the auth DB, yaml/jinja2 for utils.models.model_config, psutil for + # the orphan-cleanup process scan, etc.): pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests \ 'numpy<3' pytest pytest-asyncio httpx # Torch CPU + transformers are required by a chunk of the backend test @@ -133,7 +134,7 @@ jobs: python -m pip install --upgrade pip pip install -r studio/backend/requirements/studio.txt pip install \ - python-multipart aiofiles sqlalchemy cryptography \ + python-multipart aiofiles sqlalchemy cryptography psutil \ pyyaml jinja2 mammoth unpdf requests typer \ 'numpy<3' pytest pytest-asyncio httpx # torchvision: unsloth_zoo.vision_utils imports it at module scope. @@ -229,7 +230,9 @@ jobs: tests/sh/test_resolve_cuda_archs.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ - tests/sh/test_torch_flavor.sh; do + tests/sh/test_torch_flavor.sh \ + tests/sh/test_with_llama_cpp_dir_flag.sh \ + tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do echo "::group::$s" bash "$s" echo "::endgroup::" diff --git a/.github/workflows/studio-export-capability-ci.yml b/.github/workflows/studio-export-capability-ci.yml new file mode 100644 index 0000000000..1ee6489209 --- /dev/null +++ b/.github/workflows/studio-export-capability-ci.yml @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +# Runs studio/backend/tests/test_export_capability.py on Linux, Windows and macOS. +# +# export_capability() is per-OS (is_apple_silicon() and the PyTorch-import probe differ per +# platform) and the export backend must import without PyTorch, so this confirms the gating and +# import-safety on hosted Windows/macOS. Hosted runners have no GPU/MLX, so a real accelerator +# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block +# torch/unsloth, so the job installs only a CPU PyTorch plus import deps. + +name: Studio export capability + +on: + pull_request: + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + push: + branches: [main] + paths: + - 'studio/backend/utils/hardware/hardware.py' + - 'studio/backend/core/export/export.py' + - 'studio/backend/routes/export.py' + - 'studio/backend/main.py' + - 'studio/backend/tests/test_export_capability.py' + - '.github/workflows/studio-export-capability-ci.yml' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + capability: + name: capability (${{ matrix.os }}) + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + runs-on: ${{ matrix.os }} + timeout-minutes: 20 + env: + # No accelerator on hosted runners; keep detection on the CPU path. + CUDA_VISIBLE_DEVICES: "" + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Upgrade pip + run: python -m pip install --upgrade pip + - name: Install CPU PyTorch + # CPU wheel index so every OS gets a CPU build; keep PyPI as an extra index so torch's + # transitive deps still resolve (matching the other workflows in this repo). + run: python -m pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple "torch>=2.4,<2.13" + - name: Install backend import deps + # Enough to import utils.hardware and core.export.export; NOT unsloth (needs a GPU, and + # the import-safety test blocks it) or triton/llama.cpp (Linux-only / native builds). + run: python -m pip install + transformers peft accelerate safetensors huggingface_hub datasets + sentencepiece protobuf fastapi starlette structlog psutil + python-multipart pydantic httpx "numpy<3" pytest + - name: Export capability + import-safety tests + working-directory: studio/backend + run: python -m pytest tests/test_export_capability.py -q diff --git a/install.ps1 b/install.ps1 index f7f9540970..8c667df079 100644 --- a/install.ps1 +++ b/install.ps1 @@ -99,6 +99,7 @@ function Install-UnslothStudio { $TauriMode = $false $SkipTorch = $false $ShortcutsOnly = $false + $WithLlamaCppDir = "" $argList = $args for ($i = 0; $i -lt $argList.Count; $i++) { switch ($argList[$i]) { @@ -116,6 +117,14 @@ function Install-UnslothStudio { } $PackageName = $argList[$i] } + "--with-llama-cpp-dir" { + $i++ + if ($i -ge $argList.Count) { + Write-Host "[ERROR] --with-llama-cpp-dir requires a path argument." -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir requires a path argument.") + } + $WithLlamaCppDir = $argList[$i] + } } } @@ -2430,6 +2439,13 @@ exit 0 } $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } + if ($WithLlamaCppDir) { + if (-not (Test-Path -LiteralPath $WithLlamaCppDir -PathType Container)) { + Write-Host "[ERROR] --with-llama-cpp-dir path does not exist: $WithLlamaCppDir" -ForegroundColor Red + return (Exit-InstallFailure "--with-llama-cpp-dir path does not exist.") + } + $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR = (Resolve-Path -LiteralPath $WithLlamaCppDir).Path + } $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" # Hand the venv interpreter to setup.ps1 so it reuses the Python we already # resolved and built the venv with, instead of re-probing the system (which @@ -2445,6 +2461,7 @@ exit 0 } else { Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue } + Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } diff --git a/install.sh b/install.sh index 7a9f0be87f..0370559540 100755 --- a/install.sh +++ b/install.sh @@ -53,6 +53,11 @@ _VERBOSE=false _SHORTCUTS_ONLY=false _next_is_package=false _next_is_python=false +_next_is_llama_cpp_dir=false +# Seed from the environment so a caller who exports UNSLOTH_LOCAL_LLAMA_CPP_DIR +# (the documented piped-install style) is honored; the --with-llama-cpp-dir +# flag below overrides it when given. +_WITH_LLAMA_CPP_DIR="${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" for arg in "$@"; do if [ "$_next_is_package" = true ]; then PACKAGE_NAME="$arg" @@ -64,6 +69,11 @@ for arg in "$@"; do _next_is_python=false continue fi + if [ "$_next_is_llama_cpp_dir" = true ]; then + _WITH_LLAMA_CPP_DIR="$arg" + _next_is_llama_cpp_dir=false + continue + fi case "$arg" in --local) STUDIO_LOCAL_INSTALL=true ;; --package) _next_is_package=true ;; @@ -72,6 +82,7 @@ for arg in "$@"; do --no-torch) _NO_TORCH_FLAG=true ;; --verbose|-v) _VERBOSE=true ;; --shortcuts-only) _SHORTCUTS_ONLY=true ;; + --with-llama-cpp-dir) _next_is_llama_cpp_dir=true ;; esac done @@ -255,6 +266,10 @@ if [ "$_next_is_python" = true ]; then echo "❌ ERROR: --python requires a version argument (e.g. --python 3.12)." >&2 exit 1 fi +if [ "$_next_is_llama_cpp_dir" = true ]; then + echo "❌ ERROR: --with-llama-cpp-dir requires a path argument." >&2 + exit 1 +fi # Validate --package to prevent injection into shell/Python commands. # Must start with a letter/digit (rejects leading dashes that uv would parse as flags). @@ -3023,6 +3038,13 @@ _run_setup_with_studio_home() { "$@" fi } +if [ -n "$_WITH_LLAMA_CPP_DIR" ]; then + if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then + echo "[ERROR] --with-llama-cpp-dir path does not exist: $_WITH_LLAMA_CPP_DIR" >&2 + exit 1 + fi + _WITH_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_WITH_LLAMA_CPP_DIR" && pwd -P)" +fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then _run_setup_with_studio_home env \ SKIP_STUDIO_BASE="$_SKIP_BASE" \ @@ -3031,6 +3053,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then STUDIO_LOCAL_INSTALL=1 \ STUDIO_LOCAL_REPO="$_REPO_ROOT" \ UNSLOTH_NO_TORCH="$SKIP_TORCH" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \ bash "$SETUP_SH" bool: + """True if export can run: MLX active, or Unsloth imported (only succeeds on a GPU host).""" + return bool(_IS_MLX) or (FastLanguageModel is not None) + + +def _export_runtime_message() -> str: + """Precise reason the export runtime is unavailable, mirroring hardware.export_capability().""" + if torch is None: + return ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + return ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export on " + "CPU only.)" + ) + + +# Kept for call sites / tests referencing the PyTorch-missing text. +_PYTORCH_MISSING_MESSAGE = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." +) + _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False @@ -58,6 +101,28 @@ def _compressed_export_supported(): return False +def _torchao_export_supported(): + """True if the installed unsloth build has the portable torchao FP8/INT8 export path.""" + try: + import unsloth.save as _us + return hasattr(_us, "_normalize_torchao_method") + except Exception: + return False + + +def _has_nvidia_gpu(): + """True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it.""" + try: + from utils.hardware import hardware as _hw + return _hw.DEVICE == _hw.DeviceType.CUDA and not _hw.IS_ROCM + except Exception: + try: + import torch + return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None + except Exception: + return False + + def _hf_offline(timeout = 3): """True if export should avoid the Hub: honors the HF offline env vars, else does one cheap TCP reachability probe so a network-down load uses local files / the HF cache @@ -394,13 +459,17 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """ Export merged model (for PEFT models). Args: save_directory: Local directory to save model - format_type: "16-bit (FP16)" or "4-bit (FP4)" + format_type: "16-bit (FP16)", "4-bit (FP4)", or a compressed-tensors label + compressed_method: Optional compressed-tensors scheme alias (e.g. "fp8", + "fp8_static", "w8a8", "w4a16", "mxfp4", "mxfp8", "nvfp4"). Overrides + format_type and is resolved against unsloth.save COMPRESSED_EXPORT_SCHEMES. push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID (username/model-name) hf_token: Hugging Face token @@ -409,38 +478,108 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - if not self.is_peft: - return ( - False, - "This is not a PEFT model. Use 'Export Base Model' instead.", - None, - ) + # Merged export works for PEFT adapters and non-PEFT Local/HF base models alike + # (save_pretrained_merged is a no-op merge that just saves the base). output_path: Optional[str] = None - # compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and - # write to a sibling "-" directory (for vLLM). - _COMPRESSED = { - "FP8 (compressed-tensors)": ("fp8", "fp8"), - "NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"), + # Quantized formats save to a sibling "-". Two backends: compressed-tensors + # (llm-compressor, NVIDIA-only) and portable torchao FP8/INT8 (device-agnostic). The alias + # comes from `compressed_method` (the "all formats" dropdown) or the `format_type` label. + _LABEL_TO_ALIAS = { + "FP8 (compressed-tensors)": "fp8", + "NVFP4 (compressed-tensors)": "nvfp4", } - is_compressed = format_type in _COMPRESSED + compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type) + compressed_suffix: Optional[str] = None + # Classify the alias: torchao-portable vs compressed-tensors. + torchao_info = None + if compressed_alias and _torchao_export_supported(): + try: + import unsloth.save as _us_t + torchao_info = _us_t._normalize_torchao_method(compressed_alias) + except Exception: + torchao_info = None + is_torchao = torchao_info is not None + is_compressed = compressed_alias is not None and not is_torchao try: - if _IS_MLX: - if is_compressed: - return False, "Compressed-tensors export is not supported on macOS/MLX.", None - mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" - elif is_compressed: + if _IS_MLX and (is_compressed or is_torchao): + return ( + False, + "Quantized (FP8/FP4/INT) export is not supported on macOS/MLX. " + "Use 16-bit or GGUF.", + None, + ) + + if is_torchao: + # Portable torchao: no NVIDIA GPU, no calibration. + compressed_suffix = torchao_info[1] + + if is_compressed: + # compressed-tensors needs CUDA; enforce in the backend even if the UI gate is bypassed. + if not _has_nvidia_gpu(): + return ( + False, + "Compressed-tensors (FP8/FP4) export requires an NVIDIA GPU. On other " + "hardware use the portable FP8/INT8 (torchao) formats or 16-bit.", + None, + ) if not _compressed_export_supported(): return ( False, - "Compressed-tensors (FP8/NVFP4) export requires an Unsloth build with " + "Compressed-tensors (FP8/FP4) export requires an Unsloth build with " "compressed-tensors support. Upgrade unsloth, or choose 16-bit.", None, ) - save_method = _COMPRESSED[format_type][0] + import unsloth.save as _us + + # Prefer the llm-compressor-main shadow (transformers 5.x): it quantizes newer models + # (Qwen3.5, Gemma-4, ...) the shipped 0.10.x cannot. Route all compressed exports + # through it when available; else fall back to the workspace 0.10.x path below. + _shadow_pp = None + try: + from utils.transformers_version import llmcompressor_shadow_pythonpath + _shadow_pp = llmcompressor_shadow_pythonpath() + except Exception as e: + logger.warning(f"llm-compressor-main shadow unavailable: {e}") + if _shadow_pp: + os.environ[_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV] = _shadow_pp + else: + # No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its + # transformers ceiling, so fail fast for sidecar models; default-tier still works. + os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None) + _exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + return ( + False, + "FP8/FP4 compressed-tensors export is not available for this model: it " + f"runs under transformers {_tf_ver}, but the installed llm-compressor " + f"supports transformers <= {_us._LLM_COMPRESSOR_MAX_TRANSFORMERS} and the " + "llm-compressor-main runtime could not be provisioned (offline or " + "UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN). Export to GGUF or 16-bit instead.", + None, + ) + + try: + info = _us._normalize_compressed_method(compressed_alias) + except Exception as e: + return False, f"Unsupported compressed export '{compressed_alias}': {e}", None + if info is None: + return ( + False, + f"'{compressed_alias}' is not a recognized compressed-tensors export.", + None, + ) + compressed_suffix = info[2] + + if _IS_MLX: + mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit" + elif is_compressed or is_torchao: + save_method = compressed_alias elif format_type == "4-bit (FP4)": save_method = "merged_4bit_forced" elif self._audio_type == "whisper": @@ -464,10 +603,10 @@ class ExportBackend: save_directory, self.current_tokenizer, save_method = save_method ) - # Compressed export writes to the "-" sibling; report that as output. + # Compressed / torchao writes to the "-" sibling; report that as output. final_dir = ( - f"{save_directory}-{_COMPRESSED[format_type][1]}" - if is_compressed + f"{save_directory}-{compressed_suffix}" + if (is_compressed or is_torchao) else save_directory ) self._write_export_metadata(final_dir) @@ -507,10 +646,9 @@ class ExportBackend: token = hf_token, private = private, ) - elif is_compressed and output_path and Path(output_path).is_dir(): - # The compressed model was already built locally in output_path; upload it - # directly so we do not re-run the (expensive, OOM-prone) compression that - # push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time. + elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir(): + # Already built in output_path; upload it directly instead of re-running the + # expensive quantization that push_to_hub_merged(save_method=...) would redo. hf_api = HfApi(token = hf_token) repo_id = PushToHubMixin._create_repo( PushToHubMixin, @@ -522,7 +660,7 @@ class ExportBackend: username = repo_id.split("/")[0], base_model = getattr(self.current_model.config, "_name_or_path", "unknown"), model_type = getattr(self.current_model.config, "model_type", "llm"), - method = format_type, + method = compressed_alias or format_type, extra = "unsloth", ) ModelCard(content).push_to_hub( @@ -568,6 +706,8 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None @@ -686,7 +826,7 @@ class ExportBackend: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, @@ -697,7 +837,9 @@ class ExportBackend: Args: save_directory: Local directory to save model - quantization_method: GGUF quantization method (e.g., "Q4_K_M") + quantization_method: A single GGUF quant method (e.g., "Q4_K_M") or a list of them + (e.g., ["Q4_K_M", "Q8_0"]). A list produces one GGUF per quant from a single + model load (unsloth save_to_gguf loops internally). push_to_hub: Whether to push to Hugging Face Hub repo_id: Hub repository ID hf_token: Hugging Face token @@ -705,11 +847,13 @@ class ExportBackend: Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None - # Only forward imatrix_file to an unsloth build that accepts it; otherwise even a plain - # no-imatrix export would fail with an unexpected-keyword error against an older unsloth. + # Only forward imatrix_file to an unsloth build that accepts it, else older builds raise + # an unexpected-keyword error even for a plain no-imatrix export. if imatrix_file is not None and not _supports_kwarg( self.current_model.save_pretrained_gguf, "imatrix_file" ): @@ -724,8 +868,14 @@ class ExportBackend: output_path: Optional[str] = None model_tmp_to_cleanup: Optional[str] = None try: - # unsloth expects lowercase quant method - quant_method = quantization_method.lower() + # Normalize to a lowercased list so multiple quants come from one model load. + if isinstance(quantization_method, (list, tuple)): + quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()] + else: + quant_methods = [str(quantization_method).lower()] + if not quant_methods: + quant_methods = ["q4_k_m"] + quant_method = quant_methods if len(quant_methods) > 1 else quant_methods[0] # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it # can't drift past the pinned llama-quantize binary's gguf API. @@ -847,7 +997,7 @@ class ExportBackend: return ( True, - f"GGUF model exported successfully ({quantization_method})", + f"GGUF model exported successfully ({', '.join(quant_methods)})", output_path, ) @@ -867,19 +1017,56 @@ class ExportBackend: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: """ Export LoRA adapter only (not merged). + Args: + gguf: If True, also convert the adapter to a GGUF LoRA file (llama.cpp + convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`. + gguf_outtype: GGUF LoRA output float type; one of q8_0/f16/bf16/f32. + Returns: Tuple of (success: bool, message: str, output_path: Optional[str]) """ + if not _export_runtime_available(): + return False, _export_runtime_message(), None if not self.current_model or not self.current_tokenizer: return False, "No model loaded. Please select a checkpoint first.", None if not self.is_peft: return False, "This is not a PEFT model. No adapter to export.", None + _GGUF_LORA_OUTTYPES = ("q8_0", "f16", "bf16", "f32") + if gguf: + if _IS_MLX: + return ( + False, + "GGUF LoRA adapter export is not supported on macOS/MLX. " + "Use the safetensors adapter instead.", + None, + ) + outtype = str(gguf_outtype).lower() + if outtype not in _GGUF_LORA_OUTTYPES: + return ( + False, + f"Invalid GGUF LoRA outtype '{gguf_outtype}'. " + f"Choose one of {', '.join(_GGUF_LORA_OUTTYPES)}.", + None, + ) + # getattr so an older build without save_pretrained_gguf returns a clean message + # instead of an AttributeError (a generic 500). + _save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None) + if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"): + return ( + False, + "This Unsloth build does not support GGUF LoRA adapter export. " + "Upgrade unsloth and unsloth_zoo, or export the safetensors adapter.", + None, + ) + output_path: Optional[str] = None try: if save_directory: @@ -887,7 +1074,24 @@ class ExportBackend: logger.info(f"Saving LoRA adapter locally to: {save_directory}") ensure_dir(Path(save_directory)) - if _IS_MLX: + if gguf: + # Writes the adapter files plus "-lora-.gguf". + _apply_wsl_sudo_patch() + self.current_model.save_pretrained_gguf( + save_directory, + self.current_tokenizer, + save_method = "lora", + quantization_method = outtype, + # Forward the token so convert_lora_to_gguf.py can fetch a gated base's config. + token = hf_token or None, + ) + final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf"))) + logger.info( + "LoRA GGUF export complete. Files in %s:\n %s", + save_directory, + "\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)", + ) + elif _IS_MLX: # MLX: save adapters.safetensors + tokenizer files self.current_model.save_lora_adapters(save_directory) self.current_tokenizer.save_pretrained(save_directory) @@ -907,7 +1111,24 @@ class ExportBackend: logger.info(f"Pushing LoRA adapter to Hub: {repo_id}") - if _IS_MLX: + if gguf: + # Upload the locally-built GGUF folder; needs a local save_directory so the + # conversion is not re-run. + if not (output_path and Path(output_path).is_dir()): + return ( + False, + "GGUF LoRA Hub upload requires a local save directory; set one and " + "retry.", + None, + ) + hf_api = HfApi(token = hf_token) + hf_api.create_repo(repo_id, private = private, exist_ok = True) + hf_api.upload_folder( + folder_path = output_path, + repo_id = repo_id, + repo_type = "model", + ) + elif _IS_MLX: with tempfile.TemporaryDirectory() as tmp_dir: self.current_model.save_lora_adapters(tmp_dir) self.current_tokenizer.save_pretrained(tmp_dir) diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 052a47dd80..671ef363f5 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -456,6 +456,7 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + compressed_method: Optional[str] = None, ) -> Tuple[bool, str, Optional[str]]: """Export merged PEFT model.""" return self._run_export( @@ -467,6 +468,7 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "compressed_method": compressed_method, }, ) @@ -495,13 +497,13 @@ class ExportOrchestrator: def export_gguf( self, save_directory: str, - quantization_method: str = "Q4_K_M", + quantization_method = "Q4_K_M", push_to_hub: bool = False, repo_id: Optional[str] = None, hf_token: Optional[str] = None, imatrix_file = None, ) -> Tuple[bool, str, Optional[str]]: - """Export model in GGUF format.""" + """Export model in GGUF format. `quantization_method` may be a single method or a list.""" return self._run_export( "gguf", { @@ -521,8 +523,10 @@ class ExportOrchestrator: repo_id: Optional[str] = None, hf_token: Optional[str] = None, private: bool = False, + gguf: bool = False, + gguf_outtype: str = "q8_0", ) -> Tuple[bool, str, Optional[str]]: - """Export LoRA adapter only.""" + """Export LoRA adapter only (optionally also as a GGUF LoRA file).""" return self._run_export( "lora", { @@ -531,6 +535,8 @@ class ExportOrchestrator: "repo_id": repo_id, "hf_token": hf_token, "private": private, + "gguf": gguf, + "gguf_outtype": gguf_outtype, }, ) @@ -557,9 +563,13 @@ class ExportOrchestrator: cmd = {"type": "export", "export_type": export_type, **params} try: self._send_cmd(cmd) + # GGUF for 30B+ models can take 30+ min per quant; a multi-quant list runs them + # all in one op off a single merge, so scale the timeout by the quant count. + _qm = params.get("quantization_method") + _n = len(_qm) if isinstance(_qm, (list, tuple)) and _qm else 1 resp = self._wait_response( f"export_{export_type}_done", - timeout = 3600, # GGUF for 30B+ models can take 30+ min + timeout = 3600 * max(1, _n), ) op_success = resp.get("success", False) op_message = resp.get("message", "") diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index d473dcb54f..7828116236 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -397,6 +397,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + compressed_method = cmd.get("compressed_method"), ) elif export_type == "base": success, message, output_path = backend.export_base_model( @@ -423,6 +424,8 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: repo_id = cmd.get("repo_id"), hf_token = cmd.get("hf_token"), private = cmd.get("private", False), + gguf = cmd.get("gguf", False), + gguf_outtype = cmd.get("gguf_outtype", "q8_0"), ) else: success, message = False, f"Unknown export type: {export_type}" diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 0307336dde..7b572a28ff 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -494,6 +494,29 @@ class AnthropicPassthroughEmitter: self._usage: dict = {} self._stop_reason: str = "end_turn" self._stop_sequence: Optional[str] = None + # Optional text-form tool-call healing (client-tool passthrough only). + self._healer = None + self._healed_tool_use = False + self._healed_call_count = 0 + self._heal_disable_parallel = False + + def enable_healing( + self, + allowed_tools: set, + tools: Optional[list] = None, + *, + disable_parallel_tool_use: bool = False, + ) -> None: + """Promote text-form tool calls in streamed content to tool_use blocks. + + Only calls naming a tool in ``allowed_tools`` (the client's declared + tools) are promoted; everything else streams as text exactly as before. + Never enabled for Studio's own tool loop. + """ + from core.inference.passthrough_healing import StreamToolCallHealer + + self._healer = StreamToolCallHealer(allowed_tools, tools) + self._heal_disable_parallel = disable_parallel_tool_use def start( self, @@ -542,29 +565,42 @@ class AnthropicPassthroughEmitter: delta = choice.get("delta") or {} finish_reason = choice.get("finish_reason") + # ── Structured tool calls take precedence over healing ── + # Grammar mode worked: flush anything the healer held (it preceded the + # call in the model's output) and relay verbatim from here on. + if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant: + for kind, value in self._healer.structured_tool_call_seen(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + # ── Text content ── content = delta.get("content") - if content: - if self._current_block_type != "text": - if self._current_block_type is not None: - events.append(self._close_current_block()) - events.extend(self._open_text_block()) - events.append( - build_anthropic_sse_event( - "content_block_delta", - { - "type": "content_block_delta", - "index": self.block_index, - "delta": {"type": "text_delta", "text": content}, - }, - ) - ) + if content and self._healer is not None and not self._healer.dormant: + # Route text through the healer: held/promoted portions become + # synthetic tool_use blocks, the rest streams as text unchanged. + for kind, value in self._healer.feed(content): + if kind == "text": + events.extend(self._emit_text_delta(value)) + else: + events.extend(self._emit_healed_tool_use(value)) + elif content: + events.extend(self._emit_text_delta(content)) # ── Tool calls (streaming deltas) ── tool_calls = delta.get("tool_calls") or [] for tc in tool_calls: tc_idx = tc.get("index", 0) fn = tc.get("function") or {} + if ( + self._heal_disable_parallel + and tc_idx not in self._tool_call_states + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # disable_parallel_tool_use: a healed call already consumed the + # single allowed slot. The caller's chunk-level cap only sees + # native indexes, so drop this native call (and its later + # argument deltas, which never allocate a state either). + continue if tc_idx not in self._tool_call_states: # New tool call — close prior block, open tool_use block if self._current_block_type is not None: @@ -618,6 +654,17 @@ class AnthropicPassthroughEmitter: def finish(self) -> list[str]: events: list[str] = [] + if self._healer is not None: + # Last-chance heal of any held residue (e.g. an unclosed tool block). + for kind, value in self._healer.finalize(): + if kind == "text" and value: + events.extend(self._emit_text_delta(value)) + elif kind == "tool_call": + events.extend(self._emit_healed_tool_use(value)) + if self._healed_tool_use and self._stop_reason != "max_tokens": + # A promoted call must stop for tool use; a truncation still wins + # (its arguments may be incomplete). + self._stop_reason = "tool_use" if self._current_block_type is not None: events.append(self._close_current_block()) events.append( @@ -641,6 +688,76 @@ class AnthropicPassthroughEmitter: ) return events + def _emit_text_delta(self, content: str) -> list[str]: + events: list[str] = [] + if self._current_block_type != "text": + if self._current_block_type is not None: + events.append(self._close_current_block()) + events.extend(self._open_text_block()) + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": content}, + }, + ) + ) + return events + + def _emit_healed_tool_use(self, call: dict) -> list[str]: + # A healed call arrives complete, so its tool_use block opens, carries + # one input_json_delta, and closes immediately; an open text block is + # closed first (only the safe prefix ever streamed into it). + if ( + self._heal_disable_parallel + and (self._healed_call_count + len(self._tool_call_states)) >= 1 + ): + # Healed and native calls share the single allowed slot. + return [] + events: list[str] = [] + if self._current_block_type is not None: + events.append(self._close_current_block()) + function = call.get("function") or {} + tool_id = anthropic_tool_use_id("") + self.block_index += 1 + self._current_block_type = "tool_use" + events.append( + build_anthropic_sse_event( + "content_block_start", + { + "type": "content_block_start", + "index": self.block_index, + "content_block": { + "type": "tool_use", + "id": tool_id, + "name": function.get("name", ""), + "input": {}, + }, + }, + ) + ) + arguments = function.get("arguments") or "" + if arguments: + events.append( + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": { + "type": "input_json_delta", + "partial_json": arguments, + }, + }, + ) + ) + events.append(self._close_current_block()) + self._healed_tool_use = True + self._healed_call_count += 1 + return events + def _open_text_block(self) -> list[str]: self.block_index += 1 self._current_block_type = "text" diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index af62a29679..3ccfc5cdfe 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -801,7 +801,11 @@ _MTP_MIN_SIZE_B = 3.0 # Cap total GPU occupancy at this fraction of the card. The fit reserves an # absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction # of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve. -_CTX_FIT_VRAM_FRACTION = 0.95 +# 3%: the context-linear compute buffer is now modelled (_compute_buffer_ctx_bytes), +# so this cushion no longer covers it - only fragmentation, the per-device CUDA +# context on a multi-GPU split, and MoE routing, which measure ~2-3% (Qwen3.5-397B on +# 3 GPUs under-predicts by 2.7%). Below 3% one fragmentation spike overflows to CPU. +_CTX_FIT_VRAM_FRACTION = 0.97 # Apple unified memory is shared with the OS, so tighter than VRAM. Matches the # 0.85 MLX uses in mlx_inference.py (_configure_memory_limits); not kept in sync. @@ -1223,6 +1227,25 @@ def _backfill_usage_from_timings(usage, timings): return out +def _is_external_link(path: Path) -> bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, which Studio does not own.""" + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. @@ -1254,6 +1277,7 @@ class LlamaCppBackend: self._is_diffusion: bool = False self._diffusion_visual_bin: Optional[str] = None self._healthy = False + self._load_rss_hwm = (None, 0) # (pid, peak VmRSS) for load_progress self._stats_logger = None # vLLM-style engine-stats poller, set on load # Set by _classify_gpu_offload after _wait_for_health. self._gpu_offload_active: Optional[bool] = None @@ -1461,6 +1485,21 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + @staticmethod + def _read_rss_bytes(pid: int) -> Optional[int]: + """Resident set size of ``pid`` in bytes, from /proc//status (Linux). + 0 when the status has no VmRSS line (zombie / kernel thread); None where + /proc is unavailable (macOS/Windows) or the value is unreadable.""" + try: + with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: + for line in f: + if line.startswith("VmRSS:"): + # IndexError guards a "VmRSS:" line with no value column. + return int(line.split()[1]) * 1024 # kB -> bytes + except (FileNotFoundError, PermissionError, ValueError, IndexError, OSError): + return None + return 0 # readable but no VmRSS line + def load_progress(self) -> Optional[dict]: """Return live model-load progress, or None if not loading. @@ -1520,22 +1559,32 @@ class LlamaCppBackend: except OSError: pass - # Read VmRSS from /proc//status (kilobytes on Linux). - bytes_loaded = 0 - try: - with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f: - for line in f: - if line.startswith("VmRSS:"): - kb = int(line.split()[1]) - bytes_loaded = kb * 1024 - break - except (FileNotFoundError, PermissionError, ValueError, OSError): + # VmRSS of the llama-server; None where /proc is unavailable. + bytes_loaded = LlamaCppBackend._read_rss_bytes(pid) + if bytes_loaded is None: return None + # RSS climbs as weights page in, then drops once -ngl offloads them to + # VRAM and the mmap pages are freed. Hold a per-process high-water mark + # so the bar never regresses to ~8% mid-load (#5740). + hwm_pid, hwm = getattr(self, "_load_rss_hwm", (None, 0)) + hwm = bytes_loaded if hwm_pid != pid else max(hwm, bytes_loaded) + self._load_rss_hwm = (pid, hwm) + bytes_loaded = hwm + phase = "ready" if self._healthy else "mmap" fraction = 0.0 if bytes_total > 0: fraction = min(1.0, bytes_loaded / bytes_total) + # Once llama-server is healthy the load is complete by definition. With + # layers offloaded to VRAM (-ngl) the process releases the mmap'd weight + # pages, so VmRSS sinks back well below the shard total; the raw RSS + # fraction would then report a partial (~8%) load indefinitely and freeze + # a fraction-driven progress bar even though the model is ready (#5740). + if self._healthy: + if bytes_total > 0: + bytes_loaded = bytes_total + fraction = 1.0 return { "phase": phase, "bytes_loaded": bytes_loaded, @@ -2419,9 +2468,10 @@ class LlamaCppBackend: prev = curr # Free-VRAM fraction at which Studio pins the GPU directly instead of - # deferring to ``--fit on``. 5% headroom covers CUDA context + compute - # buffers; 0.90 dropped 91-94% fits to CPU offload (#5106). - _GPU_PIN_VRAM_FRACTION = 0.95 + # deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in + # the fit, so this only guards fragmentation + multi-GPU per-device CUDA context + # (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106). + _GPU_PIN_VRAM_FRACTION = 0.97 # Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF # dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived @@ -2977,6 +3027,27 @@ class LlamaCppBackend: _DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it _COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate + # Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682). + _CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB) + _MMPROJ_VRAM_SAFETY = 1.4 # mmproj worst-case buffer vs file size (runtime ~1.3x) + _MTP_DRAFT_COMPUTE_BYTES = 224 * 1024 * 1024 # MTP draft decode graph beyond its KV + # The flash-attn KQ mask + attention scratch grow ~linearly with context; the flat + # _estimate_compute_buffer_bytes term only covers ctx -> 0. The per-token rate + # depends on the KV cache type: a QUANTIZED cache (q8_0/q5/q4/iq4) needs a + # context-sized dequant scratch that scales with n_embd, measured at 0.74-2.02 x + # n_embd across Qwen3.5/3.6 (2B/4B/9B/27B) and Gemma-4 (12B/31B) at q8_0; an + # f16/bf16/f32 cache skips the dequant and pays only the KQ mask, a flat n_ubatch*2 + # bytes per context token regardless of n_embd (measured 1024 B/tok on Qwen-9B and + # Gemma-31B alike). So Qwen3.5-4B at 256k is 1.30 GiB at q8_0 vs 0.31 GiB at f16. + # 2.25 covers the worst quantized case (Qwen3.5-4B, ~2.0x) plus the under-modeled + # flat base; the mask safety covers the f16 base gap. Without this term, tight tiers + # at extreme context over-pin and spill to CPU (the 3% cushion is only ~0.25 GiB on + # an 8 GB card, far below the ~1-2.4 GiB quantized buffer at 256k): e.g. Qwen3.5-4B + # Q4 at 256k needs ~8.5 GiB on a real 8 GB card (weights 2.4 + KV 4.3 + compute 1.3 + # + CUDA ctx) -> CPU spill; with this reserve the auto context caps to ~210k, fits. + _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) + _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) + _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) def _estimate_compute_buffer_bytes( self, @@ -3007,6 +3078,85 @@ class LlamaCppBackend: compute = act_scratch + out_buffer * max(0, par - 1) return int(compute * self._COMPUTE_BUFFER_SAFETY) + def _compute_buffer_ctx_bytes( + self, + n_ctx: int, + n_ubatch: Optional[int] = None, + cache_type_kv: Optional[str] = None, + ) -> int: + """Context-linear growth of the per-device compute buffer (bytes), charged + on top of the flat ``_estimate_compute_buffer_bytes``. The flash-attn KQ + mask + attention scratch scale ~linearly with context and with the micro- + batch; the flat term only covers ctx -> 0. A quantized KV cache adds a + context-sized dequant scratch that scales with n_embd; f16/bf16/f32 pays only + the KQ mask, a flat n_ubatch*2 bytes per context token. ``cache_type_kv`` None + -> f16 (llama.cpp's default; an env-set quantized cache is budgeted as f16 on + the KV side, whose over-reservation absorbs the dequant scratch). Returns 0 + when dims are missing or ``n_ctx`` <= 0.""" + n_embd = self._embedding_length or 0 + if n_embd <= 0 or n_ctx <= 0: + return 0 + ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + if _kv_bytes_per_elem(cache_type_kv) < 2.0: + # Quantized cache: the dequant scratch dominates and scales with n_embd. + # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on + # GLM-5.2 and Kimi-K2.7 vs up to 2.02x on regular attention. + ub_scale = ub / self._DEFAULT_N_UBATCH + rate = ( + self._CTX_COMPUTE_BYTES_PER_EMBD_MLA + if self._key_length_mla + else self._CTX_COMPUTE_BYTES_PER_EMBD + ) + per_tok = rate * n_embd * ub_scale + else: + # f16/bf16/f32: only the KQ mask ([n_kv, n_ubatch] f16), n_embd-independent. + per_tok = ub * 2 * self._CTX_COMPUTE_F16_MASK_SAFETY + return int(per_tok * n_ctx) + + def _slots_that_fit_on_gpu( + self, + n_parallel: int, + effective_ctx: int, + gpus: list[tuple[int, int]], + total_by_idx: Optional[dict[int, int]], + base_footprint_bytes: int, + cache_type_kv: Optional[str], + pin_fraction: float, + per_device_overhead_bytes: int, + min_gpus: int, + n_ubatch: Optional[int] = None, + ) -> tuple[Optional[list[int]], bool, int]: + """Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits, + so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers + to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the + slot-independent footprint (weights + soft overhead + MTP + context-linear compute, + minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer + and KV, then re-selects GPUs like the explicit-context path. Returns (gpu_indices, + use_fit=False, slots) for the largest fitting count, else (None, True, n_parallel). + Only ever reduces; deterministic and unit-testable with synthetic VRAM maps.""" + for slots in range(n_parallel - 1, 0, -1): + cb = self._estimate_compute_buffer_bytes( + n_ubatch = n_ubatch, n_parallel = slots, per_device_tensor = False + ) + if cb <= 0: + cb = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 + total = ( + base_footprint_bytes + + cb + + self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots) + ) + gpu_indices, use_fit = self._select_gpus( + total, + gpus, + usable_fraction = pin_fraction, + total_by_idx = total_by_idx, + per_device_overhead_bytes = per_device_overhead_bytes, + min_gpus = min_gpus, + ) + if not use_fit: + return gpu_indices, False, slots + return None, True, n_parallel + def _fit_context_to_vram( self, requested_ctx: int, @@ -3022,6 +3172,7 @@ class LlamaCppBackend: kv_on_gpu: bool = True, mtp_engaged: bool = False, mtp_overhead_fn: Optional[Callable[[int], int]] = None, + compute_ctx_bytes_fn: Optional[Callable[[int], int]] = None, budget_frac: Optional[float] = None, total_mib: Optional[int] = None, ) -> int: @@ -3073,9 +3224,14 @@ class LlamaCppBackend: def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + def _cc_at(ctx: int) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + scratch); + # the flat term in model_footprint only covers ctx -> 0. + return compute_ctx_bytes_fn(ctx) if compute_ctx_bytes_fn is not None else 0 + # Already fits? kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs) - if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes: + if model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) <= budget_bytes: return requested_ctx # Weights + compute buffer alone exceed budget -- reducing ctx can't help. @@ -3096,7 +3252,7 @@ class LlamaCppBackend: while lo <= hi: mid = (lo + hi) // 2 kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs) - if kv + _mtp_at(mid) <= remaining: + if kv + _mtp_at(mid) + _cc_at(mid) <= remaining: best = mid lo = mid + 1 else: @@ -4213,6 +4369,17 @@ class LlamaCppBackend: "expected; otherwise check the llama-server log for the cause." ) + # A live server that never answered 200 on /health is not a bad GGUF: + # the load is too large for VRAM/context, or a local proxy/VPN grabbed + # the loopback probe (#5740). + if "health check timed out" in lowered: + return ( + "llama-server started but never became healthy on its local " + "/health endpoint. Try a smaller context length or a more " + "quantized GGUF, and if you use a VPN or HTTP proxy make sure " + "localhost bypasses it (NO_PROXY=127.0.0.1,localhost)." + ) + # Fallback: genuinely unknown failure (OOM, missing binary ...). return ( "llama-server failed to start. " @@ -4232,6 +4399,7 @@ class LlamaCppBackend: max_target_ctx: Optional[int] = None, total_by_idx: Optional[dict[int, int]] = None, n_ubatch: Optional[int] = None, + soft_overhead_bytes: int = 0, ) -> tuple[int, int, list[int], Optional[list[int]]]: """Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed. @@ -4243,9 +4411,11 @@ class LlamaCppBackend: ``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``. Policy (assumes >= 2 GPUs; the caller drops the toggle below that): - - Cap context to the KV that fits the pooled VRAM after the weights and - one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``, - deterministic from dims; flat fallback when dims are unavailable). + - Cap context to the KV that fits the pooled VRAM after the weights, one + per-device flat compute-graph buffer (``_estimate_compute_buffer_bytes``, + deterministic from dims; flat fallback when dims are unavailable), and the + per-device context-linear compute growth (``_compute_buffer_ctx_bytes``, + replicated on every device in tensor mode, so summed over the split). llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only cap, honored even for an explicit ``-c``. It is more accurate than the 0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused. @@ -4254,7 +4424,9 @@ class LlamaCppBackend: share fits the smallest GPU; otherwise it is weighted by usable budget so the roomier GPU absorbs more weight and the smallest keeps room for KV. ``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes - the compute buffer. + the compute buffer. ``soft_overhead_bytes`` is the CUDA-context / mmproj / + MTP-draft-graph reserve the layer path folds into ``model_size_fit``; + charged against the pooled budget so tensor mode reserves the same overhead. """ # Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a @@ -4300,16 +4472,40 @@ class LlamaCppBackend: flat_mtp_bytes = max(0, mtp_flat_reserve_bytes) if mtp_engaged and mtp_overhead_fn is None: flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3) + # soft_overhead_bytes is the CUDA-context / mmproj / MTP-draft-graph reserve + # the layer path folds into model_size_fit. Tensor mode has no --fit valve, so + # an unreserved overshoot OOMs at startup rather than offloading; charge it here + # too. Once (pooled), mirroring the layer path -- the per-device CUDA context is + # a known slight under-charge, left for real multi-GPU data. kv_budget_b = ( - (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes + (pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 + - model_size + - flat_mtp_bytes + - max(0, soft_overhead_bytes) ) def _mtp_at(ctx: int) -> int: return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0 + # Context-linear compute buffer, summed over the split. Tensor mode + # replicates the compute graph on EVERY device (measured: the per-device + # buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at + # f16, independent of n_embd), so the growth is n_dev x the per-device + # term. cache_type_kv here is always non-quantized (tensor forces f16), so + # _compute_buffer_ctx_bytes returns the light KQ-mask term, not the heavy + # quantized dequant scratch. The flat reserve_mib above only covers ctx->0; + # without this the fit over-pins and OOMs at high context on a tight pool + # (0.5-4 GiB unreserved at 262k-1M across 2-4 GPUs), the tensor-mode analog + # of the layer-split compute bug. + n_dev = len(gpu_indices) + + def _cc_ctx(ctx: int) -> int: + return n_dev * self._compute_buffer_ctx_bytes(ctx, n_ubatch, cache_type_kv) + def _fit_ctx(ctx: int) -> int: - # Largest context whose KV (+ MTP draft reserve) fits the pooled - # budget. Floors small, but never raises an explicit ctx above asked. + # Largest context whose KV (+ MTP draft reserve + context-linear + # compute) fits the pooled budget. Floors small, but never raises an + # explicit ctx above asked. if self._can_estimate_kv() and ctx > 0: ctx_floor = min(2048, ctx) if kv_budget_b <= 0: @@ -4317,11 +4513,13 @@ class LlamaCppBackend: # falls back to layer split. return ctx_floor if mtp_overhead_fn is not None: - # kv(ctx)+mtp(ctx) is not single-linear, so binary search. + # kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search. def _consumer(c: int) -> int: - return self._estimate_kv_cache_bytes( - c, cache_type_kv, n_parallel = n_parallel - ) + _mtp_at(c) + return ( + self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel) + + _mtp_at(c) + + _cc_ctx(c) + ) if _consumer(ctx) <= kv_budget_b: return ctx @@ -4335,9 +4533,10 @@ class LlamaCppBackend: hi = mid - 1 return best kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel) - if kv_at <= kv_budget_b: + total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin + if total_at <= kv_budget_b: return ctx - return max(ctx_floor, int(ctx * kv_budget_b / kv_at)) + return max(ctx_floor, int(ctx * kv_budget_b / total_at)) # KV size unknown -> can't prove a safe cap; floor. return min(4096, ctx) if ctx > 0 else 4096 @@ -4357,10 +4556,23 @@ class LlamaCppBackend: # The MTP reserve also has to fit the even split (mirror the pooled budget): # byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above. mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes - even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024) + # Context-linear compute is replicated per device; charge the whole split so + # the weighted ratio reflects it (mirrors kv_budget_b's per-device reserve). + cc_bytes = _cc_ctx(effective_ctx) if effective_ctx > 0 else 0 + even_share_mib = ( + (model_size + kv_bytes + mtp_bytes + cc_bytes) / len(gpu_indices) / (1024 * 1024) + ) tensor_split: Optional[list[int]] = None if even_share_mib > (min_usable_mib - reserve_mib): - adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices] + # Each device also holds its replicated share of the context-linear + # compute (cc_bytes/n_dev) on top of the flat reserve. The even-share + # gate above charges cc_bytes; the split weights must subtract it too, or + # the smaller card is weighted above its real usable budget and OOMs (the + # per-device analog of the layer path's per-GPU overhead in _select_gpus). + cc_per_dev_mib = (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0 + adj = [ + max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) for i in gpu_indices + ] if sum(adj) > 0: tensor_split = adj return effective_ctx, max_available_ctx, gpu_indices, tensor_split @@ -5123,6 +5335,20 @@ class LlamaCppBackend: # compute buffer); None -> the 512 default in the estimate. _effective_ubatch = _extra_args_n_ubatch(extra_args) + def _cc_bytes(ctx: int, n_gpus: int = 1) -> int: + # Context-linear compute-buffer growth (flash-attn KQ mask + + # attention scratch); the flat _compute_buffer_pipeline folded + # into model_size_fit only covers ctx -> 0. Charged per + # candidate context so the fit can't over-pin and spill. The + # rate depends on the KV cache type (quantized adds a dequant + # scratch), so pass it through. In a layer split this buffer is + # replicated on EVERY device (measured ~equal per GPU), so scale + # by the device count; a large model at high context otherwise + # under-reserves ~(n-1)x it (e.g. Qwen3.5-397B on 3 GPUs). + return max(1, n_gpus) * self._compute_buffer_ctx_bytes( + ctx, _effective_ubatch, cache_type_kv + ) + # Layer-split compute buffer (one lump; tensor mode reserves it # per device in _plan_tensor_parallel). Context-independent, so # fold it into the model footprint for the branches below. Falls @@ -5137,7 +5363,6 @@ class LlamaCppBackend: _compute_buffer_pipeline = ( self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024 ) - model_size_fit = model_size + _compute_buffer_pipeline # Layer split adds a fixed per-device overhead on every GPU. The # folded buffer covers one device; reserve the extra devices' @@ -5145,9 +5370,6 @@ class LlamaCppBackend: # (k=1 adds nothing). _pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024 - def _subset_model_size(n_gpus: int) -> int: - return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes - # Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx: # honor it, cap only if it fits no combination. Auto (native): # prefer fewer GPUs with reduced context (multi-GPU is slower). @@ -5174,6 +5396,21 @@ class LlamaCppBackend: else 0.0 ) _pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve + + # Charge the soft overhead _CTX_FIT_VRAM_FRACTION under-covers on tight + # tiers, gated so plain dense loads (#5106) only pay the CUDA-ctx base. + # CUDA/cuBLAS context is discrete-GPU only (not Metal); the mmproj and + # MTP draft-graph buffers exist on every backend. + _soft_overhead = self._CUDA_CONTEXT_RESERVE_BYTES if gpus else 0 + if effective_is_vision and mmproj_size > 0: + _soft_overhead += int(mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0)) + if _mtp_reserves_gpu: + _soft_overhead += self._MTP_DRAFT_COMPUTE_BYTES + model_size_fit = model_size + _compute_buffer_pipeline + _soft_overhead + + def _subset_model_size(n_gpus: int) -> int: + return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes + # Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below. _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) @@ -5278,7 +5515,9 @@ class LlamaCppBackend: _tp_flat_mtp, _mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048), ) - _tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024) + _tp_required_mib = (model_size + _tp_mtp_floor + _soft_overhead) / ( + 1024 * 1024 + ) if _tp_weight_budget_mib <= _tp_required_mib: logger.info( "Tensor parallelism requested but the pooled VRAM " @@ -5327,6 +5566,7 @@ class LlamaCppBackend: max_target_ctx = self._context_length or target_ctx, total_by_idx = total_by_idx, n_ubatch = _effective_ubatch, + soft_overhead_bytes = _soft_overhead, ) use_fit = False elif gpus and self._can_estimate_kv() and effective_ctx > 0: @@ -5351,6 +5591,9 @@ class LlamaCppBackend: # budget so the fit and the check below agree. pool_budget = _pool_budget_mib(subset, _cap_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( native_ctx_for_cap, pool_budget, @@ -5359,13 +5602,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: best_cap = max(best_cap, capped) if best_cap > 0: @@ -5386,13 +5632,18 @@ class LlamaCppBackend: effective_ctx, cache_type_kv, n_parallel = n_parallel ) + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) ) + # The compute buffer is replicated on every device in a + # layer split; fold it into the per-device reserve so a + # multi-GPU pin sizes each card for its own copy. gpu_indices, use_fit = self._select_gpus( requested_total, gpus, usable_fraction = _pin_fraction, total_by_idx = total_by_idx, - per_device_overhead_bytes = _pipeline_overhead_bytes, + per_device_overhead_bytes = _pipeline_overhead_bytes + + _cc_bytes(effective_ctx), min_gpus = _layer_min_gpus, ) # No silent shrink: effective_ctx stays == requested_ctx. @@ -5423,6 +5674,9 @@ class LlamaCppBackend: subset = ranked[:n_gpus] pool_budget = _pool_budget_mib(subset, pin_fraction) _ms = _subset_model_size(n_gpus) + # Compute buffer is replicated per device in a layer + # split, so scale the context term by the subset size. + _cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n) capped = self._fit_context_to_vram( effective_ctx, pool_budget, @@ -5431,13 +5685,16 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_sub, budget_frac = 1.0, total_mib = None, ) kv = self._estimate_kv_cache_bytes( capped, cache_type_kv, n_parallel = n_parallel ) - footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024) + footprint_mib = ( + _ms + kv + _mtp_bytes(capped) + _cc_sub(capped) + ) / (1024 * 1024) if footprint_mib <= pool_budget: effective_ctx = capped gpu_indices = sorted(idx for idx, _ in subset) @@ -5460,6 +5717,7 @@ class LlamaCppBackend: _subset_model_size(n_gpus) + kv + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx, n_gpus) ) / (1024 * 1024) if footprint_mib <= _pool_budget_mib(subset, pin_fraction): gpu_indices = sorted(idx for idx, _ in subset) @@ -5514,6 +5772,7 @@ class LlamaCppBackend: n_parallel = n_parallel, mtp_engaged = _mtp_reserves_gpu, mtp_overhead_fn = mtp_overhead_fn, + compute_ctx_bytes_fn = _cc_bytes, budget_frac = 1.0, total_mib = None, ) @@ -5523,6 +5782,7 @@ class LlamaCppBackend: cap, cache_type_kv, n_parallel = n_parallel ) + _mtp_bytes(cap) + + _cc_bytes(cap) ) / (1024 * 1024) # Fit returns the request unchanged when it fits OR weights # exceed budget; only the latter over-commits, so floor to 4096. @@ -5538,6 +5798,48 @@ class LlamaCppBackend: if not explicit_ctx: effective_ctx = max_available_ctx + # Prefer fewer serving slots on GPU over --fit on offload: when the extra + # --parallel slots push the footprint past the pin budget, llama-server + # offloads layers to host and decode collapses ~3x (#6718). Retry the fit + # at fewer slots, keeping the largest count that stays fully on GPU and the + # chosen context. Skips tensor mode / Metal / KV-inestimable paths. + if ( + use_fit + and n_parallel > 1 + and gpus + and self._can_estimate_kv() + and effective_ctx > 0 + ): + # Slot-independent footprint (folded compute buffer swapped out so the + # helper re-adds a slot-sized one per candidate). + _base_footprint = ( + model_size_fit + - _compute_buffer_pipeline + + _mtp_bytes(effective_ctx) + + _cc_bytes(effective_ctx) + ) + _gi_slots, _uf_slots, _slots = self._slots_that_fit_on_gpu( + n_parallel, + effective_ctx, + gpus, + total_by_idx, + _base_footprint, + cache_type_kv, + _pin_fraction, + _pipeline_overhead_bytes + _cc_bytes(effective_ctx), + _layer_min_gpus, + _effective_ubatch, + ) + if not _uf_slots: + logger.info( + "Serving slots reduced %d -> %d to keep the model on GPU " + "(avoid --fit offload) at context %d.", + n_parallel, + _slots, + effective_ctx, + ) + gpu_indices, use_fit, n_parallel = _gi_slots, False, _slots + # MTP reserve at the final context, for the logs below. _mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0 if _mtp_will_engage: @@ -5636,8 +5938,10 @@ class LlamaCppBackend: if use_fit: cmd.extend(["--fit", "on"]) elif gpu_indices is not None: - # Fits on selected GPU(s) -- offload all layers - cmd.extend(["-ngl", "-1"]) + # Fits on selected GPU(s) -- force all layers on GPU. --fit off is + # required: without it llama.cpp's default --fit on second-guesses + # and offloads ~1 GB at --parallel 4 even though the model fits. + cmd.extend(["-ngl", "-1", "--fit", "off"]) fully_gpu_offloaded = True server_caps = self.probe_server_capabilities(binary) @@ -6022,6 +6326,33 @@ class LlamaCppBackend: _split_axis_crash = self._is_tensor_split_assert( "\n".join(self._stdout_lines[-50:]) ) + if ( + _spawn_attempt == 0 + and fully_gpu_offloaded + and _startup_crashed + and not _split_axis_crash + ): + # We forced --fit off because Studio's (conservative) VRAM + # math placed the model fully on GPU. A startup crash here + # means that estimate was optimistic, so fall back to --fit + # on and let llama.cpp offload rather than fail the load. + logger.warning( + "llama-server crashed during startup (exit code %s) " + "with forced --fit off; the fit estimate was optimistic, " + "retrying once with --fit on so it can offload. " + "Crash log: %s", + self._process.returncode, + self._llama_log_path, + ) + # Flip Studio's own --fit off (added first, before any + # user extra args) to on; a user's later --fit still wins + # by last-arg. Defensive: if absent, the default is already + # --fit on, so leave it. + _run = list(run_cmd) + if "--fit" in _run: + _run[_run.index("--fit") + 1] = "on" + run_cmd = _run + continue if ( _spawn_attempt == 0 and _fit_retry_allowed @@ -7166,6 +7497,13 @@ class LlamaCppBackend: resolved_roots: list[Path] = [] for root in install_roots: try: + # A --with-llama-cpp-dir local link (symlink/junction) + # resolves into the user's own checkout. Adding it would let + # us treat the user's externally-launched llama-server as our + # orphan and kill it, so leave such roots out of the + # allowlist (we forgo orphan-reaping for local-link installs). + if _is_external_link(root): + continue resolved_roots.append(root.resolve()) except OSError: pass @@ -7475,6 +7813,10 @@ class LlamaCppBackend: time.sleep(interval) + # Leave a marker so _classify_llama_start_failure tells a live but + # never-healthy load (too large, or a proxy hijacking the loopback + # probe) apart from a bad GGUF (#5740). + self._stdout_lines.append(f"llama-server health check timed out after {timeout}s") logger.error(f"llama-server health check timed out after {timeout}s") return False diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py new file mode 100644 index 0000000000..c73134b4a2 --- /dev/null +++ b/studio/backend/core/inference/passthrough_healing.py @@ -0,0 +1,535 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tool-call healing for the client-tool passthrough. + +With server-side tools disabled (``unsloth run --disable-tools``, every +``unsloth start`` coding agent), requests carrying the client's own ``tools`` +bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small +GGUF models often emit their tool calls as TEXT (``{...}``, +Gemma ``<|tool_call>...``, ```` XML) instead of structured +``tool_calls`` -- on the passthrough that text reaches the agent as prose and +the turn dies. This module promotes such text back into structured calls on the +RESPONSE side only: the upstream request body is never touched, no extra +generation is issued, so llama-server slot/KV-cache reuse is byte-identical. + +Healing only ever fires when the request declared client tools, and only +promotes calls whose function name exactly matches a declared tool. Promotion +removes EXACTLY the promoted calls' markup spans (the parser reports them): +undeclared calls, unparseable blocks, and suppressed alternate formats keep +every byte and relay as text, so healing can never silently delete model +output. Responses without a tool signal, requests without tools, and Studio's +own enable-tools loop are untouched. Per-request opt-out: +``auto_heal_tool_calls: false``. Process kill-switch: +``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``. +""" + +import json +import os +from collections.abc import Mapping +from typing import Any, Optional + +from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal +from core.inference.tool_loop_controller import coerce_tool_arguments +from core.tool_healing import parse_tool_calls_from_text + +# Read once at import (same convention as the other UNSLOTH_* switches). +_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1" +# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process +# default with UNSLOTH_TOOL_CALL_NUDGE=1 (e.g. an `unsloth run` operator). +_NUDGE_DEFAULT = os.environ.get("UNSLOTH_TOOL_CALL_NUDGE", "0") == "1" + + +def nudge_enabled(request_flag: Optional[bool]) -> bool: + return _NUDGE_DEFAULT if request_flag is None else bool(request_flag) + + +_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS) +# A suspected-but-unclosed tool block larger than this is declared a false +# alarm and flushed, bounding memory on a model rambling XML-lookalike text. +_MAX_HOLD_CHARS = 64 * 1024 + + +def heal_gate( + auto_heal: Optional[bool], + tools: Optional[list], + tool_choice: Any = None, +) -> Optional[set]: + """Return the declared client-tool name set when healing applies, else None. + + ``tools`` is the OpenAI-shaped list forwarded to llama-server + (``[{"type": "function", "function": {"name": ...}}, ...]``). The name set + doubles as the promotion allowlist so healed calls can never invent a tool + the client did not declare. + + ``tool_choice`` (OpenAI shape) constrains the allowlist so healing never + contradicts the request: ``"none"`` forbids tool calls outright (text-form + markup stays text), and a forced ``{"type": "function", "function": + {"name": N}}`` narrows promotion to that one function. ``"auto"`` / + ``"required"`` / absent keep the full declared set. + """ + if _HEALING_DISABLED or auto_heal is False: + return None + if tool_choice == "none": + return None + names = set() + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if isinstance(function, dict) and isinstance(function.get("name"), str): + names.add(function["name"]) + if isinstance(tool_choice, dict): + function = tool_choice.get("function") + forced = function.get("name") if isinstance(function, dict) else None + if isinstance(forced, str): + names &= {forced} + return names or None + + +def _tool_schemas_by_name(tools: Optional[list]) -> dict[str, Any]: + schemas: dict[str, Any] = {} + for tool in tools or []: + if not isinstance(tool, dict): + continue + function = tool.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + if isinstance(name, str): + schemas[name] = function.get("parameters") + return schemas + + +def _string_arg_key_from_schema(schema: Any) -> Optional[str]: + if not isinstance(schema, dict): + return None + properties = schema.get("properties") + required = schema.get("required") + if not isinstance(properties, dict) or not isinstance(required, list): + return None + required_names = [name for name in required if isinstance(name, str)] + if len(required_names) != 1: + return None + key = required_names[0] + + if key not in properties: + return None + prop_schema = properties.get(key) + if isinstance(prop_schema, dict): + prop_type = prop_schema.get("type") + if isinstance(prop_type, list): + if "string" not in prop_type: + return None + elif prop_type is not None and prop_type != "string": + return None + return key + + +def _coerce_promoted_arguments( + raw_args: Any, tool_name: str, tool_schemas: Optional[dict] +) -> Optional[dict]: + if isinstance(raw_args, Mapping): + return dict(raw_args) + if isinstance(raw_args, str): + try: + parsed = json.loads(raw_args) + if isinstance(parsed, Mapping): + return dict(parsed) + except (json.JSONDecodeError, ValueError): + pass + if tool_schemas is not None: + key = _string_arg_key_from_schema(tool_schemas.get(tool_name)) + return {key: raw_args} if key else None + coerced = coerce_tool_arguments(raw_args, heal = True, tool_name = tool_name) + return coerced.arguments + + +def _promote( + calls: list, + allowed_tools: set, + id_offset: int = 0, + tool_schemas: Optional[dict] = None, +) -> list: + """Filter parsed calls to declared tools and normalize their arguments. + + Bare string arguments on the client-tool passthrough use the declared + schema's single required string property. If the schema is ambiguous, the + call stays text instead of inventing a generic key. + """ + promoted = [] + for call in calls: + function = call.get("function") if isinstance(call, dict) else None + name = function.get("name") if isinstance(function, dict) else None + if name not in allowed_tools: + continue + arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas) + if arguments is None: + continue + promoted.append( + { + "id": f"call_{id_offset + len(promoted)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments, ensure_ascii = False), + }, + } + ) + return promoted + + +def _remove_spans(text: str, spans: list) -> str: + """Text with the given non-overlapping, sorted (start, end) ranges removed.""" + pieces = [] + pos = 0 + for start, end in spans: + pieces.append(text[pos:start]) + pos = end + pieces.append(text[pos:]) + return "".join(pieces) + + +def heal_openai_message_events( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> Optional[list]: + if not isinstance(msg, dict) or msg.get("tool_calls"): + return None + content = msg.get("content") + if not isinstance(content, str) or not has_tool_signal(content): + return None + parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + events: list = [] + pos = 0 + call_count = 0 + for call, (start, end) in zip(parsed, spans): + promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas) + if promoted: + if content[pos:start]: + events.append(("text", content[pos:start])) + events.append(("tool_call", promoted[0])) + call_count += 1 + else: + events.append(("text", content[pos:end])) + pos = end + if not call_count: + return None + if content[pos:]: + events.append(("text", content[pos:])) + return events + + +def heal_openai_message( + msg: dict, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Promote text-form tool calls in a non-streaming OpenAI message. In place. + + No-op (returns False) unless the message has NO structured ``tool_calls`` + (grammar mode already worked when it does) and its content carries a tool + signal that parses into at least one declared call. Only the promoted + calls' markup spans are removed from the content; undeclared calls and + anything the parser did not consume stay in the text byte-intact. + """ + events = heal_openai_message_events(msg, allowed_tools, tools) + if not events: + return False + calls = [value for kind, value in events if kind == "tool_call"] + content = "".join(value for kind, value in events if kind == "text").strip() + msg["tool_calls"] = calls + # OpenAI requires content = null on a pure tool-call turn. + msg["content"] = content or None + return True + + +def _earliest_signal(buffer: str) -> int: + best = -1 + for signal in TOOL_XML_SIGNALS: + index = buffer.find(signal) + if index >= 0 and (best < 0 or index < best): + best = index + return best + + +def _closed_signal_span(buffer: str) -> Optional[tuple[int, int]]: + spans = [] + for open_tag, close_tag in ( + ("", ""), + ("<|tool_call>", ""), + (""), + ): + start = buffer.find(open_tag) + if start < 0: + continue + end = buffer.find(close_tag, start) + if end >= 0: + spans.append((start, end + len(close_tag))) + return min(spans, key = lambda span: span[0]) if spans else None + + +def _partial_signal_suffix(buffer: str) -> int: + """Length of the longest buffer suffix that is a proper prefix of a signal.""" + for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1): + tail = buffer[-length:] + if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS): + return length + return 0 + + +class StreamToolCallHealer: + """Buffer-and-repair state machine for streamed passthrough content. + + ``feed(text)`` / ``finalize()`` yield ``("text", str)`` events for content + to relay and ``("tool_call", dict)`` events carrying an OpenAI-shaped call + (string ``function.arguments``). Normal prose is forwarded immediately; only + a trailing partial-signal window (< max signal length) or a suspected tool + block is ever withheld, so streaming latency stays bounded. A false alarm + (the buffer can no longer become a parseable declared call) flushes the held + text verbatim. + """ + + def __init__( + self, + allowed_tools: set, + tools: Optional[list] = None, + ) -> None: + self._allowed = set(allowed_tools) + + self._tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + self._buffer = "" + self._holding = False + self._id_offset = 0 + # Structured delta.tool_calls seen upstream: grammar mode already + # worked, so healing goes dormant and text relays verbatim. + self.dormant = False + + @property + def healed(self) -> bool: + return self._id_offset > 0 + + def structured_tool_call_seen(self) -> list: + """Go dormant; flush anything held so no text is swallowed.""" + self.dormant = True + held, self._buffer, self._holding = self._buffer, "", False + return [("text", held)] if held else [] + + def feed(self, text: str) -> list: + if self.dormant: + return [("text", text)] if text else [] + self._buffer += text + return self._drain() + + def _drain(self) -> list: + events: list = [] + while True: + if not self._holding: + start = _earliest_signal(self._buffer) + if start >= 0: + if start: + events.append(("text", self._buffer[:start])) + self._buffer = self._buffer[start:] + self._holding = True + else: + keep = _partial_signal_suffix(self._buffer) + emit = self._buffer[: len(self._buffer) - keep] + if emit: + events.append(("text", emit)) + self._buffer = self._buffer[len(self._buffer) - keep :] + return events + # HOLD: handle the FIRST complete block per pass so events keep + # document order (a later declared call must not overtake an + # earlier undeclared one flushing as text). + parsed, spans = parse_tool_calls_from_text( + self._buffer, + id_offset = self._id_offset, + allow_incomplete = False, + with_spans = True, + ) + if not parsed: + closed_span = _closed_signal_span(self._buffer) + if closed_span: + _start, end = closed_span + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[end:] + self._holding = False + continue + if len(self._buffer) > _MAX_HOLD_CHARS: + events.append(("text", self._buffer)) + self._buffer = "" + self._holding = False + continue + return events + start, end = spans[0] + promoted = _promote( + [parsed[0]], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if start: + events.append(("text", self._buffer[:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + # Drop exactly the promoted markup span; everything else + # (leading text, later blocks) stays and is rescanned. + self._buffer = self._buffer[end:] + else: + # Undeclared or unusable name: its markup is DATA, flush it + # (and anything before it) verbatim, then rescan the rest. + events.append(("text", self._buffer[:end])) + self._buffer = self._buffer[end:] + self._holding = False + + def finalize(self) -> list: + """End of stream: last-chance heal of the residue, else flush it. + + Events keep document order; only the promoted calls' markup spans are + dropped, every other residue byte flushes as text. + """ + if not self._buffer: + return [] + residue, self._buffer = self._buffer, "" + holding, self._holding = self._holding, False + if self.dormant or not holding: + return [("text", residue)] + parsed, spans = parse_tool_calls_from_text( + residue, + id_offset = self._id_offset, + allow_incomplete = True, + with_spans = True, + ) + events: list = [] + pos = 0 + any_promoted = False + for call, (start, end) in zip(parsed, spans): + promoted = _promote( + [call], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + if residue[pos:start]: + events.append(("text", residue[pos:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + any_promoted = True + else: + events.append(("text", residue[pos:end])) + pos = end + if not any_promoted: + return [("text", residue)] + tail = residue[pos:].strip() + if tail: + events.append(("text", tail)) + return events + + +def _first_choice_message(data: Any) -> Optional[dict]: + """First-choice message dict of a non-streaming chat response, else None. + + Upstream error bodies can carry ``"message": null`` (or no choices at all), + so never assume the shape: a non-dict message means "nothing to heal". + """ + try: + message = data["choices"][0]["message"] + except (KeyError, IndexError, TypeError): + return None + return message if isinstance(message, dict) else None + + +def _last_assistant_text(data: Any) -> str: + """First-choice assistant content of a non-streaming chat response, or ''.""" + message = _first_choice_message(data) + content = message.get("content") if message else None + return content if isinstance(content, str) else "" + + +def _heal_would_promote( + text: str, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """Whether ``heal_openai_message`` would promote at least one call.""" + parsed = parse_tool_calls_from_text(text, allow_incomplete = True) + tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None + return bool(_promote(parsed, allowed_tools, tool_schemas = tool_schemas)) + + +def response_has_promotable_calls( + data: Any, + allowed_tools: set, + tools: Optional[list] = None, +) -> bool: + """True when a non-streaming chat response carries a usable tool call + (structured naming a DECLARED tool, or text-form that healing would + promote). Used to decide whether a nudge retry actually improved on the + original response; a hallucinated undeclared call is not an improvement.""" + message = _first_choice_message(data) + if not message: + return False + tool_calls = message.get("tool_calls") + if tool_calls: + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST one), so a mixed + # response with a single hallucinated name could still hand the client + # an undeclared tool. + return all( + isinstance(tc, dict) + and isinstance(tc.get("function"), dict) + and tc["function"].get("name") in allowed_tools + for tc in tool_calls + ) + text = message.get("content") + if not isinstance(text, str): + return False + return _heal_would_promote(text, allowed_tools, tools) + + +def nudge_should_retry( + data: Any, + allowed_tools: Optional[set], + tools: Optional[list] = None, +) -> bool: + """True when the first response tried to call a tool but nothing healed. + + Trigger only on: healing enabled (allowed_tools set), zero structured + calls, a tool signal present in the text, and zero promotable calls -- the + exact failure a single re-ask can fix. Clean prose never retries. + """ + if not allowed_tools: + return False + message = _first_choice_message(data) + if not message or message.get("tool_calls"): + return False + text = message.get("content") + if not isinstance(text, str) or not has_tool_signal(text): + return False + return not _heal_would_promote(text, allowed_tools, tools) + + +def nudge_messages(data: Any, allowed_tools: set) -> list: + """The two-message suffix appended for the single nudge retry. + + The retry body is the original body plus this suffix, so the prompt prefix + is byte-identical and llama-server's slot/prefix cache is reused (same + shape as the enable-tools loop's reprompt). + """ + tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool" + return [ + {"role": "assistant", "content": _last_assistant_text(data)}, + { + "role": "user", + "content": ( + "You have access to the declared tools. If a tool is needed to " + f"complete the action you described, call {tool_hint} now using the " + "native tool-call format with valid JSON arguments, not prose. If no " + "tool is needed, provide the final answer directly." + ), + }, + ] diff --git a/studio/backend/core/rag/config.py b/studio/backend/core/rag/config.py index 54a224d081..2de32a68e4 100644 --- a/studio/backend/core/rag/config.py +++ b/studio/backend/core/rag/config.py @@ -6,8 +6,10 @@ from __future__ import annotations import os +import re -EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "unsloth/bge-small-en-v1.5") +DEFAULT_EMBEDDING_MODEL = "unsloth/bge-small-en-v1.5" +EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL) # Under bge's 512 limit, leaving headroom for the 2 special tokens (else overflow: # llama-server 500s, ST truncates). Keep <= embedder_max - ~12. CHUNK_TOKENS = int(os.environ.get("RAG_CHUNK_TOKENS", "500")) @@ -66,6 +68,43 @@ OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048")) # wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes # the vectors, so the index must be rebuilt. EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto") + + +def effective_embedding_model() -> str: + """The embedding model actually in use: the persisted Settings override when + one is stored, else ``EMBEDDING_MODEL`` (env/default). Read at call time so a + Settings change applies without a restart.""" + try: + from utils.embedding_model_settings import get_rag_embedding_model + return get_rag_embedding_model() + except Exception: # noqa: BLE001 - settings store unavailable (tests, early boot) + return EMBEDDING_MODEL + + +def _names_gguf(model: str) -> bool: + """True when "gguf" appears as a whole name segment, so plain substrings + like "bigguf" don't count.""" + return "gguf" in re.split(r"[^a-z0-9]+", model.lower()) + + +def effective_gguf_repo() -> str: + """GGUF repo for the llama-server backend, tracking the effective model. + + An explicit ``RAG_EMBED_GGUF_REPO`` env always wins. Otherwise any custom + model (saved in Settings or via ``RAG_EMBEDDING_MODEL``) maps to its + ``-GGUF`` companion repo (the unsloth convention the default pair follows), + or is used as-is when it already names a GGUF repo. + """ + if "RAG_EMBED_GGUF_REPO" in os.environ: + return EMBED_GGUF_REPO + model = effective_embedding_model() + if model == DEFAULT_EMBEDDING_MODEL: + return EMBED_GGUF_REPO + if _names_gguf(model): + return model + return f"{model}-GGUF" + + # llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this # tiny model) and exact vs fp32, for ~30MB more on disk. EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF") diff --git a/studio/backend/core/rag/embed_llama_server.py b/studio/backend/core/rag/embed_llama_server.py index f53478463c..46a282c939 100644 --- a/studio/backend/core/rag/embed_llama_server.py +++ b/studio/backend/core/rag/embed_llama_server.py @@ -55,9 +55,15 @@ class LlamaServerBackend: self._port: int | None = None self._stdout_lines: list[str] = [] self._stdout_thread: threading.Thread | None = None + # No lock: probes are idempotent (a duplicate 1-text encode is benign) + # and dim() -> encode() -> _ensure_ready() -> _resolve_model_path() can + # re-enter on a mid-probe model change, which would self-deadlock a + # non-reentrant lock held across the probe. self._dim: int | None = None - self._dim_lock = threading.Lock() self._model_path: str | None = None + # Effective GGUF repo the cached path/dim belong to; a Settings change + # makes it stale, forcing a re-resolve + respawn (see _ensure_ready). + self._model_repo: str | None = None self._binary: str | None = None # Sticky after an auto GPU start fails: later spawns stay on CPU. self._force_cpu = False @@ -114,24 +120,77 @@ class LlamaServerBackend: "RAG_EMBED_BACKEND=llama-server requires an embeddings-capable build" ) + @staticmethod + def _resolve_local_gguf(model: str) -> str | None: + """A custom model may be a local .gguf file or a directory holding one; + resolve it without the hub. None when the value is not a local path.""" + p = Path(model).expanduser() + if p.is_file() and p.suffix.lower() == ".gguf": + return str(p) + if p.is_dir(): + files = [ + f + for f in p.iterdir() + if f.suffix.lower() == ".gguf" and "mmproj" not in f.name.lower() + ] + if not files: + raise RuntimeError(f"no .gguf file found in local model dir {model!r}") + variant = config.EMBED_GGUF_VARIANT.lower() + match = [f for f in files if variant in f.name.lower()] or files + return str(sorted(match, key = lambda f: len(f.name))[0]) + return None + def _resolve_model_path(self) -> str: """Download (or cache-hit) the variant-matching, non-mmproj GGUF embedder, - returning its local path.""" - if self._model_path is not None: + returning its local path. Re-resolves when the effective repo changed (a + custom model was saved in Settings).""" + # Captured once: if the setting changes mid-download, the path must stay + # tagged with the repo it was resolved FOR, so _current() sees the new + # setting as stale and respawns instead of serving the old model. + desired = config.effective_gguf_repo() + if self._model_path is not None and self._model_repo == desired: + return self._model_path + local = self._resolve_local_gguf(config.effective_embedding_model()) + if local is not None: + self._model_path = local + self._model_repo = desired + self._dim = None return self._model_path from huggingface_hub import hf_hub_download, list_repo_files - repo = config.EMBED_GGUF_REPO token = os.environ.get("HF_TOKEN") or None - files = [f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf")] - files = [f for f in files if "mmproj" not in f.lower()] + # A custom model derives its "-GGUF" companion repo; when that guess does + # not exist, the model repo itself may host the .gguf files. + repo = desired + candidates = [repo] + model = config.effective_embedding_model() + if model != repo: + candidates.append(model) + files: list[str] = [] + errors: list[str] = [] + for candidate in candidates: + try: + files = [ + f + for f in list_repo_files(candidate, token = token) + if f.lower().endswith(".gguf") and "mmproj" not in f.lower() + ] + except Exception as e: # noqa: BLE001 - missing/gated repo -> next candidate + errors.append(f"{candidate!r}: {e}") + continue + if files: + repo = candidate + break + errors.append(f"{candidate!r}: no .gguf files") if not files: - raise RuntimeError(f"no .gguf file found in embedder repo {repo!r}") + raise RuntimeError("no .gguf embedder found; tried " + "; ".join(errors)) variant = config.EMBED_GGUF_VARIANT.lower() match = [f for f in files if variant in f.lower()] or files filename = sorted(match, key = len)[0] logger.info("resolving GGUF embedder %s/%s", repo, filename) self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token) + self._model_repo = desired + self._dim = None return self._model_path # Min free VRAM (MiB) for the embedder; below this, auto stays on CPU. @@ -316,13 +375,19 @@ class LlamaServerBackend: def _process_alive(self) -> bool: return self._process is not None and self._process.poll() is None + def _current(self) -> bool: + """Alive AND serving the effective repo (a Settings model change makes a + live server stale).""" + return self._process_alive() and self._model_repo == config.effective_gguf_repo() + def _ensure_ready(self) -> None: - """Guarantee a live server, (re)spawning if needed. Double-checked so the - alive path takes no lock; self-heals after the chat reaper kills us.""" - if self._process_alive(): + """Guarantee a live server on the effective model, (re)spawning if needed. + Double-checked so the current path takes no lock; self-heals after the + chat reaper kills us and re-resolves after a Settings model change.""" + if self._current(): return with self._lifecycle_lock: - if self._process_alive(): + if self._current(): return self._kill_process() self._spawn() @@ -424,14 +489,18 @@ class LlamaServerBackend: return arr def dim(self, *, model_name = None) -> int: - """Embedding width, probed once via a 1-text encode and cached.""" - if self._dim is not None: - return self._dim - with self._dim_lock: - if self._dim is None: - vec = self.encode(["x"], normalize = False) - self._dim = int(vec.shape[1]) - return self._dim + """Embedding width, probed via a 1-text encode and cached per model + (_resolve_model_path clears it when the effective repo changes). + Unlocked: concurrent probes are benign, and locking would deadlock when + the probe's encode respawns onto a changed model (see __init__).""" + self._ensure_ready() + cached = self._dim + if cached is not None: + return cached + vec = self.encode(["x"], normalize = False) + width = int(vec.shape[1]) + self._dim = width + return width def warm(self, *, model_name = None) -> None: """Start the server and probe dim off the request path.""" diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 4c8d690302..345b4dd853 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -67,7 +67,7 @@ def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" global _model, _name - name = model_name or config.EMBEDDING_MODEL + name = model_name or config.effective_embedding_model() with _lock: if _model is None or _name != name: _install_torchao_stub_once() diff --git a/studio/backend/core/rag/ingestion.py b/studio/backend/core/rag/ingestion.py index 04365ab76b..cba076f1be 100644 --- a/studio/backend/core/rag/ingestion.py +++ b/studio/backend/core/rag/ingestion.py @@ -151,6 +151,19 @@ def _ocr_scanned_pages( return out, ocred +def _replace_old_document(conn, replaces: tuple[str, str | None] | None, keep_path: str) -> None: + """Drop the document this ingestion replaced (stale embedder / empty prior + ingest), called only after the replacement completed successfully.""" + if replaces is None: + return + old_id, old_path = replaces + try: + store.delete_document(conn, old_id) + _remove_upload(old_path, keep_path = keep_path) + except Exception: # noqa: BLE001 - the new document is already live + logger.warning("failed to remove replaced document %s", old_id, exc_info = True) + + def _run( job_id: str, document_id: str, @@ -159,6 +172,7 @@ def _run( model_name: str | None, ocr: bool | None = None, caption: bool | None = None, + replaces: tuple[str, str | None] | None = None, ) -> None: conn = rag_db.get_connection() try: @@ -213,6 +227,7 @@ def _run( ) if not chunks: store.set_document_status(conn, document_id, "completed", num_chunks = 0) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": 0}) return @@ -233,6 +248,7 @@ def _run( _progress(conn, job_id, "storing", 0.9) store.add_chunks(conn, scope, document_id, chunks, vectors, regions) store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks)) + _replace_old_document(conn, replaces, stored_path) _set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0) _emit(job_id, {"type": "complete", "num_chunks": len(chunks)}) @@ -274,17 +290,32 @@ def start_ingestion( sha = _sha256_file(stored_path) conn = rag_db.get_connection() try: + effective_model = model_name or config.effective_embedding_model() + # (old_document_id, old_stored_path) replaced by this upload; deleted by + # the worker only after the replacement completes, so a failed re-index + # never destroys the still-searchable original. + replaces: tuple[str, str | None] | None = None existing = store.document_by_hash(conn, scope, sha) if existing is not None: doc = store.get_document(conn, existing) empty_completed = ( doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks") ) - if empty_completed: + # Vectors from a different embedder are stale; re-uploading must + # re-index, not dedupe. NULL (legacy rows) is assumed current. Only + # completed rows are replaceable: a pending/running duplicate has a + # live worker whose writes must not land on a deleted document. + stale_model = ( + doc is not None + and doc.get("status") == "completed" + and doc.get("embedding_model") is not None + and doc.get("embedding_model") != effective_model + ) + if empty_completed or stale_model: # A prior ingest of identical bytes yielded zero chunks (e.g. a scanned - # PDF uploaded before a vision model loaded). Re-ingest, don't dedupe. - store.delete_document(conn, existing) - _remove_upload(doc.get("stored_path"), keep_path = stored_path) + # PDF uploaded before a vision model loaded), or was embedded with a + # different model. Re-ingest, don't dedupe. + replaces = (existing, doc.get("stored_path")) else: job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0) _remove_upload(stored_path) @@ -310,6 +341,7 @@ def start_ingestion( project_id = project_id, status = "pending", stored_path = stored_path, + embedding_model = effective_model, ) job_id = _new_job(conn, document_id, scope) finally: @@ -319,7 +351,10 @@ def start_ingestion( _jobs[job_id] = queue.Queue() threading.Thread( target = _run, - args = (job_id, document_id, scope, stored_path, model_name, ocr, caption), + # effective_model (not the raw model_name) pins the embedder for the + # whole job: a Settings change mid-ingestion must not switch tokenizer + # or embedder between batches of one document. + args = (job_id, document_id, scope, stored_path, effective_model, ocr, caption, replaces), daemon = True, ).start() return document_id, job_id diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index ba248cf9a6..9afddf1d9e 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -12,6 +12,7 @@ from __future__ import annotations import logging import os +import re from dataclasses import dataclass from html.parser import HTMLParser @@ -69,6 +70,39 @@ def _html(raw: str) -> list[Page]: return [_page("\n".join(parser.out), 1)] +# pymupdf4llm rebuilds text from positioned glyphs, which mangles complex-shaping +# scripts (RTL Arabic/Hebrew emerge as shaped Presentation Forms, Indic matras drop to +# U+FFFD) and can silently drop most of a heavy-RTL page. When Markdown trips these +# signals we fall back to PyMuPDF's logical-order get_text(). Thresholds mirror the chat +# extractor guard (unslothai/unsloth#5351 review). +_SHAPED_PRESENTATION_FORMS = re.compile("[\ufb1d-\ufdff\ufe70-\ufefc]") +_PDF_FALLBACK_MIN_BAD_GLYPHS = 5 +_PDF_FALLBACK_BAD_GLYPH_RATIO = 0.0005 +_PDF_INCOMPLETE_RATIO = 0.75 +_PDF_INCOMPLETE_MIN_LETTERS = 200 + + +def _markdown_corrupted(text: str) -> bool: + """True when pymupdf4llm's glyph reconstruction mangled the text: shaped RTL + Presentation Forms or U+FFFD replacements above a small floor/ratio (so a lone + legitimate shaped glyph does not force the fallback).""" + if not text: + return False + threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text)) + shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text)) + return shaped > threshold or text.count("\ufffd") > threshold + + +def _markdown_incomplete(markdown: str, plain: str) -> bool: + """True when ``markdown`` holds far fewer letters than the raw ``get_text`` layer -- a + coarse guard for heavy-RTL pages pymupdf4llm silently drops without shaped glyphs.""" + plain_letters = sum(1 for c in plain if c.isalnum()) + if plain_letters < _PDF_INCOMPLETE_MIN_LETTERS: + return False + markdown_letters = sum(1 for c in markdown if c.isalnum()) + return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters + + def _pdf_markdown(doc) -> list[str] | None: """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index i maps to page i+1. Returns None when the lib is missing, extraction fails, or the @@ -100,9 +134,19 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: try: md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None for i, page in enumerate(doc): - # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval); - # fall back to plain text when Markdown is off, unavailable, or empty here. - text = (md[i] if md else "") or page.get_text("text") or "" + plain = page.get_text("text") or "" + candidate = md[i] if md else "" + # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval), + # but drop to PyMuPDF's logical-order text when Markdown is off/empty or when + # pymupdf4llm mangled it (RTL/Indic) or dropped most of the page. + if ( + candidate + and not _markdown_corrupted(candidate) + and not _markdown_incomplete(candidate, plain) + ): + text = candidate + else: + text = plain pages.append(_page(text, i + 1)) if want_images: for img in page.get_images(full = True): @@ -308,12 +352,61 @@ def render_pdf_pages( doc.close() +def _docx_table_rows(table) -> list[str]: + """Each row as pipe-joined cell text (the locator splits anchors on pipes). + Columns stay aligned to the layout grid (merged cells fill their spanned slots, + skipped leading/trailing grid columns become empty fields). Cells are walked in + document order so a nested table, and any text after it, flattens in place.""" + from docx.table import Table + from docx.text.paragraph import Paragraph + + rows: list[str] = [] + seen: set = set() # already emitted; dedups merges spanning columns or rows + for row in table.rows: + cells: list[str] = [""] * getattr(row, "grid_cols_before", 0) + trailing: list[str] = [] # nested rows + any post-nested text, kept in order + for cell in row.cells: + # A merged cell shares one across the columns and rows it spans: + # emit its text once, then placeholders, so columns and rows stay aligned. + if cell._tc in seen: + cells.append("") + continue + seen.add(cell._tc) + # Paragraph text before the first nested table is the aligned field; the + # nested table and anything after it flatten below the row, in order. + field: list[str] = [] + after_table = False + for item in cell.iter_inner_content(): + if isinstance(item, Table): + after_table = True + trailing.extend(_docx_table_rows(item)) + elif isinstance(item, Paragraph): + text = " ".join(item.text.split()) # collapse in-cell newlines + if text: + (trailing if after_table else field).append(text) + cells.append(" ".join(field)) # empty cells kept so columns line up + cells.extend([""] * getattr(row, "grid_cols_after", 0)) + if any(c.strip() for c in cells): + rows.append(" | ".join(cells)) + rows.extend(trailing) + return rows + + def _docx(path: str) -> list[Page]: import docx + from docx.table import Table + from docx.text.paragraph import Paragraph document = docx.Document(path) - text = "\n".join(p.text for p in document.paragraphs) - return [_page(text, None)] + lines: list[str] = [] + # Walk body content in document order: paragraphs alone drop tables entirely. + for block in document.iter_inner_content(): + if isinstance(block, Paragraph): + if block.text.strip(): + lines.append(block.text) + elif isinstance(block, Table): + lines.extend(_docx_table_rows(block)) + return [_page("\n".join(lines), None)] def parse(path: str, *, want_images: bool = False): diff --git a/studio/backend/core/rag/retrieval.py b/studio/backend/core/rag/retrieval.py index fe6a033a52..6f933e089e 100644 --- a/studio/backend/core/rag/retrieval.py +++ b/studio/backend/core/rag/retrieval.py @@ -39,8 +39,12 @@ def retrieve_dense( model_name: str | None = None, ) -> list[Hit]: k = k or config.TOP_K_DENSE - vec = embeddings.encode([query], model_name = model_name, normalize = True)[0] - return [Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k)] + effective = model_name or config.effective_embedding_model() + vec = embeddings.encode([query], model_name = effective, normalize = True)[0] + return [ + Hit(cid, s, dense_score = s) + for cid, s in store.search_dense(conn, scope, vec, k, embedding_model = effective) + ] def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]: diff --git a/studio/backend/core/rag/store.py b/studio/backend/core/rag/store.py index 8e59c5fbf6..f9128d1715 100644 --- a/studio/backend/core/rag/store.py +++ b/studio/backend/core/rag/store.py @@ -109,11 +109,12 @@ def create_document( status: str = "pending", stored_path: str | None = None, document_id: str | None = None, + embedding_model: str | None = None, ) -> str: document_id = document_id or str(uuid.uuid4()) conn.execute( "INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, " - "status, stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?,?)", + "status, stored_path, created_at, embedding_model) VALUES(?,?,?,?,?,?,?,?,?,?,?)", ( document_id, scope, @@ -125,6 +126,7 @@ def create_document( status, stored_path, _now(), + embedding_model, ), ) conn.commit() @@ -261,20 +263,50 @@ def search_lexical(conn: sqlite3.Connection, scope, query: str, k: int): return [(r["chunk_id"], -r["s"]) for r in rows] -def search_dense(conn: sqlite3.Connection, scope, vector, k: int): +def search_dense( + conn: sqlite3.Connection, + scope, + vector, + k: int, + *, + embedding_model: str | None = None, +): """Cosine KNN over vec0 for one scope or several. Returns [(chunk_id, 1 - distance)]. vec0 KNN constrains its partition key by - equality, so multi-scope runs one query per scope and merges by score.""" + equality, so multi-scope runs one query per scope and merges by score. + ``embedding_model`` drops hits from documents indexed under a different + (same-width) model, whose vectors live in another space; NULL-model legacy + documents are assumed current, matching the ingestion dedupe rule.""" if not rag_db.vec_table_exists(conn): return [] + dim = rag_db.vec_table_dim(conn) + if dim is not None and dim != len(vector): + # Embedding model switched widths and nothing re-indexed yet; the stale + # table cannot answer new-model queries (vec0 errors on the MATCH). + return [] + # Over-fetch when filtering so stale-model hits don't starve the top-k. + fetch = k * 3 if embedding_model else k out: list[tuple[str, float]] = [] for s in _scopes(scope): rows = conn.execute( "SELECT chunk_id, distance FROM chunks_vec " "WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?", - (s, _f32(vector), k), + (s, _f32(vector), fetch), ).fetchall() out.extend((r["chunk_id"], 1.0 - r["distance"]) for r in rows) + if embedding_model and out: + ids = [cid for cid, _ in out] + placeholders = ",".join("?" * len(ids)) + valid = { + r["id"] + for r in conn.execute( + f"SELECT c.id FROM chunks c JOIN documents d ON d.id=c.document_id " + f"WHERE c.id IN ({placeholders}) " + f"AND (d.embedding_model IS NULL OR d.embedding_model=?)", + (*ids, embedding_model), + ).fetchall() + } + out = [t for t in out if t[0] in valid] out.sort(key = lambda t: t[1], reverse = True) return out[:k] diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fe26d48c7f..e8367ad08c 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -301,28 +301,28 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, -) -> list[dict]: + with_spans: bool = False, +): """Parse OpenAI-format tool calls from model text. Handles formats like: {"name":"web_search","arguments":{"query":"..."}} <|tool_call>call:web_search{query:"..."} ... + + With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]`` + is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup + in ``content`` (including its close tag when present), so a caller can + remove exactly the parsed markup and keep every other byte intact. """ tool_calls: list[dict] = [] - # Collect JSON- and Gemma-format candidates with their byte spans, then - # accept them in document order. Both order and spans matter: - # * tools execute in returned order, so a call appearing earlier in the - # text must be emitted first even across the two formats; - # * a tool-call marker INSIDE another call's argument string is data, not a - # call, so a candidate starting within an already accepted span is - # skipped (covers a JSON marker nested in a Gemma arg and a Gemma marker - # nested in a JSON arg alike, regardless of which format is outer). + call_spans: list[tuple] = [] + # Collect every supported call format with spans, then emit in document + # order. A marker inside another call's argument string is data, not a + # separate executable call. + parsed_items = [] # (start, span_end, name, arguments) candidates = [] # (start, brace_end, kind, match) for m in _TC_JSON_START_RE.finditer(content): - # A marker that begins inside an open value - # is that parameter's data, not its own call; skip it (same guard the - # XML-style parser below applies to nested = 0: + span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) + body = body[:close_idx] + elif not allow_incomplete: + continue + else: + body = _TC_FUNC_CLOSE_RE.sub("", body) + span_end = body_end + + arguments: dict = {} + param_starts = list(_TC_PARAM_START_RE.finditer(body)) + if len(param_starts) == 1: + pm = param_starts[0] + val = body[pm.end() :] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + continue + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[pm.group(1)] = val.strip() + else: + valid_params = True + for pidx, pm in enumerate(param_starts): + param_name = pm.group(1) + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + val = body[val_start:next_param] + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + valid_params = False + break + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) + arguments[param_name] = val.strip() + if not valid_params: + continue + + span_start = fm.start() + wrap_open = re.search(r"\s*$", content[:span_start]) + wrap_close = re.match(r"\s*", content[span_end:]) + if wrap_open and wrap_close: + span_start = wrap_open.start() + span_end += wrap_close.end() + parsed_items.append((span_start, span_end, func_name, json.dumps(arguments))) + + parsed_items.sort(key = lambda item: item[0]) + for start, span_end, name, arguments in parsed_items: tool_calls.append( { "id": f"call_{id_offset + len(tool_calls)}", @@ -369,77 +443,9 @@ def parse_tool_calls_from_text( "function": {"name": name, "arguments": arguments}, } ) - - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - if not allow_incomplete: - close_idx = body.rfind(_FUNC_CLOSE_TAG) - if close_idx < 0: - continue - body = body[:close_idx] - else: - body = _TC_FUNC_CLOSE_RE.sub("", body) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) + call_spans.append((start, span_end)) + if with_spans: + return tool_calls, call_spans return tool_calls diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 610af2472e..17dc1299ca 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1253,32 +1253,48 @@ def _adapt_for_mlx_vlm( return adapted -_MLX_STUDIO_OPTIM_MAP = { - "adamw_8bit": "adamw", - "paged_adamw_8bit": "adamw", - "adamw_bnb_8bit": "adamw", - "paged_adamw_32bit": "adamw", - "adamw_torch": "adamw", - "adamw_torch_fused": "adamw", - "adamw": "adamw", - "adafactor": "adafactor", - "sgd": "sgd", - "adam": "adam", - "muon": "muon", - "lion": "lion", -} _MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"} +# Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used +# only when mlx (Apple Silicon) is not importable so Studio config validation +# still works on non-MLX hosts. The zoo function stays the source of truth. +_MLX_STUDIO_ADAMW_ALIASES = frozenset( + ( + "adamw_8bit", + "paged_adamw_8bit", + "adamw_bnb_8bit", + "paged_adamw_32bit", + "adamw_torch", + "adamw_torch_fused", + "paged_adamw", + "adamw_32bit", + "adamw_hf", + "adamw_anyprecision", + "adamw_apex_fused", + ) +) +_MLX_STUDIO_NATIVE_OPTIMIZERS = ("adafactor", "adamw", "adam", "sgd", "muon", "lion") + + def _normalize_mlx_studio_optimizer(value): - raw = str(value or "adamw_8bit").strip().lower() try: - return _MLX_STUDIO_OPTIM_MAP[raw] - except KeyError: - supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP)) - raise ValueError( - f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}." - ) + from unsloth_zoo.mlx.trainer import _normalize_mlx_optimizer_name + return _normalize_mlx_optimizer_name(value or "adamw_8bit") + except (ImportError, ValueError): + # Missing mlx, or an older unsloth-zoo whose normalizer lacks CUDA/TRL + # aliases: map common adamw_* names locally so notebook defaults work. + opt = str(getattr(value, "value", value) or "adamw_8bit").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_STUDIO_ADAMW_ALIASES: + opt = "adamw" + if opt not in _MLX_STUDIO_NATIVE_OPTIMIZERS: + supported = ", ".join(_MLX_STUDIO_NATIVE_OPTIMIZERS) + raise ValueError( + f"Unsupported optimizer for MLX training: {value!r}. " + f"Supported optimizers: {supported}." + ) + return opt def _normalize_mlx_studio_scheduler(value): diff --git a/studio/backend/hub/services/models/folder_browser.py b/studio/backend/hub/services/models/folder_browser.py index 9b0b46509b..eb137127fb 100644 --- a/studio/backend/hub/services/models/folder_browser.py +++ b/studio/backend/hub/services/models/folder_browser.py @@ -27,6 +27,7 @@ from hub.utils.paths import ( studio_root, well_known_model_dirs, ) +from utils.paths.external_media import linux_run_media_mount_roots from hub.services.models.common import _safe_is_dir from hub.services.models.local_inventory import _resolve_hf_cache_dir @@ -175,6 +176,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -346,6 +349,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -485,6 +493,8 @@ def browse_folders_response( # Home first as the safe fallback. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default. try: _add_sug(_resolve_hf_cache_dir()) diff --git a/studio/backend/hub/storage/scan_folders.py b/studio/backend/hub/storage/scan_folders.py index 85f515da00..fdb15c7c3c 100644 --- a/studio/backend/hub/storage/scan_folders.py +++ b/studio/backend/hub/storage/scan_folders.py @@ -16,37 +16,14 @@ from datetime import datetime, timezone from storage.studio_db import get_connection from hub.utils.paths import normalize_path +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) _schema_lock = threading.Lock() _schema_ready = False -_SENSITIVE_PATH_COMPONENTS = { - ".aws", - ".azure", - ".config", - ".docker", - ".gcloud", - ".gnupg", - ".huggingface", - ".kaggle", - ".kube", - ".modelscope", - ".ngc", - ".local", - ".mozilla", - ".pki", - ".thunderbird", - ".ssh", - ".1password", - ".bitwarden", - ".password-store", - "1password", - "bitwarden", - "keychains", - "keyrings", - "mozilla", - "thunderbird", -} def _denied_path_prefixes() -> list[str]: @@ -76,8 +53,7 @@ def _denied_path_prefixes() -> list[str]: def _contains_sensitive_path_component(path: str) -> bool: - parts = os.path.normpath(path).split(os.sep) - return any(part.lower() in _SENSITIVE_PATH_COMPONENTS for part in parts) + return _shared_contains_sensitive_path_component(path) def contains_sensitive_path_component(path: str) -> bool: @@ -142,6 +118,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index e22aaba282..f05d8359ec 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -168,6 +168,16 @@ def test_resolve_browse_target_rejects_sensitive_dir(tmp_path): assert exc_info.value.status_code == 403 +def test_resolve_browse_target_rejects_sensitive_root(tmp_path): + ssh = tmp_path / "home" / ".ssh" + ssh.mkdir(parents = True) + + with pytest.raises(HTTPException) as exc_info: + folder_browser._resolve_browse_target(str(ssh), [ssh]) + + assert exc_info.value.status_code == 403 + + def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): home = tmp_path / "home" (home / ".ssh").mkdir(parents = True) @@ -181,6 +191,24 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path): assert ".ssh" not in names +def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + monkeypatch.setattr(folder_browser.Path, "home", lambda: home) + monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root]) + monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf") + monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: []) + monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: []) + + allowlist = folder_browser._build_browse_allowlist() + + assert media_root.resolve() in allowlist + assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve() + + def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path): # The endpoint creates the cache dir on demand so the desktop "Open folder" # action works even before the first download. diff --git a/studio/backend/main.py b/studio/backend/main.py index 5402e5eb7b..8762c43195 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -12,6 +12,8 @@ from pathlib import Path as _Path import asyncio from dataclasses import asdict +from typing import Any, Optional + # Suppress C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" @@ -36,6 +38,10 @@ if sys.platform == "win32": pass del _win_stream +_SYSTEM_GPU_CACHE_TTL_SECONDS = 10.0 +_system_gpu_cache_lock = threading.Lock() +_system_gpu_cache: Optional[tuple[float, dict[str, Any]]] = None + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. @@ -226,7 +232,6 @@ import shutil import warnings from contextlib import asynccontextmanager from importlib.metadata import PackageNotFoundError, version as package_version -from typing import Optional from urllib.parse import urlparse @@ -1078,8 +1083,57 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c return {"status": "shutting_down"} +def _get_cached_system_gpu_info(logger) -> dict[str, Any]: + """Return merged GPU visibility/utilization with bounded live-probe churn.""" + import time + from utils.hardware import get_backend_visible_gpu_info, get_visible_gpu_utilization + + global _system_gpu_cache + now = time.monotonic() + with _system_gpu_cache_lock: + if _system_gpu_cache is not None: + cached_at, cached_gpu_info = _system_gpu_cache + if now - cached_at < _SYSTEM_GPU_CACHE_TTL_SECONDS: + return cached_gpu_info + + try: + visibility_info = get_backend_visible_gpu_info() or {"available": False, "devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU visibility info: {e}") + visibility_info = {"available": False, "devices": []} + + try: + utilization_info = get_visible_gpu_utilization() or {"devices": []} + except Exception as e: + logger.debug(f"Failed to get GPU utilization info: {e}") + utilization_info = {"devices": []} + + util_devices = {d.get("index"): d for d in utilization_info.get("devices", [])} + enriched_devices = [] + + for dev in visibility_info.get("devices", []): + idx = dev.get("index") + util = util_devices.get(idx, {}) + + total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0 + used_vram = util.get("vram_used_gb") or 0 + + enriched_dev = dict(dev) + enriched_dev["vram_used_gb"] = used_vram + enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0 + enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct") + enriched_devices.append(enriched_dev) + + gpu_info = { + "available": visibility_info.get("available", False), + "devices": enriched_devices, + } + _system_gpu_cache = (time.monotonic(), gpu_info) + return gpu_info + + @app.get("/api/system") -async def get_system_info(current_subject: str = Depends(get_current_subject)): +def get_system_info(current_subject: str = Depends(get_current_subject)): """Get system information. Auth-gated: the response (platform, Python/GPU, memory, ML packages) can @@ -1088,31 +1142,84 @@ async def get_system_info(current_subject: str = Depends(get_current_subject)): """ import platform import psutil - from utils.hardware import get_device + import os + import time + import logging + from utils.hardware import get_device, export_capability from utils.hardware.hardware import _backend_label - visibility_info = get_backend_visible_gpu_info() - gpu_info = { - "available": visibility_info["available"], - "devices": visibility_info["devices"], - } + logger = logging.getLogger(__name__) + + gpu_info = _get_cached_system_gpu_info(logger) - # CPU & Memory memory = psutil.virtual_memory() + try: + cpu_freq = psutil.cpu_freq() + except Exception as e: + logger.debug(f"Failed to get CPU frequency: {e}") + cpu_freq = None + + try: + disk = psutil.disk_usage(os.path.abspath(os.sep)) + except Exception as e: + logger.debug(f"Failed to get disk usage: {e}") + disk = None + + try: + current_process = psutil.Process(os.getpid()) + process_used_mb = round(current_process.memory_info().rss / 1024**2) + except Exception as e: + logger.debug(f"Failed to get current process memory: {e}") + process_used_mb = 0 + + try: + boot_time = psutil.boot_time() + except Exception as e: + logger.debug(f"Failed to get boot time: {e}") + boot_time = None + + # Read versions from metadata so a 3s poll never imports heavy ML libs (or 500s on their import errors). + from importlib.metadata import PackageNotFoundError, version as pkg_version + + ml_packages = {} + for pkg in ("torch", "transformers"): + try: + ml_packages[pkg] = pkg_version(pkg) + except PackageNotFoundError: + pass + except Exception as e: + logger.debug(f"Failed to read {pkg} version: {e}") + return { "platform": platform.platform(), "python_version": platform.python_version(), - # _backend_label so /api/system reports "rocm" (not "cuda") on AMD, - # matching /api/hardware and /api/gpu-visibility. "device_backend": _backend_label(get_device()), - "cpu_count": psutil.cpu_count(), + "cpu_count": psutil.cpu_count(logical = True), + "uptime_seconds": max(0, round(time.time() - boot_time)) if boot_time else None, + "cpu": { + "logical_count": psutil.cpu_count(logical = True), + "physical_count": psutil.cpu_count(logical = False), + "usage_percent": psutil.cpu_percent(interval = None), + "frequency_mhz": round(cpu_freq.current, 2) + if cpu_freq and cpu_freq.current is not None + else None, + }, "memory": { - "total_gb": round(memory.total / 1e9, 2), - "available_gb": round(memory.available / 1e9, 2), + "total_gb": round(memory.total / 1024**3, 2), + "available_gb": round(memory.available / 1024**3, 2), "percent_used": memory.percent, + "process_used_mb": process_used_mb, + }, + "disk": { + "total_gb": round(disk.total / 1e9, 2) if disk else 0, + "free_gb": round(disk.free / 1e9, 2) if disk else 0, + "percent_used": disk.percent if disk else 0, }, "gpu": gpu_info, + "ml_packages": ml_packages, + # Export capability + torch-aware reason. See /api/system/hardware. + **export_capability(), } @@ -1135,11 +1242,13 @@ def get_hardware_info( method auto-selection. Sync def (not async): hardware/detail probes can shell out, and FastAPI runs sync endpoints in a threadpool. """ - from utils.hardware import get_gpu_summary, get_package_versions + from utils.hardware import get_gpu_summary, get_package_versions, export_capability body = { "gpu": get_gpu_summary(), "versions": get_package_versions(), + # Export capability + torch-aware reason; the Export UI grays out with the message. + **export_capability(), } if include_details: from utils.llama_cpp_update import get_installed_llama_version diff --git a/studio/backend/models/export.py b/studio/backend/models/export.py index 7e05373f11..9dc4d9451a 100644 --- a/studio/backend/models/export.py +++ b/studio/backend/models/export.py @@ -6,7 +6,7 @@ from pathlib import Path, PureWindowsPath from pydantic import BaseModel, Field, field_validator -from typing import List, Optional, Literal, Dict, Any +from typing import List, Optional, Literal, Dict, Any, Union def _validate_save_directory(value: str) -> str: @@ -168,6 +168,15 @@ class ExportMergedModelRequest(ExportCommonOptions): description = "Export precision / format for the merged model. The compressed-tensors " "options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).", ) + compressed_method: Optional[str] = Field( + None, + description = "Optional quantized-export alias. Either a compressed-tensors scheme " + "(e.g. 'fp8', 'fp8_static', 'w8a8', 'w4a16', 'mxfp4', 'mxfp8', 'nvfp4' - NVIDIA only) " + "from unsloth.save COMPRESSED_EXPORT_SCHEMES, or a portable torchao alias " + "('torchao_fp8', 'torchao_int8') from TORCHAO_EXPORT_SCHEMES that needs no NVIDIA GPU. " + "When set, it overrides format_type. Lets the export UI expose the full set of formats " + "beyond the quick buttons.", + ) class ExportBaseModelRequest(ExportCommonOptions): @@ -189,9 +198,10 @@ class ExportGGUFRequest(BaseModel): def _check_save_directory(cls, v): return _validate_save_directory(v) - quantization_method: str = Field( + quantization_method: Union[str, List[str]] = Field( "Q4_K_M", - description = 'GGUF quantization method (e.g. "Q4_K_M")', + description = 'GGUF quantization method(s). A single method (e.g. "Q4_K_M") or a list ' + '(e.g. ["Q4_K_M", "Q8_0"]) to produce multiple GGUFs from one model load.', ) push_to_hub: bool = Field( False, @@ -219,4 +229,13 @@ class ExportGGUFRequest(BaseModel): class ExportLoRAAdapterRequest(ExportCommonOptions): """Request for exporting only the LoRA adapter (not merged).""" - # Uses fields from ExportCommonOptions only + gguf: bool = Field( + False, + description = "If True, also convert the adapter to a GGUF LoRA file " + "(llama.cpp convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.", + ) + gguf_outtype: Literal["q8_0", "f16", "bf16", "f32"] = Field( + "q8_0", + description = "GGUF LoRA output float type (only used when gguf=True). " + "Q8_0 falls back to F16 per tensor for dims not divisible by the block size (32).", + ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 58701b5828..52cb632c9b 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -780,6 +780,16 @@ class ChatCompletionRequest(BaseModel): True, description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.", ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = ( + "[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the " + "model emitted a tool signal that healing could not repair, retry ONCE with " + "a short nudge appended (the retry shares the full prompt prefix, so the " + "server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips " + "the process default." + ), + ) context_overflow: Optional[Literal["error", "truncate_middle"]] = Field( None, description = ( @@ -1612,6 +1622,14 @@ class AnthropicMessagesRequest(BaseModel): False, description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.", ) + auto_heal_tool_calls: Optional[bool] = Field( + True, + description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).", + ) + nudge_tool_calls: Optional[bool] = Field( + None, + description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).", + ) model_config = {"extra": "allow"} @model_validator(mode = "before") diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index cf2cb2fa70..a7fd7cbec7 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -46,6 +46,23 @@ router = APIRouter() logger = get_logger(__name__) +def _ensure_export_supported() -> None: + """Reject a mutating export request up front (HTTP 400) when the host can't export. + + Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints + (scan/status/logs) are intentionally NOT gated so the Export page can still render the reason. + """ + from utils.hardware import export_capability + + cap = export_capability() + if not cap.get("export_supported", True): + raise HTTPException( + status_code = 400, + detail = cap.get("export_unsupported_message") + or "Export is not supported on this platform.", + ) + + @router.post("/load-checkpoint", response_model = ExportOperationResponse) async def load_checkpoint( request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject) @@ -58,6 +75,7 @@ async def load_checkpoint( a clear error instead of tearing down the user's other running workloads. """ try: + _ensure_export_supported() backend = get_export_backend() # Run in a worker thread (spawns and waits on a subprocess, can take # minutes) so the event loop stays free to serve the live log SSE stream. @@ -266,6 +284,7 @@ async def export_merged_model( Wraps ExportBackend.export_merged_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_merged_model, @@ -275,6 +294,7 @@ async def export_merged_model( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + compressed_method = request.compressed_method, ) if not success: @@ -304,6 +324,7 @@ async def export_base_model( Wraps ExportBackend.export_base_model. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_base_model, @@ -342,6 +363,7 @@ async def export_gguf( Wraps ExportBackend.export_gguf. """ try: + _ensure_export_supported() backend = get_export_backend() # A custom path wins; otherwise the imatrix toggle requests the upstream auto-download. imatrix_file = request.imatrix_path or (True if request.imatrix else None) @@ -382,6 +404,7 @@ async def export_lora_adapter( Wraps ExportBackend.export_lora_adapter. """ try: + _ensure_export_supported() backend = get_export_backend() success, message, output_path = await asyncio.to_thread( backend.export_lora_adapter, @@ -390,6 +413,8 @@ async def export_lora_adapter( repo_id = request.repo_id, hf_token = request.hf_token, private = request.private, + gguf = request.gguf, + gguf_outtype = request.gguf_outtype, ) if not success: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5f2169f86a..1cc1c83ad4 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -165,6 +165,26 @@ def _friendly_error(exc: Exception) -> str: return "An internal error occurred" +def _friendly_upstream_error(text: str) -> str: + """Rewrite a raw llama-server error body into an actionable message where we can. + + The main case is a tool-calling grammar that llama-server can't compile ("failed to + parse grammar" / "failed to initialize samplers"). This surfaces to coding agents as + a hard 400 on every tool-bearing turn. It is a llama-server limitation with some + model/quant + tool-schema combinations, and recent llama.cpp builds handle the common + coding-agent tools, so point the user at updating Studio rather than the raw body. + """ + lowered = text.lower() + if "failed to parse grammar" in lowered or "failed to initialize samplers" in lowered: + return ( + "The model couldn't compile a tool-calling grammar for this request. This is a " + "llama-server limitation with some model/quant and tool-schema combinations. " + "Update Studio (it installs the latest llama.cpp, which handles the common " + "coding-agent tools) or try a different GGUF model." + ) + return f"llama-server error: {text}" + + def _clamp_finish_reason(value) -> str: """Coerce an upstream finish_reason into OpenAI's known chat values. @@ -278,8 +298,8 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 with code="context_length_exceeded" so these paths deliver the same signal as - the non-passthrough path; any other upstream error keeps llama-server's - message verbatim.""" + the non-passthrough path; a tool-grammar compile failure gets the same actionable + guidance as the Anthropic passthrough; any other upstream error stays verbatim.""" if _classify_llama_generation_error(Exception(text)): return HTTPException( status_code = 400, @@ -292,7 +312,7 @@ def _openai_passthrough_error(status_code, text) -> "HTTPException": ) return HTTPException( status_code = status_code, - detail = f"llama-server error: {text[:500]}", + detail = _friendly_upstream_error(text[:500]), ) @@ -764,6 +784,7 @@ def _set_stream_response_read_timeout( _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 +_OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 class _CompatSameTaskTimeout: @@ -1126,6 +1147,16 @@ from core.inference.key_exchange import decrypt_api_key from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client +from core.inference.passthrough_healing import ( + StreamToolCallHealer, + heal_gate, + heal_openai_message, + heal_openai_message_events, + nudge_enabled, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override @@ -8078,6 +8109,13 @@ def _build_chat_request( if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: chat_kwargs["enable_thinking"] = bool(_tpl_kw["enable_thinking"]) explicit_enable_thinking = True + # auto_heal_tool_calls / nudge_tool_calls are not typed on + # ResponsesRequest; lift them from the extra-body so passthrough + # healing (and the opt-in nudge) honor them on both paths. + if isinstance(_extra.get("auto_heal_tool_calls"), bool): + chat_kwargs["auto_heal_tool_calls"] = _extra["auto_heal_tool_calls"] + if isinstance(_extra.get("nudge_tool_calls"), bool): + chat_kwargs["nudge_tool_calls"] = _extra["nudge_tool_calls"] if isinstance(payload.reasoning, dict): effort = payload.reasoning.get("effort") @@ -8312,16 +8350,112 @@ async def _responses_stream( parse_think_markers = _responses_should_parse_think_markers(chat_req, llama_backend) ) reasoning_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} - message_state: dict[str, Any] = {"output_index": None, "item_id": None, "opened": False} + message_state: dict[str, Any] = { + "output_index": None, + "item_id": None, + "opened": False, + "text": "", + } + # Message items already closed mid-stream (a healed tool call splits + # the assistant text into separate message items, as native Responses + # streams do). Kept for the final response.completed snapshot. + closed_message_states: list[dict] = [] # Per-tool-call state keyed by Chat Completions `tool_calls[].index`, # stable across chunks for the same call. Values: # {output_index, item_id, call_id, name, arguments, opened} tool_call_state: dict[int, dict] = {} next_output_index = 0 + # Text-form tool calls promoted back to structured calls (declared + # client tools only); dormant once grammar-mode structured deltas appear. + _allowed_tools = heal_gate( + getattr(chat_req, "auto_heal_tool_calls", None), + body.get("tools"), + body.get("tool_choice"), + ) + healer = StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + healed_tc_index = 0 + + def _healed_tc(call: dict): + # Chat-delta shape for a healed call. Indexes live in a disjoint + # range so a healed call can never merge into a structured call's + # state slot; parallel_tool_calls=false caps healed calls too (the + # upstream cap ran before injection). + nonlocal healed_tc_index + if payload.parallel_tool_calls is False and healed_tc_index >= 1: + return None + tc = { + "index": 1_000_000 + healed_tc_index, + "id": call["id"], + "type": "function", + "function": call["function"], + } + healed_tc_index += 1 + return tc def _sse(event_name: str, payload: dict) -> str: return f"event: {event_name}\ndata: {json.dumps(payload)}\n\n" + def _tool_call_delta_events(tc: dict) -> list: + # One Chat Completions tool_calls delta -> Responses SSE events, + # allocating/merging per-call state (shared by the structured loop + # and the healer's promoted calls). + events = [] + idx = tc.get("index", 0) + st = tool_call_state.get(idx) + fn = tc.get("function") or {} + if st is None: + # First chunk for this tool call -- allocate an + # output_index and emit output_item.added. + st = { + "output_index": _claim_output_index(), + "item_id": f"fc_{uuid.uuid4().hex[:12]}", + "call_id": tc.get("id") or "", + "name": fn.get("name") or "", + "arguments": "", + "opened": False, + } + tool_call_state[idx] = st + else: + # Later chunks sometimes carry id/name only once; merge + # when present. + if tc.get("id") and not st["call_id"]: + st["call_id"] = tc["id"] + if fn.get("name") and not st["name"]: + st["name"] = fn["name"] + + if not st["opened"] and st["call_id"] and st["name"]: + item_added = { + "type": "response.output_item.added", + "output_index": st["output_index"], + "item": { + "type": "function_call", + "id": st["item_id"], + "status": "in_progress", + "call_id": st["call_id"], + "name": st["name"], + "arguments": "", + }, + } + events.append(_sse("response.output_item.added", item_added)) + st["opened"] = True + + arg_delta = fn.get("arguments") or "" + if arg_delta and st["opened"]: + st["arguments"] += arg_delta + args_delta_event = { + "type": "response.function_call_arguments.delta", + "item_id": st["item_id"], + "output_index": st["output_index"], + "delta": arg_delta, + } + events.append(_sse("response.function_call_arguments.delta", args_delta_event)) + elif arg_delta: + # Buffer args until we can open the item (some models + # send id/name in the same chunk as the first arg delta; + # if not, stash). + st["arguments"] += arg_delta + return events + def _claim_output_index() -> int: nonlocal next_output_index output_index = next_output_index @@ -8406,6 +8540,98 @@ async def _responses_stream( ), ] + def _close_message_item() -> list[str]: + """Close the open message item so later text opens a fresh one. + + Emits the same done-event triplet the end-of-stream close loop + would, records the item for the final snapshot, and resets the + state in place. No-op when no message item is open. + """ + if not message_state["opened"]: + return [] + text = message_state["text"] + events = [ + _sse( + "response.output_text.done", + { + "type": "response.output_text.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "text": text, + }, + ), + _sse( + "response.content_part.done", + { + "type": "response.content_part.done", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "part": {"type": "output_text", "text": text, "annotations": []}, + }, + ), + _sse( + "response.output_item.done", + { + "type": "response.output_item.done", + "output_index": message_state["output_index"], + "item": { + "type": "message", + "id": message_state["item_id"], + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + }, + }, + ), + ] + closed_message_states.append(dict(message_state)) + message_state.update( + {"output_index": None, "item_id": None, "opened": False, "text": ""} + ) + return events + + def _healed_event_sse(events) -> list[str]: + """Serialize healer events preserving their order. + + Text around a healed call must keep its position relative to the + function_call item (output indexes are claimed in emission order), + so never split an event list into all-text-then-all-calls. A healed + call also CLOSES any open message item, so trailing text opens a + fresh message with a later output index, exactly like a native + Responses stream that interleaves messages and calls. + """ + nonlocal full_text + out: list[str] = [] + for kind, value in events: + if kind == "text": + if not value: + continue + out.extend(_ensure_message_open()) + full_text += value + message_state["text"] += value + api_monitor.append_reply(monitor_id, value) + out.append( + _sse( + "response.output_text.delta", + { + "type": "response.output_text.delta", + "item_id": message_state["item_id"], + "output_index": message_state["output_index"], + "content_index": 0, + "delta": value, + }, + ) + ) + else: + tc = _healed_tc(value) + if tc is None: + continue + out.extend(_close_message_item()) + out.extend(_tool_call_delta_events(tc)) + return out + def _snapshot_output() -> list[dict]: """Snapshot of all completed output items for response.completed.""" indexed_items: list[tuple[int, dict]] = [] @@ -8422,19 +8648,23 @@ async def _responses_stream( }, ) ) - if message_state["opened"]: + # Closed copies keep opened=True (snapshotted before reset); the + # live state contributes only when a message is currently open. + for msg_st in [*closed_message_states, message_state]: + if not msg_st["opened"]: + continue indexed_items.append( ( - message_state["output_index"], + msg_st["output_index"], { "type": "message", - "id": message_state["item_id"], + "id": msg_st["item_id"], "status": "completed", "role": "assistant", "content": [ { "type": "output_text", - "text": full_text, + "text": msg_st["text"], "annotations": [], } ], @@ -8562,7 +8792,7 @@ async def _responses_stream( "output": [], "error": { "code": resp.status_code, - "message": f"llama-server error: {err_text[:500]}", + "message": _friendly_upstream_error(err_text[:500]), }, }, }, @@ -8618,10 +8848,30 @@ async def _responses_stream( "delta": reasoning_delta, }, ) + # Heal text-form tool calls in the visible stream (never in + # reasoning text): promoted calls join the structured tc loop + # below through the same state machinery, and healer events are + # emitted IN ORDER so text after a healed call never jumps ahead + # of the function_call item. Once a structured delta arrives, + # grammar mode worked and the healer goes dormant. + if healer is not None and not healer.dormant: + healed_events = [] + if delta.get("tool_calls"): + # Held text preceded the structured call; the call's own + # deltas follow in the structured loop below. + healed_events = healer.structured_tool_call_seen() + if visible_delta: + healed_events.append(("text", visible_delta)) + elif visible_delta: + healed_events = healer.feed(visible_delta) + visible_delta = "" + for event in _healed_event_sse(healed_events): + yield event if visible_delta: for event in _ensure_message_open(): yield event full_text += visible_delta + message_state["text"] += visible_delta api_monitor.append_reply(monitor_id, visible_delta) yield _sse( "response.output_text.delta", @@ -8635,60 +8885,19 @@ async def _responses_stream( ) for tc in delta.get("tool_calls") or []: - idx = tc.get("index", 0) - st = tool_call_state.get(idx) - fn = tc.get("function") or {} - if st is None: - # First chunk for this tool call -- allocate an - # output_index and emit output_item.added. - st = { - "output_index": _claim_output_index(), - "item_id": f"fc_{uuid.uuid4().hex[:12]}", - "call_id": tc.get("id") or "", - "name": fn.get("name") or "", - "arguments": "", - "opened": False, - } - tool_call_state[idx] = st - else: - # Later chunks sometimes carry id/name only once; merge - # when present. - if tc.get("id") and not st["call_id"]: - st["call_id"] = tc["id"] - if fn.get("name") and not st["name"]: - st["name"] = fn["name"] - - if not st["opened"] and st["call_id"] and st["name"]: - item_added = { - "type": "response.output_item.added", - "output_index": st["output_index"], - "item": { - "type": "function_call", - "id": st["item_id"], - "status": "in_progress", - "call_id": st["call_id"], - "name": st["name"], - "arguments": "", - }, - } - yield _sse("response.output_item.added", item_added) - st["opened"] = True - - arg_delta = fn.get("arguments") or "" - if arg_delta and st["opened"]: - st["arguments"] += arg_delta - args_delta_event = { - "type": "response.function_call_arguments.delta", - "item_id": st["item_id"], - "output_index": st["output_index"], - "delta": arg_delta, - } - yield _sse("response.function_call_arguments.delta", args_delta_event) - elif arg_delta: - # Buffer args until we can open the item (some models - # send id/name in the same chunk as the first arg delta; - # if not, stash). - st["arguments"] += arg_delta + if ( + payload.parallel_tool_calls is False + and healed_tc_index >= 1 + and tc.get("index", 0) not in tool_call_state + ): + # A healed call already consumed the single allowed slot; + # _drop_parallel_tool_call_deltas only sees native indexes, + # so a native index-0 call would still open a second + # function_call item. Skip it (and its later argument + # deltas, which never allocate a state either). + continue + for event in _tool_call_delta_events(tc): + yield event _apply_usage(chunk_data.get("usage")) except asyncio.CancelledError: @@ -8744,10 +8953,19 @@ async def _responses_stream( "delta": final_reasoning, }, ) + # Last-chance heal of any held residue (e.g. a tool block the model + # never closed) before the trailing visible text is flushed; events + # keep healer order so trailing text stays behind a healed call. + if healer is not None: + events = (healer.feed(final_visible) if final_visible else []) + healer.finalize() + final_visible = "" + for event in _healed_event_sse(events): + yield event if final_visible: for event in _ensure_message_open(): yield event full_text += final_visible + message_state["text"] += final_visible api_monitor.append_reply(monitor_id, final_visible) yield _sse( "response.output_text.delta", @@ -8806,6 +9024,10 @@ async def _responses_stream( continue if kind == "message": + # Per-item text: message items closed mid-stream (healed-call + # rotation) already emitted their done events, so this state + # carries only its own text, not the whole stream's. + _msg_text = st["text"] yield _sse( "response.output_text.done", { @@ -8813,7 +9035,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "text": full_text, + "text": _msg_text, }, ) yield _sse( @@ -8823,7 +9045,7 @@ async def _responses_stream( "item_id": st["item_id"], "output_index": st["output_index"], "content_index": 0, - "part": {"type": "output_text", "text": full_text, "annotations": []}, + "part": {"type": "output_text", "text": _msg_text, "annotations": []}, }, ) yield _sse( @@ -8837,7 +9059,7 @@ async def _responses_stream( "status": "completed", "role": "assistant", "content": [ - {"type": "output_text", "text": full_text, "annotations": []} + {"type": "output_text", "text": _msg_text, "annotations": []} ], }, }, @@ -9443,6 +9665,7 @@ async def anthropic_messages( session_id = payload.session_id, cancel_id = payload.cancel_id, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, ) ) return await _monitored_anthropic( @@ -9462,6 +9685,8 @@ async def anthropic_messages( presence_penalty = presence_penalty, tool_choice = openai_tool_choice, disable_parallel_tool_use = _disable_parallel, + auto_heal_tool_calls = payload.auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, ) ) @@ -10032,6 +10257,7 @@ async def _anthropic_passthrough_stream( session_id = None, cancel_id = None, disable_parallel_tool_use = False, + auto_heal_tool_calls = None, ): """Streaming client-side pass-through: forward tools to llama-server and translate its stream to Anthropic SSE without executing anything.""" @@ -10068,6 +10294,16 @@ async def _anthropic_passthrough_stream( async def _stream(): emitter = AnthropicPassthroughEmitter() + # Promote text-form tool calls (declared client tools only) into + # tool_use blocks; verbatim behavior when healing is off or no tools. + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + if _allowed_tools: + emitter.enable_healing( + _allowed_tools, + openai_tools, + disable_parallel_tool_use = disable_parallel_tool_use, + ) for line in emitter.start(message_id, model_name, input_tokens = input_tokens): yield line @@ -10129,7 +10365,7 @@ async def _anthropic_passthrough_stream( yield build_anthropic_sse_event( "error", anthropic_error_body( - f"llama-server error: {_err_text}", + _friendly_upstream_error(_err_text), status = resp.status_code, ), ) @@ -10204,6 +10440,8 @@ async def _anthropic_passthrough_non_streaming( presence_penalty = None, tool_choice = "auto", disable_parallel_tool_use = False, + auto_heal_tool_calls = None, + nudge_tool_calls = None, ): """Non-streaming client-side pass-through.""" target_url = f"{llama_backend.base_url}/v1/chat/completions" @@ -10232,38 +10470,102 @@ async def _anthropic_passthrough_non_streaming( if resp.status_code != 200: raise HTTPException( status_code = resp.status_code, - detail = f"llama-server error: {resp.text[:500]}", + detail = _friendly_upstream_error(resp.text[:500]), ) data = resp.json() + # tool_choice arrives here already converted to the OpenAI shape. + _allowed_tools = heal_gate(auto_heal_tool_calls, openai_tools, tool_choice) + + # Opt-in single-retry nudge (mirrors the OpenAI passthrough): the model + # tried to call a tool but nothing usable came out; re-ask once with the + # prompt prefix intact so llama-server's KV cache is reused. + if ( + _allowed_tools + and nudge_enabled(nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, openai_tools) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await nonstreaming_client().post( + target_url, + json = retry_body, + timeout = _llama_non_streaming_generation_timeout(), + ) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, openai_tools): + data = retry_data + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + choice = (data.get("choices") or [{}])[0] message = choice.get("message") or {} finish_reason = choice.get("finish_reason") - content_blocks = [] - text = message.get("content") or "" - if text: - text = _TOOL_XML_RE.sub("", text).strip() - if text: - content_blocks.append(AnthropicResponseTextBlock(text = text)) + healing_active = bool(_allowed_tools) + healed_events = ( + heal_openai_message_events(message, _allowed_tools, openai_tools) + if healing_active + else None + ) - tool_calls = message.get("tool_calls") or [] - # disable_parallel_tool_use: keep only the first tool_use block. - if disable_parallel_tool_use and len(tool_calls) > 1: - tool_calls = tool_calls[:1] - for tc in tool_calls: - fn = tc.get("function") or {} - try: - args = json.loads(fn.get("arguments", "{}")) - except json.JSONDecodeError: - args = {} - content_blocks.append( - AnthropicResponseToolUseBlock( - id = anthropic_tool_use_id(tc.get("id")), - name = fn.get("name", ""), - input = args, + content_blocks = [] + tool_calls = [] + if healed_events: + emitted_tool_uses = 0 + for kind, value in healed_events: + if kind == "text": + text = str(value).strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + continue + if disable_parallel_tool_use and emitted_tool_uses >= 1: + continue + fn = value.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + tool_calls.append(value) + emitted_tool_uses += 1 + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(value.get("id")), + name = fn.get("name", ""), + input = args, + ) + ) + else: + text = message.get("content") or "" + if text: + # Keep unpromoted bytes when healing is active; legacy stripping is + # only for opted-out or no-client-tool requests. + if not healing_active: + text = _TOOL_XML_RE.sub("", text) + text = text.strip() + if text: + content_blocks.append(AnthropicResponseTextBlock(text = text)) + + tool_calls = message.get("tool_calls") or [] + if disable_parallel_tool_use and len(tool_calls) > 1: + tool_calls = tool_calls[:1] + for tc in tool_calls: + fn = tc.get("function") or {} + try: + args = json.loads(fn.get("arguments", "{}")) + except json.JSONDecodeError: + args = {} + content_blocks.append( + AnthropicResponseToolUseBlock( + id = anthropic_tool_use_id(tc.get("id")), + name = fn.get("name", ""), + input = args, + ) ) - ) stop_reason = openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = bool(tool_calls)) @@ -10611,45 +10913,84 @@ async def _openai_passthrough_stream( body = _build_openai_passthrough_body( payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + # Text-form tool calls from small models get promoted to structured calls on + # the way back (declared client tools only); requests without tools or with + # auto_heal_tool_calls=false keep the verbatim relay. tool_choice constrains + # the allowlist ("none" disables, a forced function narrows to it). + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + client = None + resp = None + send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None + + async def _aclose_send_task(task: Optional[asyncio.Task[Optional[httpx.Response]]]) -> None: + if task is None: + return + if not task.done(): + task.cancel() + try: + task_resp = await task + if task_resp is not None: + try: + await task_resp.aclose() + except Exception: + pass + except (asyncio.CancelledError, Exception): + pass # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: - # Dispatch BEFORE returning StreamingResponse so transport errors and - # non-200 upstream statuses surface as real HTTP errors -- OpenAI SDKs - # rely on status codes to raise APIError/BadRequestError. + # Keep the pre-header window short so accepted SSE clients receive + # immediate headers in the common timeout-reduced stall. client = httpx.AsyncClient( timeout = _llama_streaming_generation_timeout(), limits = httpx.Limits(max_keepalive_connections = 0), trust_env = False, ) - resp = None _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) + while True: try: req = client.build_request( "POST", target_url, json = body, headers = {"Connection": "close"} ) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S - resp = await _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel(client, req, cancel_event, request = request) ) + done, _ = await asyncio.wait( + {send_task}, + timeout = _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task not in done: + break + + # Dispatch returned quickly enough to preserve pre-header status. + resp = await send_task + send_task = None except httpx.RequestError as e: # llama-server subprocess crashed / starting / unreachable. logger.error("openai passthrough stream: upstream unreachable: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) raise HTTPException( status_code = 502, detail = _friendly_error(e), ) + if resp is None and send_task is not None and not send_task.done(): + break if resp is None: api_monitor.finish(monitor_id, "cancelled") + await _aclose_send_task(send_task) try: await client.aclose() except Exception: @@ -10694,6 +11035,8 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, err_text[:500]) raise _openai_passthrough_error(upstream_status, err_text) + # Keep tracker cleanup paired if pre-header dispatch is cancelled after we + # have already committed headers. async def _stream(): # Same httpx lifecycle pattern as _anthropic_passthrough_stream: # save resp.aiter_lines() so the finally block can aclose() it on @@ -10701,10 +11044,10 @@ async def _openai_passthrough_stream( lines_iter = None # Watchers unblock aiter_lines() during prefill, before in-loop # cancel/disconnect checks can run. - cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) - disconnect_watcher = asyncio.create_task( - _await_disconnect_then_close(request, resp, cancel_event) - ) + cancel_watcher = None + disconnect_watcher = None + + nonlocal resp, send_task, first_token_deadline, _truncate_budget monitor_done = False saw_finish_reason = False saw_done = False @@ -10713,9 +11056,14 @@ async def _openai_passthrough_stream( last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) + healer = ( + StreamToolCallHealer(_allowed_tools, body.get("tools")) if _allowed_tools else None + ) + healed_call_index = 0 def _synthetic_finish_line() -> str: - finish_reason = "tool_calls" if saw_tool_call_delta else "stop" + healed = healer is not None and healer.healed + finish_reason = "tool_calls" if (saw_tool_call_delta or healed) else "stop" chunk = ChatCompletionChunk( id = last_chunk_id, created = last_chunk_created, @@ -10729,7 +11077,182 @@ async def _openai_passthrough_stream( ) return f"data: {chunk.model_dump_json(exclude_none = True)}" + def _healer_sse_lines(events) -> list: + # Serialize healer events as chunks matching the upstream stream's + # id/model/created so clients see one coherent completion. + nonlocal healed_call_index + lines = [] + for kind, value in events: + if kind == "text": + if not value: + continue + delta = {"content": value} + else: + # parallel_tool_calls=false caps healed calls too (the SSE + # line cap only sees structured upstream deltas). + if payload.parallel_tool_calls is False and healed_call_index >= 1: + continue + delta = { + "tool_calls": [ + { + "index": healed_call_index, + "id": value["id"], + "type": "function", + "function": value["function"], + } + ] + } + healed_call_index += 1 + chunk = { + "id": last_chunk_id, + "object": "chat.completion.chunk", + "created": last_chunk_created, + "model": last_chunk_model, + "choices": [{"index": 0, "delta": delta, "finish_reason": None}], + } + lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) + return lines + + def _heal_transform(chunk_data: dict, raw_line: str) -> list: + """SSE lines to emit in place of one upstream line (healing on).""" + choices = chunk_data.get("choices") + if not (isinstance(choices, list) and choices and isinstance(choices[0], dict)): + return [raw_line] + choice = choices[0] + delta = choice.get("delta") + delta = delta if isinstance(delta, dict) else {} + if delta.get("tool_calls"): + # Structured call streamed: grammar mode worked. Flush any held + # text (it preceded the call) and relay verbatim from here on. + lines = _healer_sse_lines(healer.structured_tool_call_seen()) + if healed_call_index: + if payload.parallel_tool_calls is False: + # A healed call already consumed the single allowed + # slot; the upstream SSE cap keeps native index 0, so + # drop the native call here or the client gets two. + del delta["tool_calls"] + if delta or choice.get("finish_reason") or chunk_data.get("usage"): + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + # A healed call already went out on index 0..n-1; OpenAI + # clients merge tool-call deltas by index, so shift the + # native calls into the next indexes or they would merge + # into the healed call. + for tc in delta["tool_calls"]: + if isinstance(tc, dict) and isinstance(tc.get("index"), int): + tc["index"] += healed_call_index + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + content = delta.get("content") + finish = choice.get("finish_reason") + if not isinstance(content, str) or not content: + if not finish: + return [raw_line] + # Finish chunk: last-chance heal of the residue, and rewrite a + # "stop" into "tool_calls" when text-form calls were promoted. + lines = _healer_sse_lines(healer.finalize()) + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + return lines + ["data: " + json.dumps(chunk_data, ensure_ascii = False)] + return lines + [raw_line] + events = healer.feed(content) + if finish: + events += healer.finalize() + if not finish and events == [("text", content)]: + # Nothing held or promoted: the healer passed the chunk + # through whole, so keep the verbatim upstream bytes. + return [raw_line] + del delta["content"] + prefix_lines = [] + if delta: + prefix_chunk = {k: v for k, v in chunk_data.items() if k != "usage"} + prefix_choice = dict(choice) + prefix_choice["delta"] = dict(delta) + prefix_choice["finish_reason"] = None + prefix_chunk["choices"] = [prefix_choice] + prefix_lines.append("data: " + json.dumps(prefix_chunk, ensure_ascii = False)) + delta.clear() + lines = prefix_lines + _healer_sse_lines(events) + if delta or finish or chunk_data.get("usage"): + if healer.healed and finish == "stop": + choice["finish_reason"] = "tool_calls" + lines.append("data: " + json.dumps(chunk_data, ensure_ascii = False)) + return lines + try: + while True: + if send_task is not None and not send_task.done(): + try: + resp = await send_task + except httpx.RequestError as e: + logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + return + send_task = None + elif send_task is not None: + try: + resp = send_task.result() + except httpx.RequestError as e: + logger.error("openai passthrough stream: upstream unreachable: %s", e) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield f"data: {json.dumps(_openai_stream_error_chunk(e))}\n\n" + return + send_task = None + + if resp is None: + api_monitor.finish(monitor_id, "cancelled") + return + if resp.status_code == 200: + break + + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + logger.error( + "openai passthrough upstream error: status=%s body=%s", + resp.status_code, + err_text[:500], + ) + upstream_status = resp.status_code + try: + await resp.aclose() + except Exception: + pass + resp = None + if ( + _truncate_budget > 0 + and _classify_llama_generation_error(Exception(err_text)) + and _apply_overflow_truncation(body, err_text) + ): + _truncate_budget -= 1 + req = client.build_request( + "POST", target_url, json = body, headers = {"Connection": "close"} + ) + first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S + send_task = asyncio.create_task( + _send_stream_with_preheader_cancel( + client, req, cancel_event, request = request + ) + ) + continue + + upstream_error = _openai_passthrough_error(upstream_status, err_text) + error_payload = ( + upstream_error.detail + if isinstance(upstream_error.detail, dict) + else openai_error_body( + str(upstream_error.detail), + status = upstream_status, + ) + ) + api_monitor.fail(monitor_id, err_text[:500]) + yield f"data: {json.dumps(error_payload)}\n\n" + return + + cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_close(request, resp, cancel_event) + ) lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( lines_iter, @@ -10745,6 +11268,14 @@ async def _openai_passthrough_stream( data_text = raw_line[6:].strip() if data_text == "[DONE]": saw_done = True + # Upstream ended without a finish chunk: heal the residue + # first so the synthetic finish sees healer.healed. + if healer is not None and not saw_stream_error: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" if ( not saw_finish_reason and not saw_stream_error @@ -10800,13 +11331,18 @@ async def _openai_passthrough_stream( # emit a successful finish_reason after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True - monitor_event = _monitor_openai_sse_line( - monitor_id, - raw_line, - llama_backend.context_length, - ) - if monitor_event == "error": - saw_stream_error = True + # With healing active, a content-bearing line may be replaced by + # held/promoted chunks; otherwise the single upstream line + # relays verbatim (monitored exactly as emitted either way). + if ( + healer is not None + and not healer.dormant + and isinstance(chunk_data, dict) + and not saw_stream_error + ): + out_lines = _heal_transform(chunk_data, raw_line) + else: + out_lines = [raw_line] # If a trailing usage-only chunk (include_usage) arrives before # any finish chunk, emit the synthetic finish first so the order # stays finish -> usage -> [DONE], matching the other streams. @@ -10820,23 +11356,46 @@ async def _openai_passthrough_stream( and not saw_stream_error and not cancel_event.is_set() ): + if healer is not None: + # Residue must precede the finish it may upgrade. + held = _healer_sse_lines(healer.finalize()) + for held_line in held: + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" finish_line = _synthetic_finish_line() _monitor_openai_sse_line( monitor_id, finish_line, llama_backend.context_length ) yield finish_line + "\n\n" saw_finish_reason = True - # Relay verbatim to preserve llama-server's native id, - # finish_reason, delta.tool_calls, and usage chunks. - yield raw_line + "\n\n" - if monitor_event == "done": - monitor_done = True + for out_line in out_lines: + monitor_event = _monitor_openai_sse_line( + monitor_id, + out_line, + llama_backend.context_length, + ) + if monitor_event == "error": + saw_stream_error = True + # Relay to preserve llama-server's native id, + # finish_reason, delta.tool_calls, and usage chunks. + yield out_line + "\n\n" + if monitor_event == "done": + monitor_done = True + if monitor_done: break if not saw_done and not saw_stream_error and not cancel_event.is_set(): # Synthesize a finish chunk only if one was not already # emitted (e.g. before a trailing usage-only chunk), but # always close with [DONE] whenever the upstream omitted it, # so the stream ends on the [DONE] sentinel either way. + if healer is not None: + for held_line in _healer_sse_lines(healer.finalize()): + _monitor_openai_sse_line( + monitor_id, held_line, llama_backend.context_length + ) + yield held_line + "\n\n" if not saw_finish_reason: finish_line = _synthetic_finish_line() _monitor_openai_sse_line( @@ -10880,6 +11439,7 @@ async def _openai_passthrough_stream( err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: + await _aclose_send_task(send_task) await _aclose_stream_resources( watchers = (cancel_watcher, disconnect_watcher), iterator = lines_iter, @@ -10893,6 +11453,7 @@ async def _openai_passthrough_stream( # finally never ran. Release the eagerly-opened upstream resp/client # and the cancel-registry entry here; the watchers and line iterator # are created inside _stream(), so there is nothing else to close. + await _aclose_send_task(send_task) await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) @@ -10907,6 +11468,8 @@ async def _openai_passthrough_stream( unstarted_cleanup = _unstarted_cleanup, ) except BaseException: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) raise @@ -10975,6 +11538,9 @@ async def _openai_passthrough_non_streaming( _guided_fence = bool((payload.model_extra or {}).get("_unsloth_guided_fence")) _do_fence = _guided_fence and _extract_response_format(payload) is not None _cap_parallel = payload.parallel_tool_calls is False + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) try: data = resp.json() @@ -10987,6 +11553,33 @@ async def _openai_passthrough_non_streaming( api_monitor.finish(monitor_id) return Response(content = resp.content, media_type = "application/json") + # Opt-in single-retry nudge: the model clearly tried to call a tool (signal + # present) but nothing parseable/declared came out, so re-ask once with the + # original prompt prefix intact (llama-server reuses the slot's KV cache) + # plus a two-message nudge suffix. The retry replaces the original response + # only when it actually yields a usable call. + if ( + _allowed_tools + and nudge_enabled(payload.nudge_tool_calls) + and nudge_should_retry(data, _allowed_tools, body.get("tools")) + ): + retry_body = { + **body, + "messages": [*body.get("messages", []), *nudge_messages(data, _allowed_tools)], + } + try: + retry_resp = await nonstreaming_client().post( + target_url, + json = retry_body, + timeout = _llama_non_streaming_generation_timeout(), + ) + if retry_resp.status_code == 200: + retry_data = retry_resp.json() + if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): + resp, data = retry_resp, retry_data + except (httpx.RequestError, ValueError) as exc: + logger.warning("tool-call nudge retry failed; keeping original: %s", exc) + changed = False for choice in data.get("choices", []): if not isinstance(choice, dict): @@ -10995,6 +11588,17 @@ async def _openai_passthrough_non_streaming( if not isinstance(msg, dict): continue + # Small models emit tool calls as text instead of structured tool_calls; + # promote them (declared client tools only) so the agent sees a real call. + # Truncation wins over the upgrade (same rule as the streaming and + # Anthropic paths): a call cut off at max_tokens keeps + # finish_reason="length" so the client knows the arguments may be + # incomplete, while the healed call itself stays attached. + if _allowed_tools and heal_openai_message(msg, _allowed_tools, body.get("tools")): + if choice.get("finish_reason") == "stop": + choice["finish_reason"] = "tool_calls" + changed = True + # OpenAI requires content=null on a pure tool-call turn; llama-server # emits content="". if msg.get("tool_calls") and msg.get("content") == "": diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 9ea12644eb..65d64a135d 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -7,6 +7,7 @@ import asyncio import hashlib import json import os +import re import shutil import sys import uuid @@ -62,25 +63,51 @@ def _safe_is_dir(path) -> bool: return False +# Hub repo id shape ("owner/name", no leading separator); anything else is +# treated as a local filesystem path. +_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$") + + def _is_hidden_model(*values: str | None) -> bool: """True if any id/path is the RAG embedding model (EMBEDDING_MODEL or EMBED_GGUF_REPO basename) or the llama.cpp install validation probe (ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF). None are usable chat models; the probe can be cached as a side effect of installing the prebuilt llama-server and otherwise sorts smallest, so it - would be auto-selected.""" + would be auto-selected. A local-path embedder is matched by exact resolved + path only: a generic basename like "model" must not substring-hide + unrelated chat models.""" from core.rag import config as rag_config - needles = ( - rag_config.EMBEDDING_MODEL.split("/")[-1].lower(), - rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(), + needles = [ # The validation probe's repo (matches the cached repo id) and its exact # filename (matches the on-disk path). The filename carries the .gguf so # it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``. "ggml-org/models", "stories260k.gguf", - ) - return any(v and any(n in v.lower() for n in needles) for v in values) + ] + exact_paths: list[str] = [] + for model in ( + rag_config.effective_embedding_model(), + rag_config.effective_gguf_repo(), + ): + if _HF_REPO_ID_RE.match(model): + needles.append(model.split("/")[-1].lower()) + else: + resolved = _safe_resolve(Path(model).expanduser()) + if resolved: + exact_paths.append(resolved.lower()) + for v in values: + if not v: + continue + low = v.lower() + if any(n in low for n in needles): + return True + if exact_paths: + resolved = _safe_resolve(Path(v).expanduser()) + if resolved and resolved.lower() in exact_paths: + return True + return False def _safe_resolve(path: Path) -> Optional[str]: @@ -1183,6 +1210,7 @@ def _build_browse_allowlist() -> list[Path]: legacy_hf_cache_dir, well_known_model_dirs, ) + from utils.paths.external_media import linux_run_media_mount_roots from storage.studio_db import list_scan_folders candidates: list[Path] = [] @@ -1198,6 +1226,8 @@ def _build_browse_allowlist() -> list[Path]: candidates.append(resolved) _add(Path.home()) + for p in linux_run_media_mount_roots(): + _add(p) _add(_resolve_hf_cache_dir()) try: _add(hf_default_cache_dir()) @@ -1317,6 +1347,8 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]: def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path: """Resolve a requested browse path by walking from trusted allowlist roots.""" + from storage.studio_db import contains_sensitive_path_component + requested_path = _normalize_browse_request_path(path) resolved_roots: list[Path] = [] seen_roots: set[str] = set() @@ -1367,8 +1399,18 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa "under your home folder." ), ) + if contains_sensitive_path_component(str(resolved_child)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) current = resolved_child + if contains_sensitive_path_component(str(current)): + raise HTTPException( + status_code = 403, + detail = "Credential or configuration directories are not browseable.", + ) if not current.is_dir(): raise HTTPException( status_code = 400, @@ -1416,7 +1458,8 @@ async def browse_folders( then hidden (if ``show_hidden=true``). """ from utils.paths import hf_default_cache_dir, well_known_model_dirs - from storage.studio_db import list_scan_folders + from utils.paths.external_media import linux_run_media_mount_roots + from storage.studio_db import contains_sensitive_path_component, list_scan_folders # Build once; the sandbox check and suggestion chips share it. allowed_roots = _build_browse_allowlist() @@ -1469,6 +1512,8 @@ async def browse_folders( is_hidden = name.startswith(".") if is_hidden and not show_hidden: continue + if contains_sensitive_path_component(name): + continue entries.append( BrowseEntry( name = name, @@ -1522,6 +1567,8 @@ async def browse_folders( # Home first -- the safe fallback when everything else is cold. _add_sug(Path.home()) + for p in linux_run_media_mount_roots(): + _add_sug(p) # The HF cache root the process is actually using. try: _add_sug(hf_default_cache_dir()) diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 4e35fce3c2..e20fea74a3 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -167,7 +167,7 @@ def create_knowledge_base( conn, name = payload.name.strip(), description = (payload.description or None), - embedding_model = config.EMBEDDING_MODEL, + embedding_model = config.effective_embedding_model(), ) return {"id": kb_id, "name": payload.name.strip()} finally: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 0694ae31e0..862bce8be8 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -4,7 +4,7 @@ from typing import Literal, Optional from urllib.parse import unquote, urlsplit -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, ConfigDict, Field, field_validator from auth.authentication import get_current_subject @@ -47,6 +47,15 @@ from utils.preview_sharing_settings import ( get_preview_sharing_enabled, set_preview_sharing_enabled, ) +from utils.embedding_model_settings import ( + MAX_EMBEDDING_MODEL_LENGTH, + default_embedding_model, + get_rag_embedding_model, + get_stored_embedding_model, + reset_rag_embedding_model, + set_rag_embedding_model, + validate_embedding_model, +) router = APIRouter() @@ -229,6 +238,186 @@ def update_openai_auto_switch_override( return ModelOverridesResponse(overrides = get_model_overrides()) +class EmbeddingModelPayload(BaseModel): + embedding_model: str = Field(..., min_length = 1, max_length = MAX_EMBEDDING_MODEL_LENGTH) + # Token for gated/private repos during verification (not stored). + hf_token: Optional[str] = Field(default = None, max_length = 512) + # Skip HF verification (offline installs, local paths HF can't see). + force: bool = False + + +class EmbeddingModelResponse(BaseModel): + embedding_model: str + default_embedding_model: str + is_custom: bool + + +def _embedding_model_response() -> EmbeddingModelResponse: + return EmbeddingModelResponse( + embedding_model = get_rag_embedding_model(), + default_embedding_model = default_embedding_model(), + is_custom = get_stored_embedding_model() is not None, + ) + + +def _llama_backend_active() -> bool: + """True when this install embeds via the llama-server (GGUF) backend.""" + from core.rag import config as rag_config + from core.rag import embeddings + + try: + raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() + key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + except Exception: # noqa: BLE001 - backend probe must never block saving + return False + return key in embeddings._LLAMA_ALIASES + + +def _resolves_as_local_gguf(model: str) -> bool: + """True when ``model`` is a local .gguf file or a directory holding one, so + a save on the llama-server backend needs no HF verification (the artifact + itself is the proof).""" + from core.rag.embed_llama_server import LlamaServerBackend + try: + return LlamaServerBackend._resolve_local_gguf(model) is not None + except Exception: # noqa: BLE001 - dir without .gguf, filesystem oddity + return False + + +def _local_gguf_backend_error(model: str) -> str | None: + """409 detail when ``model`` is a local dir without a .gguf but this install + embeds via llama-server (macOS/CPU default), which needs one. A + sentence-transformers-only folder would verify fine yet fail at first index. + None when not applicable. ``force`` skips this check like HF verification.""" + from pathlib import Path + + if not Path(model).expanduser().is_dir(): + return None + from core.rag.embed_llama_server import LlamaServerBackend + + if not _llama_backend_active(): + return None + try: + LlamaServerBackend._resolve_local_gguf(model) + return None + except RuntimeError: + return ( + f"{model!r} contains no .gguf file, but this install embeds with the " + "llama-server backend which requires one. Add a GGUF file to the " + "folder or use a Hugging Face repo." + ) + except Exception: # noqa: BLE001 - filesystem oddity: don't block saving + return None + + +def _hf_gguf_backend_error(model: str, hf_token: Optional[str]) -> str | None: + """409 detail when the llama-server backend would find no .gguf for an HF + repo: neither the derived companion repo nor the repo itself has one. Saves + that verify as embedding models would otherwise fail at first index. + None when not applicable; ``force`` skips this like HF verification.""" + from pathlib import Path + + if Path(model).expanduser().exists(): + return None # local paths are handled by the local checks + if not _llama_backend_active(): + return None + from core.rag import config as rag_config + + candidates = [model] if rag_config._names_gguf(model) else [f"{model}-GGUF", model] + try: + from huggingface_hub import list_repo_files + except Exception: # noqa: BLE001 - hub client unavailable: don't block saving + return None + for candidate in candidates: + try: + files = list_repo_files(candidate, token = hf_token) + except Exception: # noqa: BLE001 - missing/gated repo: try next candidate + continue + if any(f.lower().endswith(".gguf") and "mmproj" not in f.lower() for f in files): + return None + checked = " or ".join(repr(c) for c in candidates) + return ( + f"No GGUF weights found in {checked}, but this install embeds with the " + "llama-server backend which requires them. Pick a model with a GGUF " + "companion repo or GGUF files in the repo itself." + ) + + +@router.get("/embedding-model", response_model = EmbeddingModelResponse) +def get_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + return _embedding_model_response() + + +@router.put("/embedding-model", response_model = EmbeddingModelResponse) +def update_embedding_model( + payload: EmbeddingModelPayload, current_subject: str = Depends(get_current_subject) +) -> EmbeddingModelResponse: + """Set the RAG embedding model. Unless ``force`` is set, the repo is verified + to be an embedding model via HF metadata; an unverifiable model (wrong type, + typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + Documents indexed under the previous model must be re-uploaded.""" + from utils.models import is_embedding_model + + try: + model = validate_embedding_model(payload.embedding_model) + except ValueError as exc: + raise log_and_http_error( + exc, + 400, + safe_error_detail(exc, fallback = "Invalid embedding model."), + event = "settings.update_embedding_model_failed", + log = logger, + ) from exc + # The env/default model needs no verification; saving it is a no-op override. + # A local GGUF on the llama-server backend is accepted as-is: it is exactly + # what the backend loads, and HF metadata cannot verify a local path. + if ( + model != default_embedding_model() + and not payload.force + and not (_llama_backend_active() and _resolves_as_local_gguf(model)) + ): + hf_token = (payload.hf_token or "").strip() or None + from core.rag import config as rag_config + + # A GGUF-named repo on the llama-server backend is loaded from its .gguf + # files, which rarely carry sentence-transformers metadata; verify the + # GGUF is available (below) rather than the ST embedding-metadata gate, + # which would wrongly 409 a valid online GGUF embedder. + gguf_named = _llama_backend_active() and rag_config._names_gguf(model) + if not gguf_named and not is_embedding_model(model, hf_token = hf_token): + raise HTTPException( + status_code = 409, + detail = ( + f"Could not verify {model!r} as an embedding model on " + "Hugging Face (it may be the wrong model type, gated, or " + "you may be offline)." + ), + ) + gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) + if gguf_error: + raise HTTPException(status_code = 409, detail = gguf_error) + set_rag_embedding_model(model) + logger.info( + "settings.embedding_model_updated subject=%s model=%s forced=%s", + current_subject, + model, + payload.force, + ) + return _embedding_model_response() + + +@router.delete("/embedding-model", response_model = EmbeddingModelResponse) +def reset_embedding_model( + current_subject: str = Depends(get_current_subject), +) -> EmbeddingModelResponse: + """Clear the override, returning to the env/default model.""" + reset_rag_embedding_model() + logger.info("settings.embedding_model_reset subject=%s", current_subject) + return _embedding_model_response() + + class PreviewLinkRotateResponse(BaseModel): rotated: bool = True diff --git a/studio/backend/storage/rag_db.py b/studio/backend/storage/rag_db.py index ce27326562..cbd6ceb617 100644 --- a/studio/backend/storage/rag_db.py +++ b/studio/backend/storage/rag_db.py @@ -15,6 +15,7 @@ column type). """ import logging +import re import sqlite3 import threading @@ -64,7 +65,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: error TEXT, num_chunks INTEGER NOT NULL DEFAULT 0, stored_path TEXT, - created_at TEXT NOT NULL + created_at TEXT NOT NULL, + embedding_model TEXT ); CREATE INDEX IF NOT EXISTS idx_documents_scope ON documents(scope); CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(scope, sha256); @@ -107,6 +109,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()} if "project_id" not in cols: conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT") + # Lazy upgrade: which embedder produced a document's vectors (NULL = legacy, + # assumed current). Dedupe re-ingests when it no longer matches. + if "embedding_model" not in cols: + conn.execute("ALTER TABLE documents ADD COLUMN embedding_model TEXT") def get_connection() -> sqlite3.Connection: @@ -143,9 +149,32 @@ def get_connection() -> sqlite3.Connection: return conn +def vec_table_dim(conn: sqlite3.Connection) -> int | None: + """Embedding width baked into ``chunks_vec``, or None when absent.""" + row = conn.execute( + "SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'" + ).fetchone() + if row is None or not row["sql"]: + return None + m = re.search(r"float\[(\d+)\]", row["sql"]) + return int(m.group(1)) if m else None + + def ensure_vec(conn: sqlite3.Connection, dim: int) -> None: """Create the dense ``chunks_vec`` table once the embedding dim is known - (vec0 bakes it into the column type). Idempotent; dim fixed per db.""" + (vec0 bakes it into the column type). A width change (embedding model + switched in Settings) drops the table: the old vectors live in a foreign + space and would only block inserts, while lexical search keeps serving old + chunks until they are re-uploaded.""" + existing = vec_table_dim(conn) + if existing is not None and existing != int(dim): + logger.warning( + "chunks_vec dim changed %d -> %d (embedding model switched); dropping " + "stale dense index. Re-upload documents to restore dense search.", + existing, + int(dim), + ) + conn.execute("DROP TABLE chunks_vec") conn.execute( f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(" f"scope TEXT partition key, " diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index ba9f5b9cbc..41a9adcc29 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -22,7 +22,15 @@ logger = logging.getLogger(__name__) from typing import Any, Iterable, Optional -from utils.paths import project_workspaces_root, studio_db_path, ensure_dir +from utils.paths import ( + ensure_dir, + project_workspaces_root, + studio_db_path, +) +from utils.paths.external_media import is_linux_run_media_path +from utils.paths.sensitive import ( + contains_sensitive_path_component as _shared_contains_sensitive_path_component, +) from utils.training_runs import extract_project_name @@ -61,6 +69,14 @@ def _denied_path_prefixes() -> list[str]: return [] +def _contains_sensitive_path_component(path: str) -> bool: + return _shared_contains_sensitive_path_component(path) + + +def contains_sensitive_path_component(path: str) -> bool: + return _contains_sensitive_path_component(path) + + _schema_lock = threading.Lock() _schema_ready = False _SQLITE_IN_CHUNK_SIZE = 900 @@ -896,6 +912,8 @@ def add_scan_folder(path: str) -> dict: raise ValueError("Path must be a directory, not a file") if not os.access(normalized, os.R_OK | os.X_OK): raise ValueError("Path is not readable") + if _contains_sensitive_path_component(normalized): + raise ValueError("Credential or configuration directories are not allowed") # Windows: normcase for the denylist check but store original casing # so consumers see the native drive-letter casing (e.g. C:\Models). @@ -903,6 +921,8 @@ def add_scan_folder(path: str) -> dict: check = os.path.normcase(normalized) if is_win else normalized for prefix in _denied_path_prefixes(): if check == prefix or check.startswith(prefix + os.sep): + if prefix == "/run" and is_linux_run_media_path(check): + continue raise ValueError(f"Path under {prefix} is not allowed") conn = get_connection() diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 92a87ce045..a6c1fcda9c 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -128,13 +128,7 @@ class TestToolActionNudge: assert "call render_html once" in nudge def test_balanced_nudge_empty_without_known_tool_categories(self): - assert ( - _build_tool_action_nudge( - tools = [], - model_name = "Llama-3.1-8B-Instruct", - ) - == "" - ) + assert _build_tool_action_nudge(tools = [], model_name = "Llama-3.1-8B-Instruct") == "" # ===================================================================== diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 42c400383e..5d14c5c5bd 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -61,11 +61,16 @@ from core.inference.llama_cpp import LlamaCppBackend MIB = 1024 * 1024 -def _backend(vocab = 248320, embd = 5120): +def _backend( + vocab = 248320, + embd = 5120, + mla = None, +): """Backend with just the dims the compute-buffer estimate reads.""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._vocab_size = vocab b._embedding_length = embd + b._key_length_mla = mla # non-None -> MLA (compressed attention) return b @@ -150,3 +155,138 @@ class TestParallel1Default: def test_default_n_parallel(self): est = _backend()._estimate_compute_buffer_bytes() / MIB assert est < 128 + + +class TestContextLinearBuffer: + """``_compute_buffer_ctx_bytes``: the flash-attn KQ-mask + attention scratch + grow ~linearly with context; the flat estimate above only covers ctx -> 0. + Measured slope (q8_0 KV, ubatch 512) was 0.74-2.02 x n_embd; 2 x n_embd is the + worst-case upper bound the term must hold to.""" + + # (model, n_embd, ctx, measured CUDA0 compute buffer MiB at that ctx, q8_0/ub512) + _MEASURED = [ + ("Qwen3.5-2B", 2048, 262144, 796), + ("Qwen3.5-4B", 2560, 262144, 1330), # worst slope, 2.02 x n_embd + ("Qwen3.5-9B", 4096, 262144, 1336), + ("Qwen3.6-27B", 5120, 262144, 1360), + ("Gemma-4-31B", 5376, 262144, 2392), + ] + + def test_zero_by_default(self): + # Omitted/zero ctx -> no term (keeps the flat callers unchanged). + assert _backend()._compute_buffer_ctx_bytes(0) == 0 + + def test_zero_when_embd_missing(self): + assert _backend(embd = None)._compute_buffer_ctx_bytes(262144) == 0 + + def test_grows_linearly_with_context(self): + b = _backend(embd = 4096) + a = b._compute_buffer_ctx_bytes(65536) + d = b._compute_buffer_ctx_bytes(131072) + assert d == pytest.approx(2 * a, rel = 1e-6) + + def test_scales_with_embd(self): + # The quantized (dequant-scratch) rate scales with n_embd; f16 (mask) does not. + small = _backend(embd = 2048)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + big = _backend(embd = 5120)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + assert big > small + + def test_scales_with_ubatch(self): + b = _backend(embd = 4096) + lo = b._compute_buffer_ctx_bytes(131072, n_ubatch = 256) + hi = b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) + assert hi > lo + + @pytest.mark.parametrize("name,embd,ctx,measured", _MEASURED) + def test_upper_bounds_measured_compute_growth(self, name, embd, ctx, measured): + # flat term + context-linear term must cover the real (q8_0) buffer at full ctx. + b = _backend(embd = embd) + flat = b._estimate_compute_buffer_bytes(n_parallel = 1) + total = (flat + b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0")) / MIB + assert total >= measured, f"{name}: under-reserved {total:.0f} < {measured}" + + def test_worst_case_rate_covers_two_x_embd(self): + # >= 2 x n_embd bytes per context token at the default micro-batch (the worst + # measured quantized slope, Qwen3.5-4B), so flat + term upper-bounds the buffer. + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "q8_0") / 100000 + assert per_tok >= 2 * embd + + +class TestContextBufferKVQuant: + """The context-linear rate depends on the KV cache type: a quantized cache adds a + context-sized dequant scratch (heavy); f16/bf16/f32 only pays the KQ mask (light). + Measured Qwen3.5-4B at 256k: 1.30 GiB (q8_0) vs 0.31 GiB (f16).""" + + def test_quantized_heavier_than_f16(self): + b = _backend(embd = 4096) + q = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0") + f = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + assert q > f + + def test_none_cache_type_is_f16(self): + # None -> f16 (llama.cpp's default); the env-quantized case is covered by the + # KV budget's f16 over-reservation, so we take the lighter mask-only rate. + b = _backend(embd = 4096) + assert b._compute_buffer_ctx_bytes( + 131072, cache_type_kv = None + ) == b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16") + + @pytest.mark.parametrize("ct", ["f16", "bf16", "f32"]) + def test_unquantized_uses_mask_only_rate(self, ct): + # f16/bf16/f32: KQ mask only, n_ubatch*2 B/tok, independent of n_embd. + b_small = _backend(embd = 2048) + b_big = _backend(embd = 8192) + per_small = b_small._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + per_big = b_big._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_small == per_big # no n_embd scaling on the f16 path + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY # ubatch 512 + assert per_small == pytest.approx(expected, rel = 1e-6) + + @pytest.mark.parametrize("ct", ["q8_0", "q5_1", "q4_0", "iq4_nl"]) + def test_quantized_types_use_heavy_rate(self, ct): + embd = 4096 + b = _backend(embd = embd) + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000 + assert per_tok == pytest.approx( + LlamaCppBackend._CTX_COMPUTE_BYTES_PER_EMBD * embd, rel = 1e-6 + ) + + def test_f16_covers_measured_mask(self): + # f16 buffer is ~mask only (~n_ubatch*2 B/tok); 0.5 x n_embd must cover the + # measured Qwen3.5-4B f16 slope (~0.4 x n_embd = 0.31 GiB at 256k). + b = _backend(embd = 2560) # Qwen3.5-4B + est = b._compute_buffer_ctx_bytes(262144, cache_type_kv = "f16") / MIB + assert est >= 320 # measured 0.31 GiB growth + + +class TestContextBufferMLA: + """MLA (compressed attention) needs a smaller quantized dequant scratch than + regular attention: measured 0.94 x n_embd on GLM-5.2 and Kimi-K2.7 vs up to + 2.02x on Qwen/Gemma. Charging the regular rate would badly over-reserve a tight + multi-GPU MLA pin (per-device scaling multiplies the error).""" + + def test_mla_lighter_than_regular(self): + reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + assert mla < reg + + @pytest.mark.parametrize( + "name,embd,ctx,measured", + [ + ("GLM-5.2", 6144, 754688, 4141), # per-device compute MiB at q8_0 + ("Kimi-K2.7", 7168, 262144, 1690), + ], + ) + def test_mla_rate_covers_measured(self, name, embd, ctx, measured): + b = _backend(embd = embd, mla = 256) + est = b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0") / MIB + assert est >= measured, f"{name}: MLA under-reserved {est:.0f} < {measured}" + + def test_mla_not_wildly_over(self): + # 1.25 x n_embd should stay within ~1.6x of the measured 0.94x (not 2.4x like + # the regular 2.25 rate would), so a multi-GPU MLA pin keeps its context. + b = _backend(embd = 6144, mla = 256) + est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB + assert est <= 4141 * 1.7 diff --git a/studio/backend/tests/test_embedding_model_settings.py b/studio/backend/tests/test_embedding_model_settings.py new file mode 100644 index 0000000000..3be4af0e32 --- /dev/null +++ b/studio/backend/tests/test_embedding_model_settings.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Test for the customizable RAG embedding model: a saved override becomes the +effective model and derives its GGUF companion for the llama-server backend.""" + +from pathlib import Path +import sys +import types as _types + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest + +import utils.embedding_model_settings as ems +from core.rag import config as rag_config + + +@pytest.fixture +def settings_store(monkeypatch): + """In-memory app_settings store patched under the module's lazy imports.""" + import storage.studio_db as studio_db + + store: dict = {} + monkeypatch.setattr( + studio_db, "get_app_setting", lambda key, fallback = None: store.get(key, fallback) + ) + monkeypatch.setattr( + studio_db, "upsert_app_settings", lambda settings: store.update(settings) or store + ) + ems._invalidate_cache() + yield store + ems._invalidate_cache() + + +def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeypatch): + """The core contract: with nothing stored the default is in effect; a saved + custom model becomes the effective embedding model and derives its -GGUF + companion (what the llama-server backend loads); reset clears the override.""" + monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False) + assert ems.get_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert rag_config.effective_gguf_repo() == rag_config.EMBED_GGUF_REPO + + assert ems.set_rag_embedding_model(" org/my-embedder ") == "org/my-embedder" + assert rag_config.effective_embedding_model() == "org/my-embedder" + assert rag_config.effective_gguf_repo() == "org/my-embedder-GGUF" + + assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL + assert ems.get_stored_embedding_model() is None diff --git a/studio/backend/tests/test_export_capability.py b/studio/backend/tests/test_export_capability.py new file mode 100644 index 0000000000..e04417f933 --- /dev/null +++ b/studio/backend/tests/test_export_capability.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for export capability gating. + +Export is supported iff ``get_device() in {CUDA, XPU, MLX}``, with a torch-aware reason otherwise +(pytorch_not_installed / no_accelerator / mlx_unavailable), and the backend must import without +PyTorch. The matrix mocks the hardware probes; wiring is checked with ast so it runs on CPU. +""" + +import ast +import builtins +from pathlib import Path + +import pytest + +import utils.hardware.hardware as hw + +_BACKEND = Path(__file__).resolve().parent.parent + + +def _src(rel): + return (_BACKEND / rel).read_text(encoding = "utf-8") + + +def _func_src(rel, name): + src = _src(rel) + node = next( + n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name + ) + return ast.get_source_segment(src, node) + + +# -- capability matrix -------------------------------------------------------------------------- + + +def _patch(monkeypatch, *, torch: bool, device, apple: bool): + monkeypatch.setattr(hw, "_has_torch", lambda: torch) + monkeypatch.setattr(hw, "get_device", lambda: device) + monkeypatch.setattr(hw, "is_apple_silicon", lambda: apple) + + +def test_cpu_with_torch_unsupported_no_accelerator(monkeypatch): + # PyTorch present but no accelerator: unsupported with no_accelerator, not "PyTorch missing". + _patch(monkeypatch, torch = True, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "no_accelerator" + assert "accelerator" in cap["export_unsupported_message"].lower() + # Must NOT tell a user with PyTorch installed to install PyTorch. + assert "PyTorch is not installed" not in cap["export_unsupported_message"] + + +def test_cuda_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.CUDA, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is True + assert cap["export_unsupported_reason"] is None + assert cap["export_unsupported_message"] is None + + +def test_xpu_with_torch_supports_export(monkeypatch): + _patch(monkeypatch, torch = True, device = hw.DeviceType.XPU, apple = False) + assert hw.export_capability()["export_supported"] is True + + +def test_mlx_without_torch_supports_export(monkeypatch): + # Apple Silicon MLX exports without PyTorch. + _patch(monkeypatch, torch = False, device = hw.DeviceType.MLX, apple = True) + assert hw.export_capability()["export_supported"] is True + + +def test_no_torch_non_apple_reports_pytorch_missing(monkeypatch): + _patch(monkeypatch, torch = False, device = hw.DeviceType.CPU, apple = False) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "pytorch_not_installed" + assert "PyTorch is not installed" in cap["export_unsupported_message"] + + +def test_apple_without_mlx_reports_mlx_unavailable(monkeypatch): + # Apple + CPU means the MLX stack is missing; reason is mlx_unavailable regardless of torch. + for has_torch in (False, True): + _patch(monkeypatch, torch = has_torch, device = hw.DeviceType.CPU, apple = True) + cap = hw.export_capability() + assert cap["export_supported"] is False + assert cap["export_unsupported_reason"] == "mlx_unavailable" + assert "MLX" in cap["export_unsupported_message"] + + +# -- import safety without PyTorch -------------------------------------------------------------- + + +def test_export_backend_imports_without_torch(monkeypatch): + """core/export/export.py must import on a --no-torch host (unsloth/torch blocked) and return a + clean 'PyTorch is not installed' message from an export attempt, not crash at import.""" + import importlib + import sys + + real_import = builtins.__import__ + + def blocking_import(name, *args, **kwargs): + top = name.split(".")[0] + if top in {"torch", "unsloth"}: + raise ImportError(f"simulated: {top} not installed") + return real_import(name, *args, **kwargs) + + # Drop any preloaded copies so the guarded import paths re-run under the block. + for m in [k for k in sys.modules if k.split(".")[0] in {"torch", "unsloth"}]: + monkeypatch.delitem(sys.modules, m, raising = False) + monkeypatch.delitem(sys.modules, "core.export.export", raising = False) + monkeypatch.setattr(builtins, "__import__", blocking_import) + + mod = importlib.import_module("core.export.export") + assert mod._IS_MLX is False + assert mod.torch is None + assert mod._export_runtime_available() is False + + be = mod.ExportBackend.__new__(mod.ExportBackend) + be.current_model = None + be.current_tokenizer = None + be.is_peft = False + be._audio_type = None + ok, message, out = be.export_merged_model("/tmp/does-not-matter") + assert ok is False + assert "PyTorch is not installed" in message + + +# -- endpoint / backend wiring (ast) ------------------------------------------------------------ + + +def test_main_endpoints_expose_export_capability(): + m = _src("main.py") + # Both system endpoints spread export_capability() into their response. + assert m.count("**export_capability()") >= 2 + assert '"/api/system/hardware"' in m and '"/api/system"' in m + + +def test_routes_guard_mutating_endpoints(): + r = _src("routes/export.py") + assert "def _ensure_export_supported()" in r + # load + all four export endpoints call the guard. + assert r.count("_ensure_export_supported()") >= 6 + + +def test_export_methods_check_runtime(): + e = _src("core/export/export.py") + assert "def _export_runtime_available()" in e + # Each export method returns the clear message when the runtime is missing. + assert e.count("_export_runtime_available()") >= 5 + assert "_PYTORCH_MISSING_MESSAGE" in e + + +def test_export_capability_reads_no_torch_helper(): + cap = _func_src("utils/hardware/hardware.py", "export_capability") + assert "_has_torch()" in cap and "DeviceType.MLX" in cap and "is_apple_silicon()" in cap diff --git a/studio/backend/tests/test_export_imatrix_compressed.py b/studio/backend/tests/test_export_imatrix_compressed.py index d914ff8651..f499390add 100644 --- a/studio/backend/tests/test_export_imatrix_compressed.py +++ b/studio/backend/tests/test_export_imatrix_compressed.py @@ -54,8 +54,7 @@ def test_merged_request_rejects_unknown_format(): def test_export_gguf_threads_imatrix_to_save_and_push(): - # imatrix_file must reach both save_pretrained_gguf and push_to_hub_gguf, but only via the - # conditional **imatrix_kw so a no-imatrix export never sends an unsupported keyword. + # imatrix_file must reach both save paths, but only via the conditional **imatrix_kw. g = _func_src("core/export/export.py", "export_gguf") assert g.count("**imatrix_kw") >= 2 assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g @@ -109,8 +108,139 @@ def test_export_merged_maps_compressed_to_save_method(): def test_compressed_hub_push_uploads_local_dir_without_recompressing(): - # A compressed Hub push must upload the already-built output_path, not re-run compression - # via push_to_hub_merged (which would compress a second time). + # A compressed / torchao Hub push must upload the built output_path, not re-quantize. m = _func_src("core/export/export.py", "export_merged_model") - assert "elif is_compressed and output_path and Path(output_path).is_dir():" in m + assert "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" in m assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m + + +# -- torchao portable FP8/INT8 (device-agnostic, no NVIDIA GPU) --------------------------------- + + +def test_merged_request_accepts_torchao_aliases(): + # Portable torchao aliases pass through compressed_method (validated in the backend registry). + for alias in ("torchao_fp8", "torchao_int8"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_routes_torchao_and_skips_nvidia_guard(): + m = _func_src("core/export/export.py", "export_merged_model") + # torchao is classified separately and its suffix comes from the torchao normalizer. + assert "_normalize_torchao_method(compressed_alias)" in m + assert "is_torchao = torchao_info is not None" in m + assert "is_compressed = compressed_alias is not None and not is_torchao" in m + # The NVIDIA guard applies to compressed-tensors only, not torchao. + assert "_has_nvidia_gpu()" in m + # torchao routes through save_method just like compressed. + assert "elif is_compressed or is_torchao:" in m + + +def test_export_merged_nvidia_guard_present(): + m = _func_src("core/export/export.py", "export_merged_model") + assert "requires an NVIDIA GPU" in m + + +def test_has_nvidia_gpu_helper_reads_hardware_module(): + h = _func_src("core/export/export.py", "_has_nvidia_gpu") + assert "DeviceType.CUDA" in h and "IS_ROCM" in h + + +def test_export_merged_relaxes_is_peft_guard(): + # Non-PEFT (Local/HF base) models can now export merged; the old hard block must be gone. + m = _func_src("core/export/export.py", "export_merged_model") + assert "Use 'Export Base Model' instead." not in m + + +def test_unsloth_save_has_torchao_registry_and_path(): + # Read unsloth/save.py as text (not import) so this runs in the CPU suite without unsloth. + save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text(encoding = "utf-8") + assert "def _normalize_torchao_method" in save_py + assert "def _unsloth_save_torchao" in save_py + assert "TORCHAO_EXPORT_SCHEMES = {" in save_py + # torchao aliases must map to (scheme, suffix) so the backend routes to the torchao path. + assert '"torchao_fp8": ("fp8", "torchao-fp8")' in save_py + assert '"torchao_int8": ("int8", "torchao-int8")' in save_py + + +# -- GGUF multi-quant list ---------------------------------------------------------------------- + + +def test_gguf_request_accepts_list_of_quants(): + r = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"]) + assert r.quantization_method == ["Q4_K_M", "Q8_0"] + r2 = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = "Q4_K_M") + assert r2.quantization_method == "Q4_K_M" + + +def test_export_gguf_normalizes_quant_list(): + g = _func_src("core/export/export.py", "export_gguf") + assert "isinstance(quantization_method, (list, tuple))" in g + assert "quant_methods" in g + + +# -- GGUF LoRA adapter export ------------------------------------------------------------------- + + +def test_lora_request_has_gguf_fields(): + from models.export import ExportLoRAAdapterRequest + + r = ExportLoRAAdapterRequest(save_directory = "/tmp/x") + assert r.gguf is False and r.gguf_outtype == "q8_0" + r2 = ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0") + assert r2.gguf is True and r2.gguf_outtype == "q8_0" + + +def test_lora_request_rejects_bad_outtype(): + from models.export import ExportLoRAAdapterRequest + with pytest.raises(ValidationError): + ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf_outtype = "q3_k") + + +def test_export_lora_wires_gguf_save_method(): + la = _func_src("core/export/export.py", "export_lora_adapter") + assert 'save_method = "lora"' in la + assert "quantization_method = outtype" in la + + +def test_orchestrator_and_worker_pass_lora_gguf(): + o = _func_src("core/export/orchestrator.py", "export_lora_adapter") + assert '"gguf": gguf' in o and '"gguf_outtype": gguf_outtype' in o + w = _src("core/export/worker.py") + assert 'gguf = cmd.get("gguf", False)' in w + assert 'gguf_outtype = cmd.get("gguf_outtype", "q8_0")' in w + + +def test_route_passes_lora_gguf(): + r = _src("routes/export.py") + assert "gguf = request.gguf" in r and "gguf_outtype = request.gguf_outtype" in r + + +# -- compressed_method ("all formats" dropdown) ------------------------------------------------- + + +def test_merged_request_accepts_compressed_method(): + # Defaults to None; any scheme alias is accepted (validation happens in the backend registry). + assert ExportMergedModelRequest(save_directory = "/tmp/x").compressed_method is None + for alias in ("fp8", "fp8_static", "w8a8", "w8a16", "w4a16", "mxfp4", "mxfp8", "nvfp4"): + r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias) + assert r.compressed_method == alias + + +def test_export_merged_resolves_alias_via_registry(): + # The scheme + suffix must come from unsloth.save's registry normalizer, not a hardcoded dict. + m = _func_src("core/export/export.py", "export_merged_model") + assert "compressed_method" in m + assert "_normalize_compressed_method(compressed_alias)" in m + assert "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m + assert "compressed_suffix" in m and 'f"{save_directory}-{compressed_suffix}"' in m + + +def test_orchestrator_and_worker_pass_compressed_method(): + o = _func_src("core/export/orchestrator.py", "export_merged_model") + assert "compressed_method" in o and '"compressed_method": compressed_method' in o + assert 'compressed_method = cmd.get("compressed_method")' in _src("core/export/worker.py") + + +def test_route_passes_compressed_method(): + assert "compressed_method = request.compressed_method" in _src("routes/export.py") diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d96f88e4a6..69ad560788 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -5,9 +5,11 @@ import asyncio import importlib.util import os import re +import sys import unittest +from contextlib import nullcontext from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import patch from fastapi import HTTPException @@ -22,6 +24,7 @@ from utils.hardware import ( estimate_required_model_memory_gb, get_backend_visible_gpu_info, get_device_map, + get_gpu_utilization, get_offloaded_device_map_entries, get_parent_visible_gpu_ids, get_visible_gpu_utilization, @@ -33,6 +36,24 @@ import utils.hardware.hardware as _hw_module _BACKEND_ROOT = Path(__file__).resolve().parent.parent +async def _inline_to_thread(func, /, *args, **kwargs): + return func(*args, **kwargs) + + +def _fake_unsloth_attention_modules(resolver): + unsloth_module = ModuleType("unsloth") + models_module = ModuleType("unsloth.models") + utils_module = ModuleType("unsloth.models._utils") + utils_module.resolve_attention_implementation = resolver + models_module._utils = utils_module + unsloth_module.models = models_module + return { + "unsloth": unsloth_module, + "unsloth.models": models_module, + "unsloth.models._utils": utils_module, + } + + def _load_route_module(name: str, relative_path: str): spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path) module = importlib.util.module_from_spec(spec) @@ -122,6 +143,139 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase): + def test_gpu_utilization_preserves_primary_shape_with_devices(self): + devices = [ + { + "index": 5, + "visible_ordinal": 0, + "gpu_utilization_pct": 11.0, + "temperature_c": 40.0, + "vram_used_gb": 4.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 16.7, + "power_draw_w": 80.0, + "power_limit_w": 300.0, + "power_utilization_pct": 26.7, + }, + { + "index": 3, + "visible_ordinal": 1, + "gpu_utilization_pct": 22.0, + "temperature_c": 50.0, + "vram_used_gb": 8.0, + "vram_total_gb": 24.0, + "vram_utilization_pct": 33.3, + "power_draw_w": 120.0, + "power_limit_w": 300.0, + "power_utilization_pct": 40.0, + }, + ] + + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), + patch.object(_hw_module, "IS_ROCM", False), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = {"raw": "5,3", "numeric_ids": [5, 3]}, + ), + patch( + "utils.hardware.hardware._smi_query", + return_value = { + "available": True, + "devices": devices, + "backend_cuda_visible_devices": "5,3", + "parent_visible_gpu_ids": [5, 3], + "index_kind": "physical", + }, + ), + ): + result = get_gpu_utilization() + + self.assertIsInstance(result, dict) + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "cuda") + self.assertEqual(result["index"], 5) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual(result["vram_total_gb"], 24.0) + self.assertEqual(result["parent_visible_gpu_ids"], [5, 3]) + self.assertEqual([device["index"] for device in result["devices"]], [5, 3]) + + def test_gpu_utilization_cpu_returns_legacy_unavailable_object(self): + with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): + result = get_gpu_utilization() + + self.assertEqual(result, {"available": False, "backend": "cpu", "devices": []}) + + def test_gpu_utilization_mlx_stays_available_without_agx_stats(self): + fake_psutil = ModuleType("psutil") + fake_psutil.virtual_memory = lambda: SimpleNamespace(total = 64 * 1024**3) + + with ( + patch.dict(sys.modules, {"psutil": fake_psutil}), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX), + patch("utils.hardware.hardware._read_apple_gpu_stats", return_value = {}), + patch( + "core.training.get_training_backend", + return_value = SimpleNamespace(_progress = None), + ), + patch("utils.hardware.apple.read_gpu_temperature_c", return_value = None), + patch("utils.hardware.apple.read_gpu_power_w", return_value = None), + ): + result = get_gpu_utilization() + + self.assertTrue(result["available"]) + self.assertEqual(result["backend"], "mlx") + self.assertIsNone(result["gpu_utilization_pct"]) + self.assertEqual(result["vram_used_gb"], 0) + self.assertEqual(result["vram_total_gb"], 64.0) + self.assertEqual(len(result["devices"]), 1) + + def test_gpu_utilization_xpu_uses_visible_devices(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = { + "available": True, + "backend": "xpu", + "parent_visible_gpu_ids": [2, 0], + "index_kind": "physical", + "devices": [ + { + "index": 2, + "visible_ordinal": 1, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 3.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 18.8, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + { + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": 1.0, + "vram_total_gb": 16.0, + "vram_utilization_pct": 6.3, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + }, + ], + }, + ), + ): + result = get_gpu_utilization() + + self.assertEqual(result["backend"], "xpu") + self.assertEqual(result["index"], 0) + self.assertEqual(result["visible_ordinal"], 0) + self.assertEqual([device["index"] for device in result["devices"]], [0, 2]) + def test_visible_gpu_utilization_filters_to_parent_visible_ids(self): smi_output = "\n".join( [ @@ -272,6 +426,14 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_get_offloaded_device_map_entries_handles_models_without_device_map(self): self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {}) + @patch( + "utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate", + new = lambda model_name, **_: model_name, + ) + @patch( + "utils.hardware.hardware._load_config_for_gpu_estimate", + new = lambda *_args, **_kwargs: None, + ) def test_estimate_required_memory_formulas(self): eight_gb = 8 * (1024**3) @@ -432,6 +594,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self): with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.resolve_requested_gpu_ids", return_value = [2, 3], @@ -464,6 +627,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase): def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "utils.hardware.hardware.estimate_required_model_memory_gb", return_value = ( @@ -582,6 +746,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): with ( patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True), + patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA), patch( "core.training.training._CTX.Queue", side_effect = [dummy_queue, dummy_queue], @@ -709,14 +874,23 @@ class TestRouteErrors(unittest.TestCase): has_audio_input = False, ) - with patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -835,9 +1009,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -849,6 +1023,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -856,7 +1037,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -899,9 +1080,9 @@ class TestRouteErrors(unittest.TestCase): with ( patch.object( - inference_route.ModelConfig, - "from_identifier", - return_value = model_config, + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), patch.object( inference_route, @@ -913,6 +1094,13 @@ class TestRouteErrors(unittest.TestCase): "get_llama_cpp_backend", return_value = SimpleNamespace(is_loaded = False), ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ), + patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), + patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), patch( "core.export.get_export_backend", return_value = SimpleNamespace(current_checkpoint = None), @@ -920,7 +1108,7 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model( + inference_route._load_model_impl( request, SimpleNamespace( app = SimpleNamespace( @@ -1102,10 +1290,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): cfg._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) @@ -1133,10 +1318,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): 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, - ), + patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)), ): result = hardware_module._determine_attention_impl_for_gpu_estimate(config) @@ -1173,10 +1355,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase): inner._attn_implementation = "eager" return "eager" - with patch( - "unsloth.models._utils.resolve_attention_implementation", - side_effect = _stub_resolver, - ): + with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)): hardware_module._determine_attention_impl_for_gpu_estimate(config) self.assertFalse(hasattr(config, "_attn_implementation")) diff --git a/studio/backend/tests/test_linux_external_media_paths.py b/studio/backend/tests/test_linux_external_media_paths.py new file mode 100644 index 0000000000..c763248f6a --- /dev/null +++ b/studio/backend/tests/test_linux_external_media_paths.py @@ -0,0 +1,287 @@ +# 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 ast +import os +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Optional + +import pytest + +from hub.storage import scan_folders +from storage import studio_db +from utils.paths import external_media + + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent + + +class _ExistingScanFolderConn: + def __init__(self): + self.params = () + + def execute( + self, + _sql, + params = (), + ): + self.params = params + return self + + def fetchone(self): + return {"id": 1, "path": self.params[0], "created_at": "fake"} + + def commit(self): + pass + + def close(self): + pass + + +class _HTTPException(Exception): + def __init__(self, status_code: int, detail: str): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + +def _stub_linux_path_checks(monkeypatch, module): + monkeypatch.setattr(module.platform, "system", lambda: "Linux") + monkeypatch.setattr(module.os.path, "realpath", os.path.normpath) + monkeypatch.setattr(module.os.path, "expanduser", lambda p: p) + monkeypatch.setattr(module.os.path, "exists", lambda _p: True) + monkeypatch.setattr(module.os.path, "isdir", lambda _p: True) + monkeypatch.setattr(module.os, "access", lambda _p, _mode: True) + + +def _stub_hub_scan_folder_db(monkeypatch): + monkeypatch.setattr(scan_folders, "_ensure_schema", lambda _conn: None) + monkeypatch.setattr(scan_folders, "get_connection", _ExistingScanFolderConn) + + +def _stub_legacy_scan_folder_db(monkeypatch): + monkeypatch.setattr(studio_db, "get_connection", _ExistingScanFolderConn) + + +def test_linux_run_media_policy_accepts_mounted_volume_descendants(monkeypatch): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB") + assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6") + + +@pytest.mark.parametrize( + "path", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_linux_run_media_policy_rejects_unrelated_run_paths(monkeypatch, path): + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + assert not external_media.is_linux_run_media_path(path) + + +def test_linux_run_media_mount_roots_lists_readable_volume_roots(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + mount = base / "dspofu" / "nvmeB" + sensitive_mount = base / "dspofu" / ".ssh" + sensitive_aws_mount = base / "dspofu" / ".aws" + other_user_mount = base / "other" / "backup" + incomplete = base / "dspofu-only" + mount.mkdir(parents = True) + sensitive_mount.mkdir() + sensitive_aws_mount.mkdir() + other_user_mount.mkdir(parents = True) + incomplete.mkdir() + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_target = base / "dspofu" / ".config" + normal_mount.mkdir(parents = True) + sensitive_target.mkdir() + alias = base / "dspofu" / "config-alias" + alias.symlink_to(sensitive_target, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_linux_run_media_mount_roots_skips_sensitive_resolved_descendant(monkeypatch, tmp_path): + base = tmp_path / "run" / "media" + normal_mount = base / "dspofu" / "nvmeB" + sensitive_descendant = normal_mount / ".ssh" / "models" + sensitive_descendant.mkdir(parents = True) + alias = base / "dspofu" / "models-alias" + alias.symlink_to(sensitive_descendant, target_is_directory = True) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + + roots = external_media.linux_run_media_mount_roots(base, user = "dspofu") + + assert roots == [normal_mount.resolve()] + + +def test_hub_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = scan_folders.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_hub_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + scan_folders.add_scan_folder(target) + + +def test_hub_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, scan_folders) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_hub_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + scan_folders.add_scan_folder("/run/media/dspofu/nvmeB/.ssh/models") + + +def test_legacy_scan_folder_accepts_linux_run_media_mount(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6" + + row = studio_db.add_scan_folder(target) + + assert row["path"] == target + + +@pytest.mark.parametrize( + "target", + [ + "/run", + "/run/media", + "/run/media/dspofu", + "/run/user/1000/models", + "/run/systemd/private", + "/run/not-media/dspofu/nvmeB", + ], +) +def test_legacy_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Path under /run is not allowed"): + studio_db.add_scan_folder(target) + + +def test_legacy_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch): + _stub_linux_path_checks(monkeypatch, studio_db) + monkeypatch.setattr(external_media.platform, "system", lambda: "Linux") + _stub_legacy_scan_folder_db(monkeypatch) + + with pytest.raises(ValueError, match = "Credential or configuration"): + studio_db.add_scan_folder("/run/media/dspofu/nvmeB/.aws/models") + + +def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path): + tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")) + function_names = { + "_build_browse_allowlist", + "_browse_relative_parts", + "_is_path_inside_allowlist", + "_match_browse_child", + "_normalize_browse_request_path", + "_resolve_browse_target", + } + functions = [ + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in function_names + ] + module = ast.Module(body = functions, type_ignores = []) + ast.fix_missing_locations(module) + + home = tmp_path / "home" + media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB" + model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6" + home.mkdir() + model_dir.mkdir(parents = True) + (media_root / ".ssh").mkdir() + + fake_paths = SimpleNamespace( + hf_default_cache_dir = lambda: tmp_path / "missing-default-hf", + legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf", + well_known_model_dirs = lambda: [], + studio_root = lambda: tmp_path / "missing-studio", + outputs_root = lambda: tmp_path / "missing-outputs", + exports_root = lambda: tmp_path / "missing-exports", + ) + fake_external_media = SimpleNamespace(linux_run_media_mount_roots = lambda: [media_root]) + fake_studio_db = SimpleNamespace( + list_scan_folders = lambda: [], + contains_sensitive_path_component = studio_db.contains_sensitive_path_component, + ) + monkeypatch.setitem(sys.modules, "utils.paths", fake_paths) + monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media) + monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db) + + ns = { + "HTTPException": _HTTPException, + "os": os, + "Path": Path, + "Optional": Optional, + "_safe_is_dir": lambda p: Path(p).is_dir(), + "_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf", + "logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None), + } + exec(compile(module, "", "exec"), ns) + + allowlist = ns["_build_browse_allowlist"]() + + assert media_root.resolve() in allowlist + assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve() + + with pytest.raises(_HTTPException) as exc: + ns["_resolve_browse_target"](str(media_root / ".ssh"), allowlist) + assert exc.value.status_code == 403 + + ssh_root = media_root / ".ssh" + with pytest.raises(_HTTPException) as exc_root: + ns["_resolve_browse_target"](str(ssh_root), [ssh_root]) + assert exc_root.value.status_code == 403 diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py index 6b26121cf8..246d810602 100644 --- a/studio/backend/tests/test_llama_cpp_start_failure_classification.py +++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py @@ -140,6 +140,16 @@ class TestOllamaAndFallback: msg = _classify("", None, None) assert "llama-server failed to start" in msg + def test_health_timeout_names_probe_not_generic(self): + # A live server that never returns 200 on /health must name the probe and + # proxy/context causes, not blame a bad GGUF (#5740). + msg = _classify( + "llama-server health check timed out after 600.0s", "/models/x.gguf", "local/x" + ) + assert "/health" in msg + assert "NO_PROXY" in msg + assert "GGUF file is valid" not in msg + class TestOsKillReturncode: """SIGKILL (-9) with no diagnostic output is the OOM killer and gets a named, diff --git a/studio/backend/tests/test_llama_cpp_wait_for_health.py b/studio/backend/tests/test_llama_cpp_wait_for_health.py index 1ba6c9f7b5..82c5b4931a 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_health.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_health.py @@ -67,6 +67,15 @@ class TestWaitForHealthResilience: monkeypatch.setattr(httpx, "get", lambda *a, **kw: ok_resp) assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True + def test_timeout_records_marker_for_classification(self, monkeypatch): + """A live-but-never-healthy server leaves a marker so the failure is + classified as a /health timeout, not a bad GGUF (#5740).""" + b = _make_backend() + b._process.poll.return_value = None + monkeypatch.setattr(httpx, "get", lambda *a, **kw: mock.Mock(status_code = 503)) + assert b._wait_for_health(timeout = 0.02, interval = 0.01) is False + assert any("health check timed out" in ln for ln in b._stdout_lines) + def test_read_error_loops_to_subprocess_poll(self, monkeypatch): """WinError 10054 (httpx.ReadError) must be swallowed; the next iteration sees the dead subprocess and returns False with a structured exit-code log.""" b = _make_backend() diff --git a/studio/backend/tests/test_load_progress_ready_fraction.py b/studio/backend/tests/test_load_progress_ready_fraction.py new file mode 100644 index 0000000000..2e499cd8c6 --- /dev/null +++ b/studio/backend/tests/test_load_progress_ready_fraction.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""load_progress() must report a complete load once llama-server is healthy. + +With layers offloaded to VRAM (-ngl) the server releases the mmap'd weight pages +after upload, so its VmRSS sinks back well below the shard total. The raw RSS +fraction would then sit at a partial (~8%) value forever and freeze a +fraction-driven progress bar even though the model is ready -- the "stuck around +8% on the second pass" symptom in #5740. In the ready phase the fraction must be +1.0 regardless of resident set size. +""" + +from __future__ import annotations + +import io +import sys +import types +from pathlib import Path +from unittest.mock import patch + +import pytest + +# Stub heavy/unavailable deps before importing the module under test, so a +# targeted run in the lightweight backend env (no structlog/httpx) still +# collects. setdefault keeps the real modules when they are installed. Mirrors +# test_llama_cpp_load_progress_matrix.py. +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +sys.modules.setdefault("structlog", types.ModuleType("structlog")) + +_httpx_stub = types.ModuleType("httpx") +for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", +): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + + +class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + +_httpx_stub.Timeout = _FakeTimeout +_httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, +) +sys.modules.setdefault("httpx", _httpx_stub) + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + + +def _backend( + gguf_path, + *, + healthy, + pid = 4321, +): + # Bare instance: exercise load_progress() without the heavy real __init__. + be = object.__new__(LlamaCppBackend) + be._process = types.SimpleNamespace(pid = pid) + be._gguf_path = str(gguf_path) + be._healthy = healthy + return be + + +def _gguf(tmp_path, size_bytes): + f = tmp_path / "model-Q4_K_M.gguf" + f.write_bytes(b"\0" * size_bytes) + return f + + +def test_ready_reports_complete_despite_low_rss(tmp_path, monkeypatch): + # Healthy, but VmRSS has dropped to ~8% of the shard total after VRAM upload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 # not 0.08 + assert p["bytes_loaded"] == p["bytes_total"] == 10000 + + +def test_mmap_phase_reports_raw_rss_fraction(tmp_path, monkeypatch): + # Still loading: the bar should track real residency, not jump to 1.0. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + p = be.load_progress() + assert p["phase"] == "mmap" + assert p["fraction"] == 0.08 + assert p["bytes_loaded"] == 800 + assert p["bytes_total"] == 10000 + + +def test_progress_fraction_is_monotonic(tmp_path, monkeypatch): + # RSS peaks during page-in, then drops after -ngl offload; the bar must hold + # its high-water mark instead of collapsing back to ~8% (#5740). + be = _backend(_gguf(tmp_path, 10000), healthy = False) + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 9000)) + assert be.load_progress()["fraction"] == 0.9 + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + p = be.load_progress() + assert p["fraction"] == 0.9 + assert p["bytes_loaded"] == 9000 + + +def test_ready_without_shard_size_still_completes(tmp_path, monkeypatch): + # bytes_total unknown (file unstattable): fraction must still read complete. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800)) + be = _backend(tmp_path / "missing.gguf", healthy = True) + p = be.load_progress() + assert p["phase"] == "ready" + assert p["fraction"] == 1.0 + assert p["bytes_total"] == 0 + + +def test_none_when_no_process(tmp_path): + be = _backend(_gguf(tmp_path, 10000), healthy = True) + be._process = None + assert be.load_progress() is None + + +def test_none_when_rss_unreadable(tmp_path, monkeypatch): + # /proc unavailable (macOS/Windows) or unreadable -> no progress payload. + monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: None)) + be = _backend(_gguf(tmp_path, 10000), healthy = False) + assert be.load_progress() is None + + +def test_read_rss_bytes_absent_pid_is_none(): + # A pid with no readable /proc entry (or no /proc at all) yields None, never + # raises. + assert LlamaCppBackend._read_rss_bytes(2**31 - 1) is None + + +def test_read_rss_bytes_valueless_line_is_none(): + # A "VmRSS:" line with no value column must not raise (IndexError) -> None. + def fake_open(path, *a, **kw): + if str(path).startswith("/proc/"): + return io.StringIO("Name:\ttest\nVmRSS:\n") + return open(path, *a, **kw) + + with patch("builtins.open", side_effect = fake_open): + assert LlamaCppBackend._read_rss_bytes(4321) is None + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason = "/proc is Linux-only") +def test_read_rss_bytes_reads_self_on_linux(): + rss = LlamaCppBackend._read_rss_bytes(__import__("os").getpid()) + assert isinstance(rss, int) and rss > 0 diff --git a/studio/backend/tests/test_local_llama_cpp_link.py b/studio/backend/tests/test_local_llama_cpp_link.py new file mode 100644 index 0000000000..c78c029d91 --- /dev/null +++ b/studio/backend/tests/test_local_llama_cpp_link.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Behavioral tests for the --with-llama-cpp-dir 'unmanaged local link' contract. + +When the canonical llama.cpp dir is a symlink (POSIX) / junction (Windows) to a +user's own checkout, Studio must treat it as externally managed: + - the in-app updater must not offer or apply a prebuilt over the link + - orphan cleanup must not kill a llama-server the user launched from that tree + +These exercise real link behavior rather than grepping the scripts. +""" + +import os +import subprocess +from pathlib import Path + +import pytest + +from utils import llama_cpp_update as u +from core.inference.llama_cpp import LlamaCppBackend + + +def _make_link(link: Path, target: Path) -> None: + """Create a directory junction (Windows) / symlink (POSIX); neither needs + elevation.""" + target.mkdir(parents = True, exist_ok = True) + if os.name == "nt": + subprocess.run( + ["cmd", "/c", "mklink", "/J", str(link), str(target)], + check = True, + capture_output = True, + text = True, + ) + else: + link.symlink_to(target, target_is_directory = True) + + +def _server_subpath() -> Path: + return Path( + "build/bin/Release/llama-server.exe" if os.name == "nt" else "build/bin/llama-server" + ) + + +class _FakeProc: + def __init__(self, pid: int, exe: str) -> None: + self.info = {"pid": pid, "name": "llama-server", "exe": exe} + self.killed = False + + def kill(self) -> None: + self.killed = True + + +def test_is_external_link_detects_link_vs_plain_dir(tmp_path: Path) -> None: + plain = tmp_path / "plain" + plain.mkdir() + assert u._is_external_link(plain) is False + + link = tmp_path / "link" + _make_link(link, tmp_path / "tgt") + assert u._is_external_link(link) is True + + +def test_active_install_is_local_link(tmp_path: Path) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + binary = str(link / _server_subpath()) + assert u._active_install_is_local_link(binary) is True + + # A plain (non-link) llama.cpp dir is Studio-managed, not a local link. + plain = tmp_path / "plain" / "llama.cpp" + plain.mkdir(parents = True) + assert u._active_install_is_local_link(str(plain / _server_subpath())) is False + + +def test_get_update_status_reports_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + st = u.get_update_status() + assert st["supported"] is False + assert st["update_available"] is False + assert st["local_link"] is True + + +def test_start_update_refuses_local_link(tmp_path: Path, monkeypatch) -> None: + link = tmp_path / "llama.cpp" + _make_link(link, tmp_path / "tgt") + monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath())) + res = u.start_update() + assert res["started"] is False + assert res["reason"] == "local_link" + + +def _run_orphan_scan(monkeypatch, studio_root: Path, fake: _FakeProc) -> int: + # psutil drives the cross-platform process scan; skip (rather than error) if a + # minimal test env lacks it. CI installs it so these tests actually run. + psutil = pytest.importorskip("psutil") + + monkeypatch.setattr( + LlamaCppBackend, + "_resolved_studio_root_and_is_legacy", + staticmethod(lambda: (studio_root.resolve(), False)), + ) + monkeypatch.setattr(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)) + monkeypatch.setattr(psutil, "process_iter", lambda attrs = None: iter([fake])) + return LlamaCppBackend._kill_orphaned_servers() + + +def test_orphan_cleanup_spares_local_link_tree(tmp_path: Path, monkeypatch) -> None: + studio_root = tmp_path / "studio-home" + studio_root.mkdir() + external = tmp_path / "external" + (external / _server_subpath().parent).mkdir(parents = True) + (external / _server_subpath()).write_text("x") + _make_link(studio_root / "llama.cpp", external) + + exe_under_link = str((external / _server_subpath()).resolve()) + fake = _FakeProc(os.getpid() + 777, exe_under_link) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 0 + assert fake.killed is False + + +def test_orphan_cleanup_kills_under_real_root(tmp_path: Path, monkeypatch) -> None: + # Control: a real (non-link) managed root still gets its orphan reaped, so + # the spare-the-link test above is meaningful (not a no-op). + studio_root = tmp_path / "studio-home" + bin_dir = studio_root / "llama.cpp" / _server_subpath().parent + bin_dir.mkdir(parents = True) + exe = studio_root / "llama.cpp" / _server_subpath() + exe.write_text("x") + + fake = _FakeProc(os.getpid() + 888, str(exe.resolve())) + killed = _run_orphan_scan(monkeypatch, studio_root, fake) + assert killed == 1 + assert fake.killed is True diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index dce5e27c08..14fc0933d0 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -76,7 +76,7 @@ def test_mlx_studio_optimizer_aliases_are_explicit(): def test_mlx_studio_rejects_unknown_optimizer(): - with pytest.raises(ValueError, match = "Unsupported optimizer for MLX training"): + with pytest.raises(ValueError, match = "Supported"): _normalize_mlx_studio_optimizer("adamw_typo") diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index aa36c6fed4..ccbd78e2b1 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -40,6 +40,7 @@ from routes.inference import ( _effective_max_tokens, _extract_content_parts, _friendly_error, + _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, @@ -57,6 +58,32 @@ from routes.inference import ( from state.tool_policy import reset_tool_policy +class TestFriendlyUpstreamError: + def test_grammar_parse_failure_gets_actionable_message(self): + raw = '{"error":{"code":400,"message":"Failed to initialize samplers: failed to parse grammar","type":"invalid_request_error"}}' + msg = _friendly_upstream_error(raw) + assert "failed to parse grammar" not in msg # raw body is not surfaced verbatim + assert "tool-calling grammar" in msg and "Update Studio" in msg + + def test_failed_to_initialize_samplers_alone_matches(self): + assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers") + + def test_unrelated_error_passes_through(self): + assert _friendly_upstream_error("out of memory") == "llama-server error: out of memory" + + def test_openai_passthrough_error_rewrites_grammar_failure(self): + # OpenAI-compatible agents (opencode/openclaw/hermes/pi via /v1/chat/completions) + # get the same actionable message as the Anthropic passthrough, not the raw body. + from routes.inference import _openai_passthrough_error + + exc = _openai_passthrough_error( + 400, '{"error":{"message":"Failed to initialize samplers: failed to parse grammar"}}' + ) + assert "tool-calling grammar" in exc.detail + # An unrelated upstream error still passes through verbatim. + assert "llama-server error:" in _openai_passthrough_error(500, "disk full").detail + + # ===================================================================== # ChatMessage — tool role, tool_calls, optional content # ===================================================================== @@ -1829,6 +1856,553 @@ class TestApiMonitorProviderAndCompletionStreams: chunks = [chunk async for chunk in response.body_iterator] return SimpleNamespace(chunks = chunks, body = "".join(chunks), monitor = monitor) + def test_passthrough_stream_preheader_dispatched_with_timeout(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 400 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_request_error_in_window(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_send(*_args, **_kwargs): + raise httpx.ConnectError("connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + assert exc.value.status_code == 502 + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = b'{"error":"bad"}') + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "bad" in entry["error"] + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_keeps_error_envelope( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(400, content = ctx_msg.encode("utf-8")) + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + payload = json.loads(body.removeprefix("data: ").strip()) + assert payload["error"]["code"] == "context_length_exceeded" + assert payload["error"]["param"] == "messages" + assert isinstance(payload["error"], dict) + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_context_error_retries_truncation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + await gate.wait() + return httpx.Response(400, content = err_body) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + assert "data: [DONE]\n\n" in "".join(chunks) + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + cancel_id = "delayed-request-error-cancel" + + async def fake_send(*_args, **_kwargs): + await gate.wait() + raise httpx.ConnectError("delayed connectivity issue") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data:" in body + assert '"error"' in body + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "Lost connection" in entry["error"] + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_preheader_cancel_cleans_pending_send(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + entered = asyncio.Event() + cancelled = asyncio.Event() + cancel_id = "preheader-cancel-cleanup" + + async def fake_send(*_args, **_kwargs): + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelled.set() + raise + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + task = asyncio.create_task( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + ) + await asyncio.wait_for(entered.wait(), timeout = 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + returned = asyncio.Event() + cancel_id = "unstarted-completed-send-cleanup" + + class Stream(httpx.AsyncByteStream): + async def __aiter__(self): + if False: + yield b"" + + stream = Stream() + upstream_response = httpx.Response(200, stream = stream) + + async def fake_send(*_args, **_kwargs): + await gate.wait() + returned.set() + return upstream_response + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + assert cancel_id in inf_mod._CANCEL_REGISTRY + + gate.set() + await asyncio.wait_for(returned.wait(), timeout = 0.2) + await asyncio.sleep(0) + await response._unstarted_cleanup() + assert upstream_response.is_closed + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + def test_external_non_streaming_json_updates_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py new file mode 100644 index 0000000000..06316e2243 --- /dev/null +++ b/studio/backend/tests/test_passthrough_healing.py @@ -0,0 +1,1358 @@ +# 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 core/inference/passthrough_healing.py: promoting text-form +tool calls back into structured calls on the client-tool passthrough. The +route-level wiring (OpenAI / Anthropic / Responses endpoints) is covered in +their own endpoint test files; this file exercises the shared state machine +and helpers directly. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.passthrough_healing import ( # noqa: E402 + StreamToolCallHealer, + heal_gate, + heal_openai_message, + nudge_messages, + nudge_should_retry, + response_has_promotable_calls, +) + +TOOLS = [ + {"type": "function", "function": {"name": "Bash", "parameters": {}}}, + {"type": "function", "function": {"name": "Read", "parameters": {}}}, +] + +BASH_COMMAND_TOOL = { + "type": "function", + "function": { + "name": "Bash", + "parameters": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + }, +} +XML_BASH = '{"name":"Bash","arguments":{"cmd":"ls"}}' +XML_UNDECLARED = '{"name":"Nuke","arguments":{}}' + + +def _events_text(events): + return "".join(text for kind, text in events if kind == "text") + + +def _events_calls(events): + return [call for kind, call in events if kind == "tool_call"] + + +class TestHealGate: + def test_returns_declared_names(self): + assert heal_gate(None, TOOLS) == {"Bash", "Read"} + assert heal_gate(True, TOOLS) == {"Bash", "Read"} + + def test_opt_out_and_no_tools(self): + assert heal_gate(False, TOOLS) is None + assert heal_gate(None, []) is None + assert heal_gate(None, None) is None + + def test_malformed_tool_entries_ignored(self): + assert heal_gate(None, ["nonsense", {"function": "x"}, {}]) is None + + def test_tool_choice_none_disables(self): + assert heal_gate(None, TOOLS, "none") is None + + def test_tool_choice_forced_function_narrows_allowlist(self): + forced = {"type": "function", "function": {"name": "Bash"}} + assert heal_gate(None, TOOLS, forced) == {"Bash"} + + def test_tool_choice_forced_undeclared_function_disables(self): + forced = {"type": "function", "function": {"name": "Nuke"}} + assert heal_gate(None, TOOLS, forced) is None + + def test_tool_choice_auto_and_required_keep_full_set(self): + assert heal_gate(None, TOOLS, "auto") == {"Bash", "Read"} + assert heal_gate(None, TOOLS, "required") == {"Bash", "Read"} + + def test_tool_choice_unrecognized_dict_keeps_full_set(self): + assert heal_gate(None, TOOLS, {"type": "function"}) == {"Bash", "Read"} + + +class TestHealOpenaiMessage: + def test_promotes_xml_and_strips_content(self): + msg = {"role": "assistant", "content": XML_BASH} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] is None + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert json.loads(call["function"]["arguments"]) == {"cmd": "ls"} + + def test_keeps_surrounding_prose(self): + msg = {"role": "assistant", "content": f"Let me check.\n{XML_BASH}"} + assert heal_openai_message(msg, {"Bash"}) is True + assert msg["content"] == "Let me check." + + def test_undeclared_name_not_promoted(self): + msg = {"role": "assistant", "content": XML_UNDECLARED} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_UNDECLARED + assert "tool_calls" not in msg + + def test_structured_calls_untouched(self): + msg = {"role": "assistant", "content": XML_BASH, "tool_calls": [{"id": "x"}]} + assert heal_openai_message(msg, {"Bash"}) is False + assert msg["content"] == XML_BASH + + def test_prose_only_untouched(self): + msg = {"role": "assistant", "content": "just an answer"} + assert heal_openai_message(msg, {"Bash"}) is False + + def test_bare_string_arguments_use_schema_key(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, [BASH_COMMAND_TOOL]) is True + args = json.loads(msg["tool_calls"][0]["function"]["arguments"]) + assert args == {"command": "echo hi"} + + def test_bare_string_arguments_decline_ambiguous_schema(self): + msg = { + "role": "assistant", + "content": '{"name":"Bash","arguments":"echo hi"}', + } + assert heal_openai_message(msg, {"Bash"}, TOOLS) is False + assert "tool_calls" not in msg + + def test_mixed_declared_and_undeclared_promotes_declared_keeps_undeclared_text(self): + # Span-exact removal: only the promoted Bash markup is dropped; the + # undeclared Nuke call's text stays in the content byte-intact. + content = f"pre {XML_BASH} mid {XML_UNDECLARED} post" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + (call,) = msg["tool_calls"] + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in msg["content"] + assert "pre" in msg["content"] and "post" in msg["content"] + assert XML_BASH not in msg["content"] + + def test_multiple_declared_calls_all_promoted(self): + content = f"{XML_BASH} and {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert len(msg["tool_calls"]) == 2 + + def test_mixed_formats_promote_in_document_order(self): + func_read = "a.txt" + content = f"{func_read} then {XML_BASH}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash", "Read"}) is True + assert [call["function"]["name"] for call in msg["tool_calls"]] == ["Read", "Bash"] + assert msg["content"] == "then" + + def test_unparseable_closed_block_not_deleted(self): + # A closed block whose body never parses is model output, + # not a promotable call; it must survive promotion of its neighbor. + garbage = "not json at all" + content = f"{XML_BASH} {garbage}" + msg = {"role": "assistant", "content": content} + assert heal_openai_message(msg, {"Bash"}) is True + assert garbage in msg["content"] + + +class TestStreamHealer: + def test_plain_text_passes_through(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("hello ") + healer.feed("world") + healer.finalize() + assert _events_text(events) == "hello world" + assert not _events_calls(events) + + def test_complete_call_in_one_chunk(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"On it. {XML_BASH}") + healer.finalize() + assert _events_text(events) == "On it. " + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert healer.healed + + def test_signal_split_across_chunks(self): + healer = StreamToolCallHealer({"Bash"}) + events = [] + for piece in ["{"name":"Bash",', '"arguments":{}}']: + events += healer.feed(piece) + events += healer.finalize() + assert _events_text(events) == "" + assert len(_events_calls(events)) == 1 + + def test_closed_malformed_tool_block_flushes_immediately(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("not json after") + assert _events_text(events) == "not json after" + assert not _events_calls(events) + + def test_mixed_formats_stream_in_document_order(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + func_read = "a.txt" + events = healer.feed(f"{func_read} then {XML_BASH}") + healer.finalize() + assert [call["function"]["name"] for call in _events_calls(events)] == ["Read", "Bash"] + assert _events_text(events).strip() == "then" + + def test_false_alarm_html_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("use the
tag") + healer.finalize() + assert _events_text(events) == "use the
tag" + assert not _events_calls(events) + + def test_partial_signal_tail_held_then_flushed_at_end(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed("trailing text -> call B, never both calls then the text. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} middle {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds == ["tool_call", "text", "tool_call"] + assert events[1][1] == " middle " + + def test_undeclared_then_declared_keeps_document_order(self): + # The undeclared block precedes the declared call; its raw text must + # be emitted BEFORE the promoted call event, never after. + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_UNDECLARED} then {XML_BASH}") + healer.finalize() + kinds = [k for k, _ in events] + assert kinds.index("tool_call") == len(kinds) - 1 + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + assert XML_UNDECLARED in _events_text(events) + + def test_declared_promoted_then_late_undeclared_flushes_raw(self): + # Streaming causality: the declared call completed and was already + # emitted before the undeclared one arrived. The undeclared markup + # must still reach the client as raw text (no data loss). + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(f"{XML_BASH} then ") + assert len(_events_calls(events)) == 1 + events += healer.feed(XML_UNDECLARED) + healer.finalize() + assert XML_UNDECLARED in _events_text(events) + assert len(_events_calls(events)) == 1 + + def test_undeclared_tool_flushes_raw(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(XML_UNDECLARED) + healer.finalize() + assert _events_text(events) == XML_UNDECLARED + assert not _events_calls(events) + + def test_two_calls_and_text_between(self): + healer = StreamToolCallHealer({"Bash", "Read"}) + xml_read = '{"name":"Read","arguments":{"path":"f"}}' + events = healer.feed(f"{XML_BASH} then {xml_read}") + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["Bash", "Read"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events).strip() == "then" + + def test_incomplete_call_healed_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed('{"name":"Bash","arguments":{"cmd":"ls"}}') + assert events == [] # held + events = healer.finalize() + (call,) = _events_calls(events) + assert call["function"]["name"] == "Bash" + + def test_teaching_text_flushes_at_finalize(self): + healer = StreamToolCallHealer({"Bash"}) + events = healer.feed(" is the marker syntax") + healer.finalize() + assert _events_text(events) == " is the marker syntax" + assert not _events_calls(events) + + def test_hold_bound_flushes(self): + healer = StreamToolCallHealer({"Bash"}) + blob = "" + "x" * (64 * 1024 + 10) + events = healer.feed(blob) + healer.finalize() + assert _events_text(events) == blob + assert not _events_calls(events) + + def test_dormant_after_structured_delta(self): + healer = StreamToolCallHealer({"Bash"}) + held = healer.feed("prefix call Bash somehow???") + assert nudge_should_retry(data, {"Read"}) is True + + def test_no_retry_on_clean_prose(self): + assert nudge_should_retry(self._resp("all done"), {"Bash"}) is False + + def test_no_retry_when_heal_would_succeed(self): + assert nudge_should_retry(self._resp(XML_BASH), {"Bash"}) is False + + def test_no_retry_with_structured_calls(self): + data = self._resp("", tool_calls = [{"id": "x"}]) + assert nudge_should_retry(data, {"Bash"}) is False + + def test_no_retry_when_healing_disabled(self): + assert nudge_should_retry(self._resp("???"), None) is False + + def test_nudge_messages_shape(self): + data = self._resp("garbage") + suffix = nudge_messages(data, {"Bash", "Read"}) + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == "garbage" + assert "`Bash` or `Read`" in suffix[1]["content"] + + def test_retry_with_undeclared_structured_call_is_not_an_improvement(self): + # The retry replaces the original only when it carries a USABLE call: + # a structured call naming an undeclared tool must not count. + undeclared = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}} + ] + declared = [ + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}} + ] + assert response_has_promotable_calls(self._resp("", undeclared), {"Bash"}) is False + assert response_has_promotable_calls(self._resp("", declared), {"Bash"}) is True + + def test_retry_with_mixed_structured_calls_is_not_an_improvement(self): + # ALL structured calls must be declared: the caller forwards the whole + # list (and a parallel cap could keep only the FIRST), so a mixed retry + # could still hand the client an undeclared tool. + mixed = [ + {"id": "x", "type": "function", "function": {"name": "Nuke", "arguments": "{}"}}, + {"id": "y", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}, + ] + assert response_has_promotable_calls(self._resp("", mixed), {"Bash"}) is False + assert ( + response_has_promotable_calls(self._resp("", list(reversed(mixed))), {"Bash"}) is False + ) + + @pytest.mark.parametrize( + "data", + [ + None, + "not a dict", + {}, + {"choices": []}, + {"choices": [{}]}, + {"choices": [{"message": None}]}, # llama-server error bodies do this + {"choices": [{"message": "not a dict"}]}, + {"choices": [{"message": {"content": None}}]}, + {"error": {"message": "boom"}}, + ], + ) + def test_malformed_response_shapes_never_raise(self, data): + # A malformed upstream body must degrade to "nothing to heal/nudge", + # never crash the request with an AttributeError. + assert nudge_should_retry(data, {"Bash"}) is False + assert response_has_promotable_calls(data, {"Bash"}) is False + suffix = nudge_messages(data, {"Bash"}) + assert suffix[0] == {"role": "assistant", "content": ""} + + +# ── Route-level wiring (OpenAI passthrough) ───────────────────────────── +# Mirrors the fake-llama-server patterns in test_openai_tool_passthrough.py. + +import asyncio # noqa: E402 +import threading # noqa: E402 +from types import SimpleNamespace # noqa: E402 + +import httpx # noqa: E402 + +from core.inference.api_monitor import ApiMonitor # noqa: E402 +from models.inference import ChatCompletionRequest, ChatMessage # noqa: E402 +from routes.inference import ( # noqa: E402 + _openai_passthrough_non_streaming, + _openai_passthrough_stream, +) + +LOOKUP_TOOL = { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object", "properties": {}}}, +} +LOOKUP_XML = '{"name":"lookup","arguments":{"q":"x"}}' + + +def _payload(**kwargs): + defaults = dict( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [LOOKUP_TOOL], + ) + defaults.update(kwargs) + return ChatCompletionRequest(**defaults) + + +def _llama_backend(): + return SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + + +def _upstream_message( + content, + tool_calls = None, + finish_reason = "stop", +): + message = {"role": "assistant", "content": content} + if tool_calls is not None: + message["tool_calls"] = tool_calls + return { + "id": "chatcmpl-up", + "object": "chat.completion", + "created": 1, + "model": "gguf", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}, + } + + +class ScriptedClient: + """Fake nonstreaming_client() returning scripted JSON bodies, counting POSTs.""" + + def __init__(self, bodies): + self.bodies = list(bodies) + self.posts = [] + + async def post( + self, + _url, + json = None, + timeout = None, + ): + self.posts.append(json) + return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)]) + + +async def _drive_non_streaming(monkeypatch, payload, bodies): + import routes.inference as inf_mod + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _openai_passthrough_non_streaming( + _llama_backend(), payload, "gguf", monitor_id = None + ) + return client, json.loads(response.body) + + +async def _drive_stream(monkeypatch, payload, lines): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in lines: + yield line + + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monkeypatch.setattr(inf_mod, "api_monitor", ApiMonitor(max_entries = 3)) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + _llama_backend(), + payload, + "gguf", + "chatcmpl-test", + monitor_id = None, + ) + return [chunk async for chunk in response.body_iterator] + + +def _stream_payloads(chunks): + out = [] + for chunk in chunks: + for line in chunk.splitlines(): + if line.startswith("data: ") and line[6:] != "[DONE]": + out.append(json.loads(line[6:])) + return out + + +class TestOpenaiNonStreamingRoute: + def test_heals_xml_to_tool_calls(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(LOOKUP_XML)] + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "tool_calls" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + assert choice["message"]["content"] is None + assert data["usage"]["total_tokens"] == 3 # usage preserved + assert len(client.posts) == 1 # healing never re-requests + + asyncio.run(_run()) + + def test_bare_string_uses_client_schema_key(self, monkeypatch): + async def _run(): + content = '{"name":"Bash","arguments":"echo hi"}' + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tools = [BASH_COMMAND_TOOL]), + [_upstream_message(content)], + ) + (call,) = data["choices"][0]["message"]["tool_calls"] + assert json.loads(call["function"]["arguments"]) == {"command": "echo hi"} + + asyncio.run(_run()) + + def test_opt_out_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False), + [_upstream_message(LOOKUP_XML)], + ) + choice = data["choices"][0] + assert choice["message"]["content"] == LOOKUP_XML + assert "tool_calls" not in choice["message"] + assert choice["finish_reason"] == "stop" + + asyncio.run(_run()) + + def test_no_tools_untouched(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, _payload(tools = None), [_upstream_message(LOOKUP_XML)] + ) + assert data["choices"][0]["message"]["content"] == LOOKUP_XML + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await _drive_non_streaming(monkeypatch, _payload(), [_upstream_message(xml)]) + assert data["choices"][0]["message"]["content"] == xml + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_structured_calls_untouched(self, monkeypatch): + async def _run(): + native = [ + { + "id": "call_up", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message("", tool_calls = native, finish_reason = "tool_calls")], + ) + assert data["choices"][0]["message"]["tool_calls"] == native + + asyncio.run(_run()) + + def test_length_finish_reason_preserved(self, monkeypatch): + async def _run(): + # Truncated generation: the healed call stays attached but the + # client must still see the truncation, so length is never + # upgraded to tool_calls. + _, data = await _drive_non_streaming( + monkeypatch, + _payload(), + [_upstream_message(LOOKUP_XML, finish_reason = "length")], + ) + choice = data["choices"][0] + assert choice["finish_reason"] == "length" + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + + asyncio.run(_run()) + + def test_tool_choice_none_relays_verbatim(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = "none"), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_tool_choice_forcing_other_function_not_promoted(self, monkeypatch): + async def _run(): + _, data = await _drive_non_streaming( + monkeypatch, + _payload(tool_choice = {"type": "function", "function": {"name": "other"}}), + [_upstream_message(LOOKUP_XML)], + ) + message = data["choices"][0]["message"] + assert message["content"] == LOOKUP_XML + assert "tool_calls" not in message + + asyncio.run(_run()) + + def test_mixed_declared_and_undeclared_promotes_and_keeps_text(self, monkeypatch): + async def _run(): + rogue = '{"name":"rogue","arguments":{}}' + mixed = f"{LOOKUP_XML} also {rogue}" + _, data = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(mixed)] + ) + choice = data["choices"][0] + (call,) = choice["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert rogue in choice["message"]["content"] + assert choice["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_healed_then_native_stream_indexes_disjoint(self, monkeypatch): + async def _run(): + # A healed text-form call goes out first (index 0); a native + # structured delta follows. Clients merge deltas by index, so the + # native call must be shifted off index 0 or the two would merge. + native_line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_native","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + native_line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + indexes = {} + for payload_data in _stream_payloads(chunks): + for ch in payload_data.get("choices", []): + for tc in (ch.get("delta") or {}).get("tool_calls") or []: + indexes.setdefault(tc["index"], tc.get("id")) + assert indexes.get(0, "").startswith("call_") and indexes[0] != "call_native" + assert indexes.get(1) == "call_native" + + asyncio.run(_run()) + + def test_role_delta_precedes_healed_stream_content(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + "}}]}", + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + first_delta = payloads[0]["choices"][0]["delta"] + assert first_delta == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + + asyncio.run(_run()) + + def test_same_chunk_role_content_finish_delays_finish_until_after_healed_tool( + self, monkeypatch + ): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"role":"assistant","content":' + + json.dumps(LOOKUP_XML) + + '},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(stream = True), lines) + payloads = _stream_payloads(chunks) + assert payloads[0]["choices"][0]["finish_reason"] is None + assert payloads[0]["choices"][0]["delta"] == {"role": "assistant"} + assert "tool_calls" in payloads[1]["choices"][0]["delta"] + assert payloads[-1]["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + +GARBAGE_SIGNAL = "call lookup somehow???" + + +class TestNudgeRetryOpenai: + def test_retry_recovers_call(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 2 # exactly one retry + # Prefix byte-identical, nudge suffix appended (KV-cache reuse guard). + original, retry = client.posts + assert retry["messages"][: len(original["messages"])] == original["messages"] + suffix = retry["messages"][len(original["messages"]) :] + assert [m["role"] for m in suffix] == ["assistant", "user"] + assert suffix[0]["content"] == GARBAGE_SIGNAL + # The healed retry response is returned. + (call,) = data["choices"][0]["message"]["tool_calls"] + assert call["function"]["name"] == "lookup" + assert data["choices"][0]["finish_reason"] == "tool_calls" + + asyncio.run(_run()) + + def test_retry_still_garbage_returns_original(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(GARBAGE_SIGNAL + "2")], + ) + assert len(client.posts) == 2 + assert data["choices"][0]["message"]["content"] == GARBAGE_SIGNAL + assert "tool_calls" not in data["choices"][0]["message"] + + asyncio.run(_run()) + + def test_default_off_single_post(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, _payload(), [_upstream_message(GARBAGE_SIGNAL)] + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_on_clean_prose(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message("all done")], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + def test_no_retry_when_heal_succeeds(self, monkeypatch): + async def _run(): + client, data = await _drive_non_streaming( + monkeypatch, + _payload(nudge_tool_calls = True), + [_upstream_message(LOOKUP_XML)], + ) + assert len(client.posts) == 1 + assert data["choices"][0]["message"]["tool_calls"] + + asyncio.run(_run()) + + def test_heal_opt_out_disables_nudge_too(self, monkeypatch): + async def _run(): + client, _ = await _drive_non_streaming( + monkeypatch, + _payload(auto_heal_tool_calls = False, nudge_tool_calls = True), + [_upstream_message(GARBAGE_SIGNAL)], + ) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestNudgeRetryAnthropic: + async def _drive( + self, + monkeypatch, + bodies, + nudge = None, + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + nudge_tool_calls = nudge, + ) + return client, json.loads(response.body) + + def test_retry_recovers_tool_use(self, monkeypatch): + async def _run(): + client, data = await self._drive( + monkeypatch, + [_upstream_message(GARBAGE_SIGNAL), _upstream_message(LOOKUP_XML)], + nudge = True, + ) + assert len(client.posts) == 2 + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_healed_tool_use_precedes_trailing_text(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} done")]) + assert [block["type"] for block in data["content"]] == ["tool_use", "text"] + assert data["content"][1]["text"] == "done" + + asyncio.run(_run()) + + def test_default_off(self, monkeypatch): + async def _run(): + client, _ = await self._drive(monkeypatch, [_upstream_message(GARBAGE_SIGNAL)]) + assert len(client.posts) == 1 + + asyncio.run(_run()) + + +class TestAnthropicPassthroughHealingText: + """Non-streaming Anthropic passthrough must relay unpromoted (undeclared) + text-form calls as text, matching the OpenAI passthrough contract. Once + heal_openai_message promotes the declared call it span-trims only that + markup and deliberately leaves the undeclared bytes in the content; the + legacy blanket _TOOL_XML_RE strip must not delete them. + """ + + async def _drive(self, monkeypatch, upstream): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient([upstream]) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + ) + return json.loads(response.body) + + def test_mixed_declared_and_undeclared_relays_undeclared_as_text(self, monkeypatch): + async def _run(): + content = f"Running now. {LOOKUP_XML} then {XML_UNDECLARED} done." + data = await self._drive(monkeypatch, _upstream_message(content)) + # Declared lookup call is promoted into a structured tool_use block. + (tool_use,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_use["name"] == "lookup" + text = " ".join(b["text"] for b in data["content"] if b["type"] == "text") + assert XML_UNDECLARED in text + assert "Running now." in text and "done." in text + assert LOOKUP_XML not in text + + asyncio.run(_run()) + + +class TestAnthropicEmitterHealing: + def _events( + self, + emitter, + chunks, + finish = True, + ): + lines = [] + for chunk in chunks: + lines += emitter.feed_chunk(chunk) + if finish: + lines += emitter.finish() + return [json.loads(ln.split("data: ", 1)[1]) for ln in lines if "data: " in ln] + + def _emitter( + self, + allowed = ("lookup",), + **kwargs, + ): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() + emitter.enable_healing(set(allowed), **kwargs) + return emitter + + def _chunk( + self, + content = None, + tool_calls = None, + finish_reason = None, + ): + delta = {} + if content is not None: + delta["content"] = content + if tool_calls is not None: + delta["tool_calls"] = tool_calls + return {"choices": [{"delta": delta, "finish_reason": finish_reason}]} + + def test_xml_becomes_tool_use_block_and_stop_reason(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(finish_reason = "stop"), + ], + ) + starts = [e for e in events if e.get("type") == "content_block_start"] + (tool_start,) = [e for e in starts if e["content_block"]["type"] == "tool_use"] + assert tool_start["content_block"]["name"] == "lookup" + assert tool_start["content_block"]["id"].startswith("toolu_") + (args,) = [ + e["delta"]["partial_json"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "input_json_delta" + ] + assert json.loads(args) == {"q": "x"} + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "tool_use" + + def test_mid_block_signal_closes_text_block_first(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = f"Let me check {LOOKUP_XML}"), + self._chunk(finish_reason = "stop"), + ], + ) + kinds = [ + (e["type"], (e.get("content_block") or e.get("delta") or {}).get("type")) + for e in events + if e["type"].startswith("content_block") + ] + # text opens, streams the safe prefix, closes; then the tool_use block. + assert kinds[0] == ("content_block_start", "text") + assert kinds[1] == ("content_block_delta", "text_delta") + assert kinds[2] == ("content_block_stop", None) + assert kinds[3] == ("content_block_start", "tool_use") + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "Let me check " + + def test_false_alarm_streams_as_text(self): + events = self._events( + self._emitter(), + [self._chunk(content = "use the
tag"), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == "use the
tag" + (message_delta,) = [e for e in events if e.get("type") == "message_delta"] + assert message_delta["delta"]["stop_reason"] == "end_turn" + + def test_signal_split_across_chunks(self): + events = self._events( + self._emitter(), + [ + self._chunk(content = "{"name":"lookup","arguments":{"q":"y"}}' + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [self._chunk(content = two), self._chunk(finish_reason = "stop")], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_disable_parallel_drops_native_after_healed(self): + # A healed call consumed the single allowed slot; a later native + # structured call (index 0, so it survives the caller's chunk-level + # cap) must not open a second tool_use block. + structured = [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + events = self._events( + self._emitter(disable_parallel_tool_use = True), + [ + self._chunk(content = LOOKUP_XML), + self._chunk(tool_calls = structured), + self._chunk(finish_reason = "tool_calls"), + ], + ) + starts = [ + e + for e in events + if e.get("type") == "content_block_start" and e["content_block"]["type"] == "tool_use" + ] + assert len(starts) == 1 + + def test_no_healing_means_verbatim_text(self): + from core.inference.anthropic_compat import AnthropicPassthroughEmitter + + emitter = AnthropicPassthroughEmitter() # enable_healing never called + events = self._events( + emitter, + [self._chunk(content = LOOKUP_XML), self._chunk(finish_reason = "stop")], + ) + texts = [ + e["delta"]["text"] + for e in events + if e.get("type") == "content_block_delta" and e["delta"]["type"] == "text_delta" + ] + assert "".join(texts) == LOOKUP_XML + + +class TestAnthropicNonStreamingRoute: + async def _drive( + self, + monkeypatch, + bodies, + auto_heal = None, + tools = None, + tool_choice = "auto", + ): + import routes.inference as inf_mod + from routes.inference import _anthropic_passthrough_non_streaming + + client = ScriptedClient(bodies) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: client) + response = await _anthropic_passthrough_non_streaming( + _llama_backend(), + [{"role": "user", "content": "hi"}], + tools if tools is not None else [LOOKUP_TOOL], + 0.7, + 0.95, + None, + 256, + "msg_test", + "gguf", + tool_choice = tool_choice, + auto_heal_tool_calls = auto_heal, + ) + return client, json.loads(response.body) + + def test_promotes_xml_to_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive(monkeypatch, [_upstream_message(LOOKUP_XML)]) + (block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert block["name"] == "lookup" + assert block["input"] == {"q": "x"} + assert data["stop_reason"] == "tool_use" + assert not any(b["type"] == "text" for b in data["content"]) + + asyncio.run(_run()) + + def test_opt_out_keeps_legacy_strip(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(f"plan {LOOKUP_XML}")], auto_heal = False + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" # XML stripped, nothing promoted + + asyncio.run(_run()) + + def test_undeclared_tool_not_promoted(self, monkeypatch): + async def _run(): + xml = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(xml)]) + assert data["stop_reason"] == "end_turn" + assert not any(b["type"] == "tool_use" for b in data["content"]) + # Healing preserves what it does not promote: the undeclared call + # reaches the client as text instead of being silently stripped. + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert text_block["text"] == xml + + asyncio.run(_run()) + + def test_mixed_undeclared_text_preserved_after_heal(self, monkeypatch): + async def _run(): + # Declared call promoted to tool_use; the undeclared call's markup + # stays in the text block (the legacy strip must not run after a + # span-exact heal), matching the OpenAI passthrough. + rogue = '{"name":"rogue","arguments":{}}' + _, data = await self._drive(monkeypatch, [_upstream_message(f"{LOOKUP_XML} {rogue}")]) + (tool_block,) = [b for b in data["content"] if b["type"] == "tool_use"] + assert tool_block["name"] == "lookup" + (text_block,) = [b for b in data["content"] if b["type"] == "text"] + assert rogue in text_block["text"] + assert data["stop_reason"] == "tool_use" + + asyncio.run(_run()) + + def test_length_beats_tool_use(self, monkeypatch): + async def _run(): + _, data = await self._drive( + monkeypatch, [_upstream_message(LOOKUP_XML, finish_reason = "length")] + ) + assert data["stop_reason"] == "max_tokens" + assert any(b["type"] == "tool_use" for b in data["content"]) + + asyncio.run(_run()) + + def test_tool_choice_none_keeps_legacy_strip(self, monkeypatch): + async def _run(): + # Anthropic {"type": "none"} arrives here converted to "none": + # the request forbade tool calls, so nothing is promoted and the + # legacy XML strip applies as before healing existed. + _, data = await self._drive( + monkeypatch, + [_upstream_message(f"plan {LOOKUP_XML}")], + tool_choice = "none", + ) + assert data["stop_reason"] == "end_turn" + (block,) = data["content"] + assert block["type"] == "text" + assert block["text"] == "plan" + + asyncio.run(_run()) + + +class TestOpenaiStreamingRoute: + def test_heals_streamed_xml(self, monkeypatch): + async def _run(): + pieces = ["", '{"name":"lookup",', '"arguments":{"q":"x"}}', ""] + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":[{"index":0,"delta":{"content":%s}}]}' + % json.dumps(p) + for p in pieces + ] + lines += [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["function"]["name"] == "lookup" + assert json.loads(call["function"]["arguments"]) == {"q": "x"} + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + # None of the XML leaked as visible content. + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert "" not in text + assert chunks[-1] == "data: [DONE]\n\n" + + asyncio.run(_run()) + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + async def _run(): + # parallel_tool_calls=false: a healed call consumed the single + # allowed slot, and the upstream SSE cap keeps native index 0, so + # the route must drop the later native call itself. + xml = '{"name":"lookup","arguments":{"q":"x"}}' + native = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + 'data: {"id":"c1","model":"gguf","created":1,"choices":' + '[{"index":0,"delta":{"content":%s}}]}' % json.dumps(xml), + native, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(parallel_tool_calls = False), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + (call,) = tool_deltas + assert call["id"] == "call_0" # the healed call; native was dropped + + asyncio.run(_run()) + + def test_false_alarm_text_flushes(self, monkeypatch): + async def _run(): + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"use the
tag"}}]}', + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + text = "".join( + (c.get("delta") or {}).get("content") or "" + for p in payloads + for c in p.get("choices", []) + ) + assert text == "use the
tag" + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["stop"] + + asyncio.run(_run()) + + def test_incomplete_xml_healed_at_done(self, monkeypatch): + async def _run(): + # No close tag and no finish chunk: healed at the [DONE] boundary, + # synthetic finish must say tool_calls. + lines = [ + 'data: {"id":"c1","choices":[{"index":0,"delta":{"content":"{\\"name\\":\\"lookup\\",\\"arguments\\":{}}"}}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + payloads = _stream_payloads(chunks) + tool_deltas = [ + tc + for p in payloads + for c in p.get("choices", []) + for tc in (c.get("delta") or {}).get("tool_calls") or [] + ] + assert len(tool_deltas) == 1 + finishes = [ + c["finish_reason"] + for p in payloads + for c in p.get("choices", []) + if c.get("finish_reason") + ] + assert finishes == ["tool_calls"] + + asyncio.run(_run()) + + def test_structured_upstream_calls_relay_verbatim(self, monkeypatch): + async def _run(): + line = ( + 'data: {"id":"c1","choices":[{"index":0,"delta":{"tool_calls":' + '[{"index":0,"id":"call_up","type":"function","function":' + '{"name":"lookup","arguments":"{}"}}]}}]}' + ) + lines = [ + line, + 'data: {"id":"c1","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}', + "data: [DONE]", + ] + chunks = await _drive_stream(monkeypatch, _payload(), lines) + assert chunks[0] == line + "\n\n" # byte-for-byte relay + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_rag_embed_llama_server.py b/studio/backend/tests/test_rag_embed_llama_server.py index 8321068afd..0e1f74cefe 100644 --- a/studio/backend/tests/test_rag_embed_llama_server.py +++ b/studio/backend/tests/test_rag_embed_llama_server.py @@ -373,6 +373,8 @@ def test_ensure_ready_respawns_dead_process(monkeypatch): def fake_spawn(): spawned["n"] += 1 b._process = _FakeProc(alive = True) + # _current() now also checks the served repo, so mark it current. + b._model_repo = config.effective_gguf_repo() monkeypatch.setattr(b, "_spawn", fake_spawn) b._ensure_ready() diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 4c46f49495..14ab0efe2e 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -86,3 +86,212 @@ def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch): _table_pdf(pdf) pages = parsers.parse(str(pdf)) assert pages and "Quarter" in pages[0].text + + +def _long_text_pdf(path): + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + body = "The quick brown fox jumps over the lazy dog. " * 12 # >200 letters + page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), body, fontsize = 11) + doc.save(str(path)) + doc.close() + + +def test_pdf_markdown_corruption_falls_back_to_plain(tmp_path, monkeypatch): + # pymupdf4llm can emit shaped RTL Presentation Forms for Arabic/Hebrew; the parser + # detects that and uses PyMuPDF's logical-order text instead of the mangled Markdown. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + shaped = "".join(chr(c) for c in range(0xFE8D, 0xFEA0)) * 20 # heavy shaped forms + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: [shaped] * doc.page_count) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "Quarter" in text # real logical-order text recovered + assert not parsers._markdown_corrupted(text) # shaped garbage not carried through + + +def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch): + # If pymupdf4llm silently drops most of a page, the parser prefers the fuller raw layer. + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", True) + monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: ["x"] * doc.page_count) + pdf = tmp_path / "long.pdf" + _long_text_pdf(pdf) + text = "\n".join(p.text for p in parsers.parse(str(pdf))) + assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown + + +def _docx_with_table(path): + import docx + + document = docx.Document() + document.add_paragraph("Intro before table.") + table = document.add_table(rows = 2, cols = 2) + table.cell(0, 0).text = "NAME" + table.cell(0, 1).text = "SCORE" + table.cell(1, 0).text = "Alice" + table.cell(1, 1).text = "97pts" + document.add_paragraph("Outro after table.") + document.save(str(path)) + + +def test_docx_extracts_table_cells(tmp_path): + # document.paragraphs alone drops tables; the parser walks body content in order so + # table cells survive (pipe-joined, which the preview locator anchors on). + pytest.importorskip("docx") + from core.rag import parsers + + docx_path = tmp_path / "t.docx" + _docx_with_table(docx_path) + text = "\n".join(p.text for p in parsers.parse(str(docx_path))) + assert all(v in text for v in ("NAME", "SCORE", "Alice", "97pts")) # cells kept + assert "Alice | 97pts" in text # row cells joined + assert text.index("Intro") < text.index("NAME") < text.index("Outro") # order kept + + +def test_docx_table_keeps_columns_and_collapses_cell_newlines(tmp_path): + # Empty cells are kept (so columns stay aligned across rows) and a cell's internal + # newlines are collapsed to spaces (so a multi-paragraph cell can't break the row). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "A" + table.cell(0, 1).text = "" # empty middle cell + table.cell(0, 2).text = "C" + multiline = table.cell(1, 0) + multiline.text = "line1" + multiline.add_paragraph("line2") # cell now holds an internal newline + table.cell(1, 1).text = "mid" + table.cell(1, 2).text = "end" + path = tmp_path / "aligned.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "A | | C" in text # empty cell preserved -> columns line up + assert "line1 line2 | mid | end" in text # internal newline collapsed to a space + + +def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path): + # A horizontally merged cell repeats across the spanned columns: emit its text once + # then a placeholder, so the row keeps as many fields as its siblings (columns stay + # aligned) without duplicating the merged text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "WIDE" + table.cell(0, 2).text = "END" + table.cell(0, 0).merge(table.cell(0, 1)) # span the first two columns + table.cell(1, 0).text = "a" + table.cell(1, 1).text = "b" + table.cell(1, 2).text = "c" + path = tmp_path / "merged.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns + assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c" + assert "a | b | c" in text + + +def test_docx_table_pads_omitted_grid_columns(tmp_path): + # A row that skips leading grid columns exposes the gap via grid_cols_before; pad it + # with empty fields so the value stays under the right header instead of shifting left. + pytest.importorskip("docx") + import docx + from docx.oxml.ns import qn + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 2, cols = 3) + table.cell(0, 0).text = "H1" + table.cell(0, 1).text = "H2" + table.cell(0, 2).text = "H3" + tr = table.rows[1]._tr # drop the first cell and mark it skipped via + tr.remove(tr.tc_lst[0]) + trPr = tr.get_or_add_trPr() + trPr.insert(0, trPr.makeelement(qn("w:gridBefore"), {qn("w:val"): "1"})) + table.rows[1].cells[0].text = "X" # sits in column 2 + path = tmp_path / "gap.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert " | X | " in text # leading gap padded so X lines up under H2, not H1 + + +def test_docx_flattens_nested_table(tmp_path): + # cell.text ignores tables nested inside a cell; walk cell.tables so nested rows are + # not silently dropped from the indexed text. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + outer = document.add_table(rows = 1, cols = 1).cell(0, 0) + outer.text = "outer" + nested = outer.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + path = tmp_path / "nested.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert "NESTED-A | NESTED-B" in text # nested table flattened, not dropped + + +def test_docx_nested_table_keeps_in_cell_order(tmp_path): + # A cell holding paragraph, nested table, paragraph must serialize in that order + # (cell.text alone would emit both paragraphs before the nested rows). + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + cell = document.add_table(rows = 1, cols = 1).cell(0, 0) + cell.text = "before" + nested = cell.add_table(rows = 1, cols = 2) + nested.cell(0, 0).text = "NESTED-A" + nested.cell(0, 1).text = "NESTED-B" + cell.add_paragraph("after") + path = tmp_path / "nested_order.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.index("before") < text.index("NESTED-A") < text.index("after") + + +def test_docx_table_vertical_merge_emitted_once(tmp_path): + # A vertically merged cell maps every continuation row back to the origin ; + # emit it once and leave placeholders below so a row-spanning label isn't repeated. + pytest.importorskip("docx") + import docx + + from core.rag import parsers + + document = docx.Document() + table = document.add_table(rows = 3, cols = 2) + table.cell(0, 0).merge(table.cell(1, 0)).merge(table.cell(2, 0)).text = "SECTION" + table.cell(0, 1).text = "r0" + table.cell(1, 1).text = "r1" + table.cell(2, 1).text = "r2" + path = tmp_path / "vmerge.docx" + document.save(str(path)) + + text = "\n".join(p.text for p in parsers.parse(str(path))) + assert text.count("SECTION") == 1 # not repeated on each spanned row + assert "SECTION | r0" in text and " | r1" in text and " | r2" in text diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 4147746b54..a7ceb49ed9 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -1986,3 +1986,177 @@ class TestTranslatedMessagesValidate: msgs = _normalise_responses_input(payload) for m in msgs: ChatMessage(**m.model_dump(exclude_none = True)) + + +# ===================================================================== +# Streaming passthrough healing — text-form calls promoted in order +# ===================================================================== + + +class TestResponsesStreamHealing: + """Route-level healing on the /v1/responses stream: text-form tool calls + are promoted through the same per-call item state machinery as structured + deltas, and healer events keep their order (text around a healed call must + not move relative to the function_call item).""" + + _XML = '{"name":"lookup","arguments":{"q":"x"}}' + _TOOL = {"type": "function", "name": "lookup", "parameters": {"type": "object"}} + + @staticmethod + def _ordered_events(lines): + events = [] + for line in lines: + if not line.startswith("event: "): + continue + name, _, rest = line.partition("\n") + payload = json.loads(rest.split("data: ", 1)[1].strip()) + events.append((name[len("event: ") :], payload)) + return events + + def _run_stream(self, monkeypatch, content, **payload_kwargs): + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, [{"choices": [{"delta": {"content": content}}]}] + ) + payload = ResponsesRequest(input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + return self._ordered_events(asyncio.run(run())) + + def test_text_around_healed_call_keeps_order(self, monkeypatch): + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + pos_before = pos_item = pos_after = None + for i, (name, payload) in enumerate(events): + if name == "response.output_text.delta": + if "before" in payload["delta"] and pos_before is None: + pos_before = i + if "after" in payload["delta"]: + pos_after = i + if ( + name == "response.output_item.added" + and payload["item"]["type"] == "function_call" + and pos_item is None + ): + pos_item = i + assert payload["item"]["name"] == "lookup" + assert pos_before is not None and pos_item is not None and pos_after is not None + assert pos_before < pos_item < pos_after + + def test_call_before_trailing_text_claims_lower_output_index(self, monkeypatch): + events = self._run_stream(monkeypatch, f"{self._XML} done.") + item_added = [ + (name, payload) for name, payload in events if name == "response.output_item.added" + ] + # The call came first in the model output, so its item is added first + # and claims the lower output_index; the trailing text's message item + # follows. + assert [payload["item"]["type"] for _, payload in item_added] == [ + "function_call", + "message", + ] + call_idx = item_added[0][1]["output_index"] + msg_idx = item_added[1][1]["output_index"] + assert call_idx < msg_idx + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert "done." in text + assert "" not in text + + def test_tool_choice_none_streams_raw_text(self, monkeypatch): + events = self._run_stream(monkeypatch, self._XML, tool_choice = "none") + assert not any( + payload["item"]["type"] == "function_call" + for name, payload in events + if name == "response.output_item.added" + ) + text = "".join( + payload["delta"] for name, payload in events if name == "response.output_text.delta" + ) + assert text == self._XML + + def test_healed_call_splits_message_items(self, monkeypatch): + # Text on both sides of a healed call becomes TWO message items: the + # healed function_call closes the first, trailing text opens a fresh + # one with a later output index (native Responses stream shape). + events = self._run_stream(monkeypatch, f"before {self._XML} after.") + added = [ + (payload["output_index"], payload["item"]["type"], payload["item"].get("id")) + for name, payload in events + if name == "response.output_item.added" + ] + assert [item_type for _, item_type, _ in added] == [ + "message", + "function_call", + "message", + ] + assert [idx for idx, _, _ in added] == sorted(idx for idx, _, _ in added) + assert added[0][2] != added[2][2] # distinct message item ids + # Text deltas attribute to their OWN message item. + deltas = [ + (payload["item_id"], payload["delta"]) + for name, payload in events + if name == "response.output_text.delta" + ] + assert [d for i, d in deltas if i == added[0][2]] == ["before "] + assert [d for i, d in deltas if i == added[2][2]] == [" after."] + # The completed snapshot lists all three items with per-item text. + completed = [payload for name, payload in events if name == "response.completed"] + output = completed[0]["response"]["output"] + assert [item["type"] for item in output] == ["message", "function_call", "message"] + assert output[0]["content"][0]["text"] == "before " + assert output[2]["content"][0]["text"] == " after." + + def test_parallel_cap_drops_native_after_healed(self, monkeypatch): + # parallel_tool_calls=false: a healed call consumed the single allowed + # slot; a later native structured call (index 0, so it survives + # _drop_parallel_tool_call_deltas) must not open a second + # function_call item. + TestResponsesStreamAdapter._install_stream_mock( + monkeypatch, + [ + {"choices": [{"delta": {"content": self._XML}}]}, + { + "choices": [ + { + "delta": { + "tool_calls": [ + { + "index": 0, + "id": "call_up", + "function": {"name": "lookup", "arguments": "{}"}, + } + ] + } + } + ] + }, + ], + ) + payload = ResponsesRequest( + input = "hi", + stream = True, + tools = [self._TOOL], + parallel_tool_calls = False, + ) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream( + payload, messages, TestResponsesStreamAdapter._Request() + ) + return await TestResponsesStreamAdapter._collect(response) + + events = self._ordered_events(asyncio.run(run())) + calls = [ + payload + for name, payload in events + if name == "response.output_item.added" and payload["item"]["type"] == "function_call" + ] + assert len(calls) == 1 + assert calls[0]["item"]["name"] == "lookup" diff --git a/studio/backend/tests/test_slot_offload_fit.py b/studio/backend/tests/test_slot_offload_fit.py new file mode 100644 index 0000000000..ac606e4627 --- /dev/null +++ b/studio/backend/tests/test_slot_offload_fit.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the offload-avoidance serving-slot reduction (`_slots_that_fit_on_gpu`). + +When a pinned context does not fit at the requested `--parallel` slot count, Studio would +flip to `--fit on` and llama-server offloads layers to host RAM, collapsing decode ~3x +(oobabooga #6718). Instead the loader retries the on-GPU fit at fewer slots and keeps the +largest count that stays fully on GPU (`-ngl -1`). These tests drive the real helper with +synthetic VRAM maps; the KV term is mocked so totals are controlled and the reduction logic +is asserted directly (no GPU, network, or subprocess). +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from core.inference.llama_cpp import LlamaCppBackend + +MIB = 1024 * 1024 +CTX = 90624 +FRAC = LlamaCppBackend._GPU_PIN_VRAM_FRACTION # 0.97; usable = free - 0.03*total + + +def _backend( + vocab = 248320, + embd = 5120, + kv_fixed_mib = 0, +): + """Backend with the dims the compute buffer reads; KV mocked to a fixed size so the + only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15).""" + b = LlamaCppBackend.__new__(LlamaCppBackend) + b._vocab_size = vocab + b._embedding_length = embd + b._key_length_mla = None + b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB + b._can_estimate_kv = lambda: True + return b + + +def _run( + b, + n_parallel, + base_mib, + gpus, + total_by_idx, + overhead_mib = 0, +): + return b._slots_that_fit_on_gpu( + n_parallel, + CTX, + gpus, + total_by_idx, + int(base_mib * MIB), + "q8_0", + FRAC, + int(overhead_mib * MIB), + 1, + 512, + ) + + +class TestSlotsThatFitOnGpu: + """Compute-buffer per slot (vocab 248320, embd 5120): cb(1)=46, cb(2)=604, cb(3)=1162, + cb(4)=1719 MiB. Single 24 GB card usable = 24576 - 0.03*24576 = 23839 MiB.""" + + def test_reduces_to_largest_fitting_slot(self): + # base+KV = 22500: par4 (24219) over 23839, par3 (23662) fits -> 3 slots on GPU. + gi, use_fit, slots = _run(_backend(), 4, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 3 + + def test_floor_when_only_one_slot_fits(self): + # base 23400: par2 (24004) over, par1 (23446) fits -> drop all the way to 1. + gi, use_fit, slots = _run(_backend(), 4, 23400, [(0, 24576)], {0: 24576}) + assert use_fit is False and gi == [0] and slots == 1 + + def test_none_fit_stays_offload(self): + # Even a single slot (24046) exceeds usable -> genuine offload, unchanged. + gi, use_fit, slots = _run(_backend(), 4, 24000, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 4 + + def test_roomy_would_keep_all_but_helper_only_reduces(self): + # On a roomy card par4 fits, so load_model never calls this helper; if called it + # still only searches < n_parallel and never raises the count above the request. + gi, use_fit, slots = _run(_backend(), 4, 5000, [(0, 183000)], {0: 183000}) + assert use_fit is False and slots == 3 and slots < 4 + + def test_single_slot_request_is_noop(self): + # n_parallel == 1: nothing to reduce (range empty) -> report offload unchanged. + gi, use_fit, slots = _run(_backend(), 1, 22500, [(0, 24576)], {0: 24576}) + assert use_fit is True and gi is None and slots == 1 + + def test_multi_gpu_reduces_across_devices(self): + # Needs 2 GPUs: usable/GPU = 23839, cumulative 47677. base+KV 46200: par4 (47919) + # over, par3 (47362) fits across both -> 3 slots spanning [0, 1]. + gi, use_fit, slots = _run( + _backend(), 4, 46200, [(0, 24576), (1, 24576)], {0: 24576, 1: 24576} + ) + assert use_fit is False and gi == [0, 1] and slots == 3 + + def test_kv_counted_per_candidate(self): + # A non-zero (slot-independent) KV shifts the threshold: with 3000 MiB KV and + # base 19500 (= 22500 total at par-independent terms) the same par3 fit holds. + gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576}) + assert use_fit is False and slots == 3 diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 1f09f9a091..0d71b89d87 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -746,10 +746,14 @@ def test_tp_plan_weighted_split_on_asymmetric_big_model(): b, (ec, mac, gi, ts) = _plan(50) reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB assert gi == [0, 1] - # split weighted by (usable - buffer); with no totals usable is free*frac + # split weighted by (usable - flat buffer - per-device context compute); with + # no totals usable is free*frac. The per-device cc is subtracted so the smaller + # card isn't weighted above its real usable budget (see below). + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + assert cc_per_dev > 0 assert ts == [ - int(48000 * _CTX_FIT_VRAM_FRACTION - reserve), - int(24000 * _CTX_FIT_VRAM_FRACTION - reserve), + int(48000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), + int(24000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev), ] assert ec < 131072 # capped below native @@ -819,6 +823,75 @@ def test_tp_plan_mtp_reserves_extra_and_shrinks_context(): assert ec_mtp < ec_no +def test_tp_plan_reserves_context_linear_compute_buffer(): + # Tensor mode replicates the compute graph on every device; measured on + # Qwen3.5-9B at f16 the per-device buffer grows ~n_ubatch*2 B/token (~1024 + # B/tok), so the fit must reserve n_dev x that on top of the flat reserve or + # it over-pins and OOMs at high context. The chosen KV must leave room for it. + b, (ec, mac, gi, ts) = _plan(50) + cc = len(gi) * b._compute_buffer_ctx_bytes(ec, None, "f16") + assert cc > 0 + assert b._estimate_kv_cache_bytes(ec) + cc <= _kv_budget_b(50) + + +def test_tp_plan_context_shrinks_vs_compute_unaware(): + # With the context-linear term the pinned context is strictly below what a + # KV-only (compute-unaware) fit at the same budget would allow. + b, (ec, *_r) = _plan(50) + b2 = _kv_seeded_backend() + b2._embedding_length = 0 # kills the context-linear compute term (returns 0) + ec_naive, *_r2 = b2._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + assert ec < ec_naive + + +def test_tp_plan_soft_overhead_shrinks_context(): + # The CUDA-ctx / mmproj / MTP-draft reserve the layer path folds into the fit + # budget (model_size_fit) must also shrink the tensor context. Tensor mode has + # no --fit valve, so an unreserved overshoot OOMs at startup instead of + # offloading. A non-zero soft_overhead must pin a strictly smaller context. + b = _kv_seeded_backend() + ec_no, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) + ec_soft, *_r2 = b._plan_tensor_parallel( + _ASYM, int(50 * _GB), 131072, soft_overhead_bytes = 2 * _GB + ) + assert 2048 < ec_soft < ec_no + + +def test_tp_plan_soft_overhead_reserved_against_budget(): + # The pinned context must leave the whole soft reserve free on top of KV and + # the replicated context compute, so the real footprint stays within the pool. + b = _kv_seeded_backend() + soft = 2 * _GB + ec, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072, soft_overhead_bytes = soft) + cc = len(_ASYM) * b._compute_buffer_ctx_bytes(ec, None, None) + assert b._estimate_kv_cache_bytes(ec) + cc + soft <= _kv_budget_b(50) + + +def test_tp_plan_weighted_split_keeps_small_gpu_within_budget(): + # Regression: the weighted split must subtract each device's replicated context + # compute (cc_bytes/n_dev), not just the flat reserve. Otherwise the smaller + # card is weighted above its usable budget and OOMs at launch. Model the split: + # llama.cpp distributes weights+KV by the tensor-split weights; every device + # also holds the flat reserve plus its per-device context compute. + b, (ec, mac, gi, ts) = _plan(50) + assert ts is not None and len(ts) == len(gi) == 2 + reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB + cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024) + free_by_idx = {0: 48000, 1: 24000} + split_content_mib = (int(50 * _GB) + b._estimate_kv_cache_bytes(ec)) / (1024 * 1024) + total_weight = sum(ts) + for w, idx in zip(ts, gi): + placed = split_content_mib * w / total_weight + usable = free_by_idx[idx] * _CTX_FIT_VRAM_FRACTION + assert placed + reserve + cc_per_dev <= usable + 1 # +1 MiB for int rounding + + # Lock the regression: under the old formula (flat reserve only) the smaller + # card was placed over its budget; the cc term is what pulls it back. + old_adj = [int(free_by_idx[i] * _CTX_FIT_VRAM_FRACTION - reserve) for i in gi] + old_small_placed = split_content_mib * old_adj[1] / sum(old_adj) + assert old_small_placed + reserve + cc_per_dev > free_by_idx[1] * _CTX_FIT_VRAM_FRACTION + + def test_tp_plan_no_kv_metadata_floors_context(): b = LlamaCppBackend() # no KV metadata -> can't size safely ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072) diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 931d8a705d..39fdd151be 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -71,6 +71,28 @@ class TestFunctionStyleTrailingText: call = _only(text) assert call == {"name": "python", "arguments": {"code": 'print("")'}} + def test_closed_function_with_trailing_prose_heal_path(self): + # Regression: the heal / finalize path (allow_incomplete=True) used to fold + # and the trailing prose into the argument and drop + # the prose from visible content. It must now match the strict path -- keep a + # clean argument and leave the trailing prose outside the call span. + text = "cats trailing words" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + fn = calls[0]["function"] + assert fn["name"] == "web_search" + assert json.loads(fn["arguments"]) == {"query": "cats"} + # The trailing prose sits outside the removed span, so it stays visible. + from core.tool_healing import ( + parse_tool_calls_from_text as _parse_with_spans, + ) + + _calls, spans = _parse_with_spans(text, allow_incomplete = True, with_spans = True) + out = text + for s, e in sorted(spans, reverse = True): + out = out[:s] + out[e:] + assert out == " trailing words" + def test_incomplete_function_without_close_is_still_rejected(self): text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] @@ -160,3 +182,18 @@ class TestHealingPathUnaffected: calls = parse_tool_calls_from_text(text, allow_incomplete = True) assert len(calls) == 1 assert calls[0]["function"]["name"] == "web_search" + + def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): + # allow_incomplete exists for truncated output; a call that DID close + # must parse identically to strict mode, leaving prose after + # out of the last parameter and out of the removal span. + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = "cats trailing" + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert text[span[0] : span[1]] == ( + "cats" + ) diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index a99eb4c45c..2d3dc5fbff 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -12,6 +12,7 @@ from __future__ import annotations import sys from pathlib import Path +from unittest.mock import MagicMock import pytest @@ -71,12 +72,58 @@ def test_default_spec_matches_table(monkeypatch): assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC -def test_skips_torchao_on_windows_rocm(): +@pytest.mark.parametrize( + ("rocm_windows_torch_installed", "installed_torch_is_windows_rocm"), + [ + (True, False), + (False, True), + ], +) +def test_skips_torchao_on_windows_rocm( + monkeypatch, tmp_path, rocm_windows_torch_installed, installed_torch_is_windows_rocm +): """The overrides step must skip torchao on Windows ROCm: no working build exists there (it imports an absent c10d backend and crashes transformers.quantizers), so the installer skips it and relies on the runtime stub instead.""" - source = _INSTALL_SCRIPT.read_text(encoding = "utf-8") - # Branches on the Windows-ROCm marker set by _ensure_rocm_torch ... - assert "elif _rocm_windows_torch_installed:" in source - # ... and reports the skip in the progress label. - assert "dependency overrides (skipped, Windows ROCm)" in source + mod = _load_module(monkeypatch) + installed_specs: list[str] = [] + progress_labels: list[str] = [] + + def _record_pip_install(*args, **kwargs): + installed_specs.extend(str(arg) for arg in args) + return 0 + + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + + subprocess_result = MagicMock() + subprocess_result.returncode = 0 + subprocess_result.stdout = "" + + monkeypatch.setenv("SKIP_STUDIO_BASE", "1") + monkeypatch.setattr(mod, "IS_WINDOWS", True) + monkeypatch.setattr(mod, "IS_MACOS", False) + monkeypatch.setattr(mod, "IS_MAC_ARM", False) + monkeypatch.setattr(mod, "NO_TORCH", False) + monkeypatch.setattr(mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed) + monkeypatch.setattr( + mod, "_installed_torch_is_windows_rocm", lambda: installed_torch_is_windows_rocm + ) + monkeypatch.setattr(mod, "_bootstrap_uv", lambda: False) + monkeypatch.setattr(mod, "_repair_bad_anyio", lambda: None) + monkeypatch.setattr(mod, "_ensure_rocm_torch", lambda: None) + monkeypatch.setattr(mod, "_ensure_cuda_torch", lambda: None) + monkeypatch.setattr(mod, "_has_usable_nvidia_gpu", lambda: True) + monkeypatch.setattr(mod, "run", lambda *args, **kwargs: None) + monkeypatch.setattr(mod, "pip_install", _record_pip_install) + monkeypatch.setattr(mod, "_progress", lambda label: progress_labels.append(label)) + monkeypatch.setattr(mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin) + monkeypatch.setattr(mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin) + monkeypatch.setattr(mod.subprocess, "run", lambda *args, **kwargs: subprocess_result) + + assert mod.install_python_stack() == 0 + + assert not any(spec.startswith("torchao") for spec in installed_specs) + assert "dependency overrides (skipped, Windows ROCm)" in progress_labels diff --git a/studio/backend/utils/embedding_model_settings.py b/studio/backend/utils/embedding_model_settings.py new file mode 100644 index 0000000000..798ae6d364 --- /dev/null +++ b/studio/backend/utils/embedding_model_settings.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persisted RAG embedding-model override (Settings -> General). + +The stored value takes precedence over the ``RAG_EMBEDDING_MODEL`` env default in +``core.rag.config``. Vectors from different models live in different spaces, so +documents already indexed under the old model must be re-uploaded after a change +(the UI warns about this). +""" + +from __future__ import annotations + +import threading +import time +from typing import Any + +EMBEDDING_MODEL_SETTING_KEY = "rag_embedding_model" +MAX_EMBEDDING_MODEL_LENGTH = 512 + +# The effective model is consulted on the embedder hot path (once per embed / +# tokenize call during ingestion), so the stored value is cached briefly instead +# of hitting sqlite each time. Writes invalidate immediately in-process; other +# readers converge within the TTL. +_CACHE_TTL_S = 2.0 +_cached: tuple[float, str | None] | None = None +# Bumped on every write/invalidate. A reader captures it before the DB read and +# only fills the cache if it is unchanged afterward, so a read that overlapped a +# save cannot repopulate the cache with the pre-save value for the whole TTL. +_generation = 0 +_lock = threading.Lock() + + +def _invalidate_cache() -> None: + global _cached, _generation + with _lock: + _cached = None + _generation += 1 + + +def default_embedding_model() -> str: + """The env/default model from rag config (``RAG_EMBEDDING_MODEL`` or bge).""" + from core.rag import config + return config.EMBEDDING_MODEL + + +def _coerce_embedding_model(value: Any) -> str | None: + if not isinstance(value, str): + return None + cleaned = value.strip() + if not cleaned or len(cleaned) > MAX_EMBEDDING_MODEL_LENGTH: + return None + # Newlines/control chars are never valid in a repo id or path. + if any(ord(ch) < 32 for ch in cleaned): + return None + return cleaned + + +def validate_embedding_model(value: Any) -> str: + cleaned = _coerce_embedding_model(value) + if cleaned is None: + raise ValueError( + "Embedding model must be a Hugging Face repo id (e.g. " + "'unsloth/bge-small-en-v1.5') or a local model path, up to " + f"{MAX_EMBEDDING_MODEL_LENGTH} characters." + ) + return cleaned + + +def get_stored_embedding_model() -> str | None: + """The persisted override, or None when unset/invalid.""" + global _cached + now = time.monotonic() + with _lock: + cached = _cached + if cached is not None and now - cached[0] < _CACHE_TTL_S: + return cached[1] + gen = _generation + try: + from storage.studio_db import get_app_setting + stored = get_app_setting(EMBEDDING_MODEL_SETTING_KEY, None) + except Exception: + # Transient store failure: keep the last known value instead of + # silently reverting the embed/search hot path to the default model, + # which would mix vector spaces mid-ingestion. + with _lock: + if _cached is not None: + _cached = (time.monotonic(), _cached[1]) + return _cached[1] + return None + value = _coerce_embedding_model(stored) + with _lock: + # Only cache when no save landed while we were reading; otherwise this + # value may be pre-save, and caching it would mask the new one for the + # TTL. The next reader re-reads the committed value. + if _generation == gen: + _cached = (time.monotonic(), value) + return value + + +def get_rag_embedding_model() -> str: + """Effective embedding model: persisted override, else env/default.""" + return get_stored_embedding_model() or default_embedding_model() + + +def set_rag_embedding_model(value: Any) -> str: + parsed = validate_embedding_model(value) + from storage.studio_db import upsert_app_settings + + # Saving the default is not an override; keeps is_custom (and the UI's + # reset affordance) honest. + stored = parsed if parsed != default_embedding_model() else None + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: stored}) + _invalidate_cache() + return parsed + + +def reset_rag_embedding_model() -> str: + """Clear the override; returns the (env/default) model now in effect.""" + from storage.studio_db import upsert_app_settings + + upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: None}) + _invalidate_cache() + return default_embedding_model() diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 5f2b2abbcf..62b537fbac 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -44,6 +44,12 @@ from .vram_estimation import ( estimate_training_vram, ) + +def export_capability() -> dict: + """Return live export capability from the hardware module.""" + return _hardware.export_capability() + + __all__ = [ "DeviceType", "DEVICE", @@ -51,6 +57,7 @@ __all__ = [ "IS_ROCM", "detect_hardware", "get_device", + "export_capability", "is_apple_silicon", "clear_gpu_cache", "get_gpu_memory_info", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 88a1784b88..8d6c919ebd 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -263,6 +263,49 @@ def get_device() -> DeviceType: return DEVICE +def export_capability() -> dict: + """Whether model export can run here, with a torch-aware reason when it cannot. + + Export runs through Unsloth, which hard-requires an accelerator (it calls ``torch.cuda`` at + import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The + reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch. + + Returns {export_supported, export_unsupported_reason, export_unsupported_message}. + """ + if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX): + return { + "export_supported": True, + "export_unsupported_reason": None, + "export_unsupported_message": None, + } + # No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch" + # would be wrong advice on a Mac even when torch is also absent. + if is_apple_silicon(): + reason = "mlx_unavailable" + message = ( + "Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run " + "`unsloth studio update` to restore MLX and enable export." + ) + elif not _has_torch(): + reason = "pytorch_not_installed" + message = ( + "PyTorch is not installed. Model export requires PyTorch with a supported accelerator " + "(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export." + ) + else: + reason = "no_accelerator" + message = ( + "Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported " + "accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export " + "on CPU only.)" + ) + return { + "export_supported": False, + "export_unsupported_reason": reason, + "export_unsupported_message": message, + } + + def clear_gpu_cache(): """ Clear GPU memory cache for the current device. @@ -710,82 +753,159 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa return None, None +def _gpu_utilization_payload( + device: DeviceType, devices: list[Dict[str, Any]], **metadata: Any +) -> Dict[str, Any]: + """Keep the legacy primary-GPU shape and append all visible devices.""" + backend = _backend_label(device) + normalized = [] + for ordinal, raw in enumerate(devices): + dev = dict(raw) + dev.setdefault("available", True) + dev.setdefault("backend", backend) + if dev.get("visible_ordinal") is None: + dev["visible_ordinal"] = ordinal + normalized.append(dev) + + normalized.sort(key = lambda dev: dev.get("visible_ordinal", dev.get("index", 0))) + payload: Dict[str, Any] = { + "available": bool(normalized), + "backend": backend, + "devices": normalized, + } + payload.update(metadata) + if normalized: + payload.update(normalized[0]) + payload["available"] = True + payload["backend"] = normalized[0].get("backend", backend) + payload["devices"] = normalized + return payload + + def get_gpu_utilization() -> Dict[str, Any]: - """Return a live snapshot of device utilization information.""" + """Live utilization snapshot for the primary GPU plus all visible GPUs.""" device = get_device() + if device == DeviceType.XPU: + result = get_visible_gpu_utilization() + return _gpu_utilization_payload( + device, + result.get("devices", []), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + if device == DeviceType.CUDA: - result = _smi_query("get_primary_gpu_utilization") - if result is not None: - result["backend"] = _backend_label(device) - if IS_ROCM: - # Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.). - _reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec()) - return result - # SMI unavailable. On Windows, use Performance Counters (Task Manager - # source) for system-wide VRAM, covering cross-process usage torch can't see. + parent_visible_spec = _get_parent_visible_gpu_spec() + result = _smi_query( + "get_visible_gpu_utilization", + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result is not None and "devices" in result: + devices = result["devices"] + numeric_ids = parent_visible_spec.get("numeric_ids") + if IS_ROCM and numeric_ids is not None: + _reconcile_rocm_unified_memory(result, numeric_ids) + + return _gpu_utilization_payload( + device, + devices, + backend_cuda_visible_devices = result.get("backend_cuda_visible_devices"), + parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []), + index_kind = result.get("index_kind"), + ) + + # Fallback Windows ROCm if IS_ROCM and platform.system() == "Windows": _win_used, _win_total = _rocm_windows_perf_counter_vram_gb() if _win_used is not None and _win_total is not None: _win_util = _rocm_windows_perf_counter_gpu_util_pct() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _win_util, - "temperature_c": None, - "vram_used_gb": _win_used, - "vram_total_gb": _win_total, - "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) - if _win_total > 0 - else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed. + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _win_util, + "temperature_c": None, + "vram_used_gb": _win_used, + "vram_total_gb": _win_total, + "vram_utilization_pct": round((_win_used / _win_total) * 100, 1) + if _win_total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) + + # Fallback Linux ROCm if IS_ROCM and platform.system() == "Linux": _linux_used, _linux_total = _rocm_linux_sysfs_vram_gb() if _linux_used is not None and _linux_total is not None: _linux_util = _rocm_linux_sysfs_gpu_busy_pct() _linux_temp = _rocm_linux_sysfs_temp_c() _linux_power = _rocm_linux_sysfs_power_w() - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": _linux_util, - "temperature_c": _linux_temp, - "vram_used_gb": _linux_used, - "vram_total_gb": _linux_total, - "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) - if _linux_total > 0 - else None, - "power_draw_w": _linux_power, - "power_limit_w": None, - "power_utilization_pct": None, - } - # Last resort: torch mem_get_info (process-local). - _visible_spec = _get_parent_visible_gpu_spec() - _numeric_ids = _visible_spec.get("numeric_ids") or [0] - _primary_idx = [_numeric_ids[0]] if _numeric_ids else [0] - _torch_devices = _torch_get_per_device_info(_primary_idx) - if _torch_devices: - _td = _torch_devices[0] - _total = _td["total_gb"] - _used = _td["used_gb"] - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": _used, - "vram_total_gb": _total, - "vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None, - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": _linux_util, + "temperature_c": _linux_temp, + "vram_used_gb": _linux_used, + "vram_total_gb": _linux_total, + "vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1) + if _linux_total > 0 + else None, + "power_draw_w": _linux_power, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - # MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%. + # Last resort: torch mem_get_info (process-local) for all visible GPUs + _visible_spec = _get_parent_visible_gpu_spec() + _numeric_ids = _visible_spec.get("numeric_ids") or [] + if not _numeric_ids: + visible_count = _torch_get_physical_gpu_count() or 0 + _numeric_ids = list(range(visible_count)) + + _torch_devices = _torch_get_per_device_info(_numeric_ids) + if _torch_devices: + gpu_array = [] + for _td in _torch_devices: + _total = _td["total_gb"] + _used = _td["used_gb"] + gpu_array.append( + { + "available": True, + "backend": _backend_label(device), + "index": _td["index"], + "name": _td.get("name", "Unknown"), + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": _used, + "vram_total_gb": _total, + "vram_utilization_pct": round((_used / _total) * 100, 1) + if _total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return _gpu_utilization_payload(device, gpu_array) + + # MLX if device == DeviceType.MLX: try: import psutil @@ -793,9 +913,8 @@ def get_gpu_utilization() -> Dict[str, Any]: total_bytes = psutil.virtual_memory().total except Exception as e: logger.error(f"Error getting MLX GPU utilization: {e}") - return {"available": False, "backend": device.value, "error": str(e)} - if not agx: - return {"available": False, "backend": device.value} + return {"available": False, "backend": device.value, "devices": [], "error": str(e)} + allocated_bytes = agx.get("vram_used_bytes", 0) or 0 vram_used_gb = allocated_bytes / (1024**3) total_gb = total_bytes / (1024**3) @@ -814,37 +933,51 @@ def get_gpu_utilization() -> Dict[str, Any]: from . import apple - return { - "available": True, - "backend": device.value, - "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, - "temperature_c": apple.read_gpu_temperature_c(), - "vram_used_gb": round(vram_used_gb, 2), - "vram_total_gb": round(total_gb, 2), - "vram_utilization_pct": ( - round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None - ), - "power_draw_w": apple.read_gpu_power_w(), - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": device.value, + "index": 0, + "visible_ordinal": 0, + "gpu_utilization_pct": agx.get("utilization_pct") if agx else None, + "temperature_c": apple.read_gpu_temperature_c(), + "vram_used_gb": round(vram_used_gb, 2), + "vram_total_gb": round(total_gb, 2), + "vram_utilization_pct": round((vram_used_gb / total_gb) * 100, 1) + if total_gb > 0 + else None, + "power_draw_w": apple.read_gpu_power_w(), + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) mem = get_gpu_memory_info() if device != DeviceType.CPU and mem.get("available"): - return { - "available": True, - "backend": _backend_label(device), - "gpu_utilization_pct": None, - "temperature_c": None, - "vram_used_gb": round(mem.get("allocated_gb", 0), 2), - "vram_total_gb": round(mem.get("total_gb", 0), 2), - "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), - "power_draw_w": None, - "power_limit_w": None, - "power_utilization_pct": None, - } + return _gpu_utilization_payload( + device, + [ + { + "available": True, + "backend": _backend_label(device), + "index": mem.get("device", 0), + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + ) - return {"available": False, "backend": _backend_label(device)} + return {"available": False, "backend": _backend_label(device), "devices": []} def _apply_unified_memory_correction( diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index 8648b053d5..c16ae91467 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -324,12 +324,74 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]: } +def _is_external_link(path: Optional[Path]) -> bool: + """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink + or a Windows directory junction / reparse point. Such a link resolves into + the user's own llama.cpp checkout, so Studio must never auto-update it.""" + if path is None: + return False + try: + if os.path.islink(path): + return True + except OSError: + return False + if os.name == "nt": + try: + import stat + attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined] + return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT) + except (OSError, AttributeError): + return False + return False + + +def _active_install_is_local_link(binary: Optional[str]) -> bool: + """True when the active llama-server resolves through a --with-llama-cpp-dir + local link at the canonical llama.cpp directory. An update would write + through that link into the user's own checkout (or fail), so the install is + treated as externally managed: no update is offered or applied. Checks only + up to and including the ``llama.cpp`` dir so a symlinked HOME / studio root + above it can't trip a false positive.""" + if not binary: + return False + for parent in Path(binary).parents: + if _is_external_link(parent): + return True + if parent.name == "llama.cpp": + break + return False + + +def _local_link_status() -> dict: + """Status payload for a local-link install: unmanaged, no update offered.""" + with _job_lock: + job = dict(_job) + return { + "supported": False, + "update_available": False, + "stale": False, + "installed_tag": None, + "latest_tag": None, + "published_repo": None, + "installed_at_utc": None, + "age_days": None, + "source_build": False, + "local_link": True, + "update_size_bytes": None, + "job": job, + } + + def get_update_status(*, force_refresh: bool = False) -> dict: """Report whether a newer prebuilt exists plus the current job state. force_refresh bypasses the 24h release cache for an explicit "check now". """ binary = _find_binary() + # A --with-llama-cpp-dir local link is the user's own tree; never offer to + # replace it. Bail before any network/freshness work. + if _active_install_is_local_link(binary): + return _local_link_status() marker = read_install_marker(binary) with _job_lock: @@ -537,6 +599,19 @@ def start_update() -> dict: """Kick off a background update. Idempotent: a second call while one is running returns the in-flight job rather than starting another.""" binary = _find_binary() + # Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt + # here would write through the link into the user's own checkout (or fail) + # and silently drop the link the flag created. + if _active_install_is_local_link(binary): + return { + "started": False, + "reason": "local_link", + "message": ( + "llama.cpp is a local directory linked with --with-llama-cpp-dir; " + "Studio won't replace it. Update your own llama.cpp checkout instead." + ), + "job": get_update_status()["job"], + } marker = read_install_marker(binary) script = _installer_script() if script is None: diff --git a/studio/backend/utils/paths/external_media.py b/studio/backend/utils/paths/external_media.py new file mode 100644 index 0000000000..1f1754664f --- /dev/null +++ b/studio/backend/utils/paths/external_media.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""External media path helpers.""" + +from __future__ import annotations + +import getpass +import os +import platform +from pathlib import Path + +from utils.paths.sensitive import ( + contains_sensitive_path_component, + is_sensitive_path_component, +) + + +def _is_linux_media_mount_path(path: str, media_root: Path | str) -> bool: + normalized = os.path.normpath(os.path.realpath(os.path.expanduser(path))) + root = os.path.normpath(os.path.realpath(os.path.expanduser(str(media_root)))) + try: + rel = os.path.relpath(normalized, root) + except ValueError: + return False + if rel == "." or rel == ".." or rel.startswith(f"..{os.sep}"): + return False + parts = [part for part in rel.split(os.sep) if part] + return len(parts) >= 2 and all(part not in (".", "..") for part in parts[:2]) + + +def is_linux_run_media_path(path: str) -> bool: + """True for Linux removable-media paths under /run/media//.""" + if platform.system() != "Linux": + return False + return _is_linux_media_mount_path(path, "/run/media") + + +def _current_username() -> str | None: + try: + user = getpass.getuser().strip() + except Exception: + return None + return user or None + + +def _contains_sensitive_media_component(path: Path, media_root: Path) -> bool: + try: + rel = path.relative_to(media_root) + except ValueError: + rel = path + return contains_sensitive_path_component(str(rel)) + + +def linux_run_media_mount_roots( + base: Path | str = "/run/media", *, user: str | None = None +) -> list[Path]: + """Readable /run/media// roots for the folder browser.""" + if platform.system() != "Linux": + return [] + user = user or _current_username() + if not user or user in (".", "..") or os.sep in user: + return [] + base_path = Path(base) + try: + resolved_base = base_path.resolve() + except (OSError, RuntimeError, ValueError): + return [] + + roots: list[Path] = [] + seen: set[str] = set() + user_dir = base_path / user + try: + if not user_dir.is_dir(): + return [] + volume_dirs = list(user_dir.iterdir()) + except (OSError, RuntimeError, ValueError): + return [] + for volume_dir in volume_dirs: + if is_sensitive_path_component(volume_dir.name): + continue + try: + resolved = volume_dir.resolve() + except (OSError, RuntimeError, ValueError): + continue + if not _is_linux_media_mount_path(str(resolved), resolved_base): + continue + if _contains_sensitive_media_component(resolved, resolved_base): + continue + key = os.path.normcase(os.path.realpath(str(resolved))) + if key in seen: + continue + try: + is_dir = resolved.is_dir() + except OSError: + continue + if is_dir and os.access(resolved, os.R_OK | os.X_OK): + seen.add(key) + roots.append(resolved) + return roots diff --git a/studio/backend/utils/paths/sensitive.py b/studio/backend/utils/paths/sensitive.py new file mode 100644 index 0000000000..7d32a5f4cf --- /dev/null +++ b/studio/backend/utils/paths/sensitive.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Shared sensitive path-component policy.""" + +from __future__ import annotations + +import os + + +SENSITIVE_PATH_COMPONENTS = { + ".aws", + ".azure", + ".config", + ".docker", + ".gcloud", + ".gnupg", + ".huggingface", + ".kaggle", + ".kube", + ".modelscope", + ".ngc", + ".local", + ".mozilla", + ".pki", + ".thunderbird", + ".ssh", + ".1password", + ".bitwarden", + ".password-store", + "1password", + "bitwarden", + "keychains", + "keyrings", + "mozilla", + "thunderbird", +} + + +def is_sensitive_path_component(name: str) -> bool: + return name.lower() in SENSITIVE_PATH_COMPONENTS + + +def contains_sensitive_path_component(path: str) -> bool: + parts = os.path.normpath(path).split(os.sep) + return any(is_sensitive_path_component(part) for part in parts) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index ac0ecfbfcd..2a63caa5b2 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -226,6 +226,11 @@ _VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510") # Backwards-compat alias _VENV_T5_DIR = _VENV_T5_550_DIR +# llm-compressor-main shadow for FP8/FP4 export of newer-transformers models. Like the .venv_t5_* +# sidecars but also shadows llm-compressor main + compressed-tensors; installed --no-deps so it +# reuses the workspace torch (torch-agnostic). +_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor") + # Tier precedence: higher rank wins in _higher_tier. _TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3} @@ -1518,6 +1523,152 @@ def _ensure_venv_t5_exists() -> bool: return _ensure_venv_t5_550_exists() +# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) --------------------- +# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize +# Qwen3.5 / Gemma-4 / Llama. +_LLMC_MAIN_TRANSFORMERS = "5.10.2" +_LLMC_MAIN_SHA = "973c9c539a84dd9efaf74e115ede5ca419704c18" +_LLMC_MAIN_COMPRESSED_TENSORS = "0.17.2a20260702" +# Installed --no-deps (torch untouched); the full runtime set llm-compressor main needs, pinned. +_VENV_LLMCOMPRESSOR_SPECS = ( + f"transformers=={_LLMC_MAIN_TRANSFORMERS}", + f"llmcompressor @ git+https://github.com/vllm-project/llm-compressor@{_LLMC_MAIN_SHA}", + f"compressed-tensors=={_LLMC_MAIN_COMPRESSED_TENSORS}", + "huggingface-hub==1.21.0", + "hf-xet==1.5.1", + "tokenizers==0.22.2", + "safetensors==0.8.0", + "accelerate==1.14.0", + "datasets==5.0.0", + "pydantic==2.13.4", + "pydantic-core==2.46.4", + "typing-inspection==0.4.2", + "loguru==0.7.3", + "pyyaml==6.0.3", + "nvidia-ml-py==13.610.43", + "pillow==12.3.0", + "auto-round==0.13.1", + "regex==2026.6.28", +) +# Fingerprint of the pin set; bump the trailing schema version to force a rebuild on layout changes. +_LLMC_SHADOW_FINGERPRINT = ( + f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1" +) +_LLMC_SHADOW_MARKER = ".unsloth_llmc_fingerprint" + + +def _llmcompressor_main_disabled() -> bool: + """True if the operator forbids the llm-compressor-main shadow (air-gapped / locked-down).""" + return os.environ.get("UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def _llmcompressor_shadow_is_valid() -> bool: + """True if the shadow dir exists with a marker matching the current pin fingerprint.""" + marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER + try: + return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT + except Exception: + return False + + +def _ensure_venv_llmcompressor_exists() -> bool: + """Ensure .venv_llmcompressor/ has the pinned llm-compressor-main stack. Install if missing. + + All specs are installed with --no-deps into a --target dir (mirrors the transformers sidecars), + so the workspace torch is never touched. Returns True on success. + """ + if _llmcompressor_shadow_is_valid(): + return True + if _llmcompressor_main_disabled(): + logger.warning( + "llm-compressor-main shadow needed but UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN is set; " + "compressed export of newer-transformers models will fail fast." + ) + return False + if _env_offline(): + logger.warning( + "llm-compressor-main shadow missing and HF/offline mode is set; cannot provision it." + ) + return False + + logger.warning( + "Provisioning llm-compressor-main shadow at %s (one-time, ~a few hundred MB, no torch) ...", + _VENV_LLMCOMPRESSOR_DIR, + ) + shutil.rmtree(_VENV_LLMCOMPRESSOR_DIR, ignore_errors = True) + os.makedirs(_VENV_LLMCOMPRESSOR_DIR, exist_ok = True) + + # Prefer uv (faster) then pip; install every spec at once, --no-deps, prereleases allowed + # (compressed-tensors ships as a pre-release). + base = [ + "--target", + _VENV_LLMCOMPRESSOR_DIR, + "--no-deps", + "--prerelease=allow", + *_VENV_LLMCOMPRESSOR_SPECS, + ] + cmds = [] + if shutil.which("uv"): + cmds.append(["uv", "pip", "install", "--python", sys.executable, *base]) + cmds.append( + [ + sys.executable, + "-m", + "pip", + "install", + *[a for a in base if a != "--prerelease=allow"], + "--pre", + ] + ) + + last_out = "" + for cmd in cmds: + result = subprocess.run( + cmd, + stdout = subprocess.PIPE, + stderr = subprocess.STDOUT, + text = True, + env = child_env_without_native_path_secret(), + **_windows_hidden_subprocess_kwargs(), + ) + last_out = result.stdout or "" + if result.returncode == 0: + try: + (Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER).write_text( + _LLMC_SHADOW_FINGERPRINT + ) + except Exception: + pass + logger.info("Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR) + return True + logger.warning("llm-compressor-main shadow install failed with %s; trying next", cmd[0]) + + logger.error( + "Failed to provision llm-compressor-main shadow (spec: llmcompressor@%s). Output:\n%s", + _LLMC_MAIN_SHA, + last_out[-4000:], + ) + return False + + +def llmcompressor_shadow_pythonpath() -> str | None: + """Provision (lazily) the llm-compressor-main shadow and return its sys.path entry, or None. + + Returns None when the shadow is disabled (UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN), offline, or + provisioning failed - callers then fall back to the fail-fast path. + """ + if _llmcompressor_main_disabled(): + return None + if _ensure_venv_llmcompressor_exists(): + return _VENV_LLMCOMPRESSOR_DIR + return None + + def _activate_venv(venv_dir: str, label: str) -> None: """Prepend *venv_dir* to sys.path, purge stale modules, reimport.""" if venv_dir not in sys.path: diff --git a/studio/frontend/src/app/router.tsx b/studio/frontend/src/app/router.tsx index 9272dfab2b..10f4947aac 100644 --- a/studio/frontend/src/app/router.tsx +++ b/studio/frontend/src/app/router.tsx @@ -11,7 +11,6 @@ import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId"; import { Route as chatRoute } from "./routes/chat"; import { Route as exportRoute } from "./routes/export"; import { Route as imagesRoute } from "./routes/images"; -import { Route as gridTestRoute } from "./routes/grid-test"; import { Route as indexRoute } from "./routes/index"; import { Route as loginRoute } from "./routes/login"; import { Route as hubRoute } from "./routes/hub"; @@ -26,7 +25,6 @@ const routeTree = rootRoute.addChildren([ onboardingRoute, loginRoute, changePasswordRoute, - gridTestRoute, hubRoute, settingsRoute, studioRoute, diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 220f66d277..b70cf9d796 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -78,6 +78,9 @@ const CHAT_ONLY_ALLOWED = new Set([ "/login", "/signup", "/change-password", + // Export stays reachable on chat-only hosts so the page can show its own grayed-out reason + // instead of a silent redirect; it self-gates via export capability, so nothing runs. + "/export", ]); function isChatOnlyAllowed(pathname: string): boolean { diff --git a/studio/frontend/src/app/routes/grid-test.tsx b/studio/frontend/src/app/routes/grid-test.tsx deleted file mode 100644 index c4b6b505a1..0000000000 --- a/studio/frontend/src/app/routes/grid-test.tsx +++ /dev/null @@ -1,69 +0,0 @@ -// 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 { DashboardGrid, DashboardLayout } from "@/components/layout"; -import { - Card, - CardContent, - CardDescription, - CardHeader, - CardTitle, -} from "@/components/ui/card"; -import { createRoute } from "@tanstack/react-router"; -import { requireAuth } from "../auth-guards"; -import { Route as rootRoute } from "./__root"; - -export const Route = createRoute({ - getParentRoute: () => rootRoute, - path: "/grid-test", - beforeLoad: () => requireAuth(), - component: GridTestPage, -}); - -function GridTestPage() { - return ( - -
-
-

Grid Test - 3 Columns

-

- max-w-7xl, gap-6, responsive 1→2→3 -

-
- - - {[1, 2, 3].map((i) => ( - - - Card {i} - ~400px at 1280px viewport - - -
- - - ))} - - -
-

4 Columns

-

~296px per card at 1280px

-
- - - {[1, 2, 3, 4].map((i) => ( - - - Card {i} - Smaller cards - - -
- - - ))} - -
- - ); -} diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 73e2b6c5d8..b7adb40d04 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -291,13 +291,12 @@ export function AppSidebar() { const chatOnly = usePlatformStore((s) => s.isChatOnly()); const chatOnlyReason = usePlatformStore((s) => s.chatOnlyReason); - // When Train/Export are greyed out (chat-only host), explain why on hover - // instead of disabling them silently. mlx_unavailable is the common macOS case - // after a reinstall/update dropped MLX and is recoverable via `unsloth studio update`. - const trainExportDisabledHint: string | undefined = !chatOnly + // Explain a greyed-out Train (chat-only host) on hover instead of disabling silently. Export is + // no longer disabled here: it stays navigable so its page can show a precise grayed-out reason. + const trainDisabledHint: string | undefined = !chatOnly ? undefined : chatOnlyReason === "mlx_unavailable" - ? "Training needs MLX. Run `unsloth studio update` to enable Train and Export." + ? "Training needs MLX. Run `unsloth studio update` to enable Train." : chatOnlyReason === "intel_mac" ? "Training needs Apple Silicon or a GPU. Intel Macs are chat-only." : chatOnlyReason === "no_gpu" @@ -1216,7 +1215,7 @@ export function AppSidebar() { pathname === "/studio" || pathname.startsWith("/studio/") } disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1245,7 +1244,7 @@ export function AppSidebar() { label={t("shell.navigation.train")} active={pathname === "/studio" || pathname.startsWith("/studio/")} disabled={chatOnly} - tooltip={trainExportDisabledHint} + tooltip={trainDisabledHint} spinner={trainingInProgress} onClick={() => { if (chatOnly) return; @@ -1266,11 +1265,8 @@ export function AppSidebar() { icon={DownloadSquare01Icon} label={t("shell.navigation.export")} active={pathname === "/export" || pathname.startsWith("/export/")} - disabled={chatOnly} - tooltip={trainExportDisabledHint} spinner={exportInProgress} onClick={() => { - if (chatOnly) return; navigate({ to: "/export" }); closeMobileIfOpen(); }} diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index b8d75b2eb6..336dfece69 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1,6 +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 +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -103,6 +104,7 @@ import { type FormatFilter, estimateQuantBytes, fitsDevice, + hfModelFitsDevice, isMlxId, isMobileVariant, isRecommendableFormat, @@ -1414,6 +1416,9 @@ export function HubModelPicker({ }, []); // When on, On Device GGUF repos show their quantizations without a click. const expandQuantizations = useChatRuntimeStore((s) => s.expandQuantizations); + // Shared with the Hub page: list only models sized within the device budget. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); // Repos the user clicked to collapse while expand-by-default is on. Kept in // memory only, so it resets on reload (and when the setting is toggled). const [collapsedGguf, setCollapsedGguf] = useState>( @@ -1805,34 +1810,19 @@ export function HubModelPicker({ formatFilter === "all" ? rows.filter((r) => isRecommendableFormat(r.id, r.isGguf, isMac)) : rows.filter((r) => matchesFormatFilter(r.id, r.isGguf, formatFilter)); - if (recommendedSort !== "recommended") return rows; + // The "recommended" sort always applies the device-fit filter; the shared + // "Fits on device" tick extends it to the other sorts too. + if (recommendedSort !== "recommended" && !fitOnDeviceOnly) return rows; return rows.filter((r) => { // Downloaded models always show, regardless of device fit. if (downloadedSet.has(r.id.toLowerCase())) return true; - // Unified-memory hosts (Mac / no discrete GPU) still report system RAM, - // so fall back to that budget instead of skipping the fit check entirely. - const hasDeviceBudget = - gpu.memoryTotalGb > 0 || gpu.systemRamAvailableGb > 0; - if (!hasDeviceBudget) return true; - // GGUF/MLX repos rarely expose safetensors metadata, so fall back to the - // GGUF param count, then the repo name, for a size estimate. Anything we - // still cannot size is hidden (requireKnown) so over-budget models like a - // 1T GGUF don't slip into Recommended. - const params = r.totalParams ?? paramsFromId(r.id); - const sizeBytes = - r.estimatedSizeBytes ?? - (params ? estimateQuantBytes(params) : undefined); - return fitsDevice({ - sizeBytes, - gpuGb: gpu.memoryTotalGb, - systemRamGb: gpu.systemRamAvailableGb, - requireKnown: true, - }); + return hfModelFitsDevice(r, gpu); }); }, [ recommendedSearch.results, downloadedSet, recommendedSort, + fitOnDeviceOnly, formatFilter, isMac, gpu, @@ -2111,23 +2101,6 @@ export function HubModelPicker({ [visibleCachedModelRows], ); - // Recommended models that match the current search query - const filteredRecommendedIds = useMemo(() => { - if (!showHfSection) return []; - const q = normalizeForSearch(debouncedQuery.trim()); - return recommendedIds - .filter((id) => normalizeForSearch(id).includes(q)) - .filter((id) => - matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), - ); - }, [ - showHfSection, - debouncedQuery, - recommendedIds, - formatFilter, - isKnownGgufRepo, - ]); - // Param counts come straight off the unsloth listings the picker already // loaded, so no extra per-id fetch is needed for the VRAM badges. const recommendedParamCountById = useMemo(() => { @@ -2138,6 +2111,42 @@ export function HubModelPicker({ return map; }, [results, recommendedSearch.results]); + // Recommended models that match the current search query + const filteredRecommendedIds = useMemo(() => { + if (!showHfSection) return []; + const q = normalizeForSearch(debouncedQuery.trim()); + return recommendedIds + .filter((id) => normalizeForSearch(id).includes(q)) + .filter((id) => + matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter), + ) + // Curated defaults obey the fit toggle like the live HF rows, else large + // defaults resurface in search results with the filter on. + .filter( + (id) => + !fitOnDeviceOnly || + downloadedSet.has(id.toLowerCase()) || + hfModelFitsDevice( + { + id, + totalParams: recommendedParamCountById.get(id), + isGguf: isKnownGgufRepo(id), + }, + gpu, + ), + ); + }, [ + showHfSection, + debouncedQuery, + recommendedIds, + formatFilter, + isKnownGgufRepo, + fitOnDeviceOnly, + downloadedSet, + recommendedParamCountById, + gpu, + ]); + const recommendedSet = useMemo( () => new Set(filteredRecommendedIds), [filteredRecommendedIds], @@ -2148,6 +2157,12 @@ export function HubModelPicker({ if (!showHfSection || section !== "recommended") return []; return results .filter(isChatSupported) + .filter( + (r) => + !fitOnDeviceOnly || + downloadedSet.has(r.id.toLowerCase()) || + hfModelFitsDevice(r, gpu), + ) .map((result) => result.id) .filter((id) => !isHiddenModelId(id)) .filter((id) => id.toLowerCase().startsWith("unsloth/")) @@ -2174,6 +2189,9 @@ export function HubModelPicker({ isKnownGgufRepo, isChatSupported, formatFilter, + fitOnDeviceOnly, + downloadedSet, + gpu, isMac, task, ]); @@ -2463,6 +2481,35 @@ export function HubModelPicker({ // selected-item checkmark never overlaps the label. const sortMenuContentClassName = "!p-1 !rounded-[14px] [&_[role=option]]:!pl-2 [&_[role=option]]:!py-1.5 [&_[role=option]]:!text-xs [&_[role=option]]:!rounded-[10px]"; + // Device-fit toggle lives inside the sort menu (shared with the Hub page). + // The whole row is the click target (a button): a Checkbox renders as a + // + + + Hides models larger than this device's memory budget. Downloaded models + stay visible. + + + ); const sectionSortDropdown = section === "recommended" ? ( ) : section === "downloaded" ? ( ) : ( ); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts index 24f0edc784..7c2ed266c0 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts @@ -114,3 +114,35 @@ export function fitsDevice(opts: { } return requireKnown ? false : true; } + +/** Fit predicate for one Hub listing row, shared by the chat model selector + * and the Hub page "Fits on device" filter. GGUF repos: metadata size (actual + * weights) or the smallest-quant estimate from the param count. Safetensors / + * MLX repos: always the params-based smallest-quant estimate, matching the + * VRAM badge's quantized-load assumption; their estimatedSizeBytes is the + * full-precision checkpoint and would wrongly hide models the quantized load + * path can run. Anything unsizable is hidden (requireKnown) so over-budget + * models with no metadata don't slip through. An unknown device budget keeps + * everything. */ +export function hfModelFitsDevice( + model: { + id: string; + totalParams?: number; + estimatedSizeBytes?: number; + isGguf?: boolean; + }, + gpu: { memoryTotalGb: number; systemRamAvailableGb: number }, +): boolean { + if (gpu.memoryTotalGb <= 0 && gpu.systemRamAvailableGb <= 0) return true; + const params = model.totalParams ?? paramsFromId(model.id); + const quantBytes = params ? estimateQuantBytes(params) : undefined; + const sizeBytes = isGgufId(model.id, model.isGguf) + ? (model.estimatedSizeBytes ?? quantBytes) + : (quantBytes ?? model.estimatedSizeBytes); + return fitsDevice({ + sizeBytes, + gpuGb: gpu.memoryTotalGb, + systemRamGb: gpu.systemRamAvailableGb, + requireKnown: true, + }); +} diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx new file mode 100644 index 0000000000..f02da6612e --- /dev/null +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -0,0 +1,160 @@ +// 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 { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useSystemInfo } from "@/hooks/use-system"; +import { useT } from "@/i18n"; +import { cn } from "@/lib/utils"; +import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react"; +import { motion } from "motion/react"; +import { useRef } from "react"; + +function clampPercent(value: number): number { + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number): string { + const digits = value >= 10 ? 1 : 2; + return `${value.toFixed(digits)} GB`; +} + +export function FloatingMonitor() { + const t = useT(); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); + + const constraintsRef = useRef(null); + + if (!isOpen) return null; + + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const ramPercent = clampPercent(systemInfo.memory?.percent_used ?? 0); + + const devices = systemInfo.gpu?.devices ?? []; + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramPercent = clampPercent( + vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0, + ); + + const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; + + return ( +
+ +
+
+ + + {t("settings.resources.liveMonitor.title")} + +
+
+
+ +
+ + +
+
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGb(ramUsed)} / {formatGb(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGb(vramUsed)} / {formatGb(vramTotal)} +
+ +
+ )} + + +
+ ); +} diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index ca4bb7afde..7c9685d6d0 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -42,6 +42,8 @@ export const CHAT_EXPAND_QUANTIZATIONS_KEY = "unsloth_chat_expand_quantizations"; export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY = "unsloth_chat_show_all_quantizations"; +export const MODELS_FIT_ON_DEVICE_ONLY_KEY = + "unsloth_models_fit_on_device_only"; export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; @@ -671,6 +673,9 @@ type ChatRuntimeStore = { expandQuantizations: boolean; /** Persisted: show non-downloaded quantizations too, not just downloaded. */ showAllQuantizations: boolean; + /** Persisted, shared by the chat model selector and the Hub page: list only + * models whose size fits this device's memory budget. */ + fitOnDeviceOnly: boolean; /** A local model picked while `loadOnSelection` is off: staged, not loaded. * The settings sheet shows its load knobs and a Load button. */ pendingSelection: PendingModelSelection | null; @@ -793,6 +798,7 @@ type ChatRuntimeStore = { setLoadOnSelection: (value: boolean) => void; setExpandQuantizations: (value: boolean) => void; setShowAllQuantizations: (value: boolean) => void; + setFitOnDeviceOnly: (value: boolean) => void; setPendingSelection: (selection: PendingModelSelection | null) => void; /** Stage a pick for a deferred load: revert knobs to the loaded baseline, * record the selection, and open the settings sheet. */ @@ -1111,6 +1117,7 @@ export const useChatRuntimeStore = create((set, get) => ({ loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false), showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true), + fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false), pendingSelection: null, loadedIsMultimodal: false, loadedIsDiffusion: false, @@ -1582,6 +1589,10 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, showAllQuantizations); set({ showAllQuantizations }); }, + setFitOnDeviceOnly: (fitOnDeviceOnly) => { + saveBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, fitOnDeviceOnly); + set({ fitOnDeviceOnly }); + }, setPendingSelection: (pendingSelection) => set({ pendingSelection }), stageModel: (selection) => { // Refuse staging mid-load: post-load cleanup would silently drop the queued diff --git a/studio/frontend/src/features/export/api/export-api.ts b/studio/frontend/src/features/export/api/export-api.ts index d1b4e88a6d..be9767a24f 100644 --- a/studio/frontend/src/features/export/api/export-api.ts +++ b/studio/frontend/src/features/export/api/export-api.ts @@ -127,6 +127,8 @@ export async function loadCheckpoint(params: { export async function exportMerged(params: { save_directory: string; format_type?: string; + /** Compressed-tensors scheme alias (e.g. "fp8", "w4a16", "mxfp4"); overrides format_type. */ + compressed_method?: string | null; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -158,7 +160,8 @@ export async function exportBase(params: { export async function exportGGUF(params: { save_directory: string; - quantization_method: string; + /** A single GGUF quant method or a list (list produces multiple GGUFs from one model load). */ + quantization_method: string | string[]; push_to_hub?: boolean; repo_id?: string | null; hf_token?: string | null; @@ -179,6 +182,10 @@ export async function exportLoRA(params: { repo_id?: string | null; hf_token?: string | null; private?: boolean; + /** Also convert the adapter to a GGUF LoRA file (llama.cpp `--lora`). */ + gguf?: boolean; + /** GGUF LoRA output float type (f32/f16/bf16/q8_0/auto); only used when gguf=true. */ + gguf_outtype?: string; }): Promise { const response = await authFetch("/api/export/export/lora", { method: "POST", diff --git a/studio/frontend/src/features/export/components/export-run-panel.tsx b/studio/frontend/src/features/export/components/export-run-panel.tsx index 9e0edf7303..29ba5a703b 100644 --- a/studio/frontend/src/features/export/components/export-run-panel.tsx +++ b/studio/frontend/src/features/export/components/export-run-panel.tsx @@ -28,7 +28,11 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useEffect, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; -import { EXPORT_METHODS, type ExportMethod } from "../constants"; +import { + EXPORT_METHODS, + type ExportMethod, + findMergedFormat, +} from "../constants"; import type { ExportLogEntry } from "../api/export-api"; import { getExportLogLineClass } from "../lib/log-style"; import { @@ -200,6 +204,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) { const summaryMethodLabel = summary?.methodLabel ?? methodTitle; const summaryQuants = summary?.quantLevels ?? quantLevels; const summaryMethod = summary?.method ?? exportMethod; + const summaryFormats = (summary?.mergedFormats ?? []).map( + (v) => findMergedFormat(v)?.label ?? v, + ); const showProgress = isExporting || isTerminal; return ( @@ -392,14 +399,32 @@ export function ExportRunPanel(props: ExportRunPanelProps) { ? "Export finished and pushed to Hugging Face Hub." : "Export finished successfully."} - {run.result?.outputPath ? ( - - {run.result.outputPath} - - ) : null} + {(() => { + // List every folder written; a multi-format merged run created one per format. + const paths = run.result?.outputPaths ?? []; + const items = + paths.length > 0 + ? paths + : run.result?.outputPath + ? [{ label: "", path: run.result.outputPath }] + : []; + const showLabels = items.length > 1; + return items.map((o, i) => ( +
+ {showLabels && o.label ? ( + + {o.label} + + ) : null} + + {o.path} + +
+ )); + })()}
)} @@ -432,6 +457,14 @@ export function ExportRunPanel(props: ExportRunPanelProps) { Export Method {summaryMethodLabel}
+ {summaryMethod === "merged" && summaryFormats.length > 0 && ( +
+ Formats + + {summaryFormats.join(", ")} + +
+ )} {summaryMethod === "gguf" && summaryQuants.length > 0 && (
Quantizations diff --git a/studio/frontend/src/features/export/constants.ts b/studio/frontend/src/features/export/constants.ts index 058de1edc4..be9080ba14 100644 --- a/studio/frontend/src/features/export/constants.ts +++ b/studio/frontend/src/features/export/constants.ts @@ -55,34 +55,172 @@ export const QUANT_OPTIONS: { { value: "f16", label: "F16" }, ]; -/** Merged-export precision formats. The compressed-tensors ones run llm-compressor for vLLM. */ -export type MergedFormat = - | "16-bit (FP16)" - | "FP8 (compressed-tensors)" - | "NVFP4 (compressed-tensors)"; +/** + * Merged-export precision formats, sorted by bit width. Three backends: + * - "plain": standard save (16-bit); `formatType` is the backend `format_type`. + * - "compressed": llm-compressor compressed-tensors (vLLM), NVIDIA-only; `value` is the alias. + * - "torchao": portable FP8/INT8, no NVIDIA GPU needed; `value` is the alias. + * `common` entries are quick pills, the rest the "More formats" dropdown; `needsNvidia` entries + * are hidden on non-NVIDIA hardware. + */ +export type MergedBackend = "plain" | "compressed" | "torchao"; -export const MERGED_FORMATS: { - value: MergedFormat; +export type MergedFormatOption = { + value: string; label: string; + bits: number; + backend: MergedBackend; + group: string; + common: boolean; + needsNvidia: boolean; + needsCalibration?: boolean; hint: string; -}[] = [ + /** Backend `format_type` for a "plain" save (unused for compressed/torchao). */ + formatType?: string; +}; + +/** Kept as a string alias for back-compat with callers that typed the old union. */ +export type MergedFormat = string; + +export const MERGED_FORMATS: MergedFormatOption[] = [ + // 16-bit { - value: "16-bit (FP16)", + value: "16-bit", label: "16-bit", + bits: 16, + backend: "plain", + group: "16-bit", + common: true, + needsNvidia: false, hint: "Full precision, runs anywhere.", + formatType: "16-bit (FP16)", + }, + // 8-bit + { + value: "fp8", + label: "FP8", + bits: 8, + backend: "compressed", + group: "FP8", + common: true, + needsNvidia: true, + hint: "Dynamic per-token FP8 (W8A8) for vLLM. Data-free.", }, { - value: "FP8 (compressed-tensors)", - label: "FP8 (vLLM)", - hint: "compressed-tensors FP8 for vLLM. Needs an NVIDIA GPU.", + value: "torchao_fp8", + label: "FP8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic FP8 (torchao). Produces on any hardware; loads in vLLM.", }, { - value: "NVFP4 (compressed-tensors)", - label: "NVFP4 (vLLM)", - hint: "compressed-tensors NVFP4 for vLLM. Needs an NVIDIA GPU; calibrates.", + value: "w8a8", + label: "INT8 (W8A8)", + bits: 8, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "8-bit weights and 8-bit activations for vLLM. Data-free.", + }, + { + value: "torchao_int8", + label: "INT8 (portable)", + bits: 8, + backend: "torchao", + group: "Portable", + common: true, + needsNvidia: false, + hint: "Device-agnostic INT8 (torchao). Produces on any hardware; loads in vLLM.", + }, + { + value: "fp8_static", + label: "FP8 Static", + bits: 8, + backend: "compressed", + group: "FP8", + common: false, + needsNvidia: true, + needsCalibration: true, + hint: "Static per-tensor FP8. Calibrates on data.", + }, + { + value: "w8a16", + label: "INT8 (W8A16)", + bits: 8, + backend: "compressed", + group: "INT", + common: false, + needsNvidia: true, + hint: "8-bit weight-only. Data-free.", + }, + { + value: "mxfp8", + label: "MXFP8", + bits: 8, + backend: "compressed", + group: "MXFP", + common: false, + needsNvidia: true, + hint: "Microscaling FP8. Needs a newer compressed-tensors stack.", + }, + // 4-bit + { + value: "w4a16", + label: "INT4 (W4A16)", + bits: 4, + backend: "compressed", + group: "INT", + common: true, + needsNvidia: true, + hint: "4-bit weight-only (GPTQ-style) for vLLM. Data-free.", + }, + { + value: "mxfp4", + label: "MXFP4", + bits: 4, + backend: "compressed", + group: "MXFP", + common: true, + needsNvidia: true, + hint: "Microscaling FP4 (W4A4) for vLLM. Data-free.", + }, + { + value: "nvfp4", + label: "NVFP4", + bits: 4, + backend: "compressed", + group: "FP4", + common: true, + needsNvidia: true, + needsCalibration: true, + hint: "NVIDIA FP4 (W4A4) for vLLM. Calibrates on data.", }, ]; +/** Look up a merged format option by its stable value. */ +export function findMergedFormat(value: string): MergedFormatOption | undefined { + return MERGED_FORMATS.find((f) => f.value === value); +} + +/** Backend payload for one merged format: plain -> formatType, compressed/torchao -> the alias. */ +export function mergedFormatPayload(value: string): { + formatType: string; + compressedMethod: string | null; +} { + const opt = findMergedFormat(value); + if (!opt || opt.backend === "plain") { + return { + formatType: opt?.formatType ?? "16-bit (FP16)", + compressedMethod: null, + }; + } + return { formatType: "16-bit (FP16)", compressedMethod: opt.value }; +} + /** * llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16. * K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0 diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 7b5c6ec6d4..07606a26ed 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -24,6 +24,19 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { + DropdownMenu, + DropdownMenuCheckboxItem, + DropdownMenuContent, + DropdownMenuLabel, + DropdownMenuSeparator, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { + Alert, + AlertDescription, + AlertTitle, +} from "@/components/ui/alert"; import { Separator } from "@/components/ui/separator"; import { Spinner } from "@/components/ui/spinner"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; @@ -63,11 +76,14 @@ import { type ExportMethod, GUIDE_STEPS, MERGED_FORMATS, - type MergedFormat, + type MergedFormatOption, + mergedFormatPayload, QUANT_OPTIONS, buildQuantSizeLabels, getEstimatedSize, } from "./constants"; +import { useHardwareInfo } from "@/hooks/use-hardware-info"; +import { usePlatformStore } from "@/config/env"; import { isExportPanelActive, useExportRuntimeStore, @@ -78,6 +94,10 @@ import { exportTourSteps } from "./tour"; const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]); +// GGUF LoRA output float types (Q8_0 first / default). Q8_0 falls back to F16 per tensor for dims +// not divisible by the block size (32); no "auto" - the choice is explicit. +const LORA_GGUF_OUTTYPES = ["q8_0", "f16", "bf16", "f32"] as const; + type SourceTab = "local" | "checkpoint" | "hf"; type SourceMode = "checkpoint" | "model"; @@ -109,7 +129,16 @@ function buildRelativeSaveDirectory( : sourceBaseModelName; return `${safePathSegment(rawName)}-GGUF`; } - return `${selectedModelIdx ?? "model"}/${checkpoint}`; + // Merged / LoRA: a checkpoint keeps the "/" layout under outputs. + if (sourceMode === "checkpoint" && selectedModelIdx && checkpoint) { + return `${selectedModelIdx}/${checkpoint}`; + } + // Local / HF source (no checkpoint): name from the model id to avoid "model/null". + const rawName = + sourceMode === "checkpoint" + ? checkpoint ?? selectedModelIdx ?? sourceBaseModelName + : sourceBaseModelName; + return `${safePathSegment(rawName)}-${exportMethod === "lora" ? "adapter" : "merged"}`; } function siblingGgufDirectory(sourcePath: string): string | null { @@ -177,9 +206,57 @@ export function ExportPage() { }); // GGUF importance matrix (required for the IQ quants) and merged-export precision. const [useImatrix, setUseImatrix] = useState(false); - const [mergedFormat, setMergedFormat] = useState("16-bit (FP16)"); - // IQ quants are imatrix-only, so force it on when one is selected; otherwise we would submit - // an IQ quant with no imatrix and llama.cpp would reject it. + // Merged precision: one or more MERGED_FORMATS values, exported in one run. Seed from a live run + // so navigating away and back (which remounts this page) keeps the selection, like exportMethod. + const [selectedFormats, setSelectedFormats] = useState(() => { + const s = useExportRuntimeStore.getState(); + return isExportPanelActive(s) && + s.summary?.method === "merged" && + s.summary.mergedFormats.length > 0 + ? s.summary.mergedFormats + : ["16-bit"]; + }); + // LoRA-only export: optionally also emit a GGUF LoRA adapter, and its output float type. + const [loraAsGguf, setLoraAsGguf] = useState(false); + const [loraGgufOuttype, setLoraGgufOuttype] = useState("q8_0"); + // GGUF method: export the full model as GGUF quants, or (for an adapter checkpoint) a GGUF LoRA. + const [ggufTarget, setGgufTarget] = useState<"model" | "lora">("model"); + + const hardware = useHardwareInfo(); + // GGUF LoRA conversion is rejected on the macOS / MLX path, so gate it out on a Mac host. + const isMacHost = usePlatformStore((s) => s.deviceType) === "mac"; + // Real CUDA (not ROCm); gates the NVIDIA-only compressed-tensors formats. + const hasNvidia = hardware.cuda != null && hardware.rocm == null; + // Only gray out on an authoritative unsupported response; while unloaded the backend route guard + // stays authoritative. The backend supplies the precise reason; the fallback below is a backstop. + const exportUnsupported = + hardware.loaded && hardware.exportSupported === false; + const exportUnsupportedMessage = + hardware.exportUnsupportedMessage ?? + "Export requires a supported accelerator (NVIDIA, AMD, or Intel GPU, or Apple Silicon) with PyTorch or MLX installed."; + const availableFormats = useMemo( + () => + MERGED_FORMATS.filter((f) => { + // compressed-tensors (llm-compressor) is the NVIDIA path; shown only on an NVIDIA GPU. + if (f.backend === "compressed") return hasNvidia; + // Portable torchao is the fallback for hosts without the NVIDIA compressed path, i.e. a + // CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors) and on macOS/MLX (the + // backend rejects quantized export there). + if (f.backend === "torchao") return !hasNvidia && !isMacHost; + // Plain 16-bit is available everywhere. + return true; + }), + [hasNvidia, isMacHost], + ); + const toggleFormat = useCallback((value: string) => { + setSelectedFormats((prev) => + prev.includes(value) + ? prev.filter((v) => v !== value) + : [...prev, value], + ); + }, []); + // availableFormats already drops NVIDIA-only formats on other hardware, so no pruning needed. + // IQ quants are imatrix-only: force imatrix on when one is selected, else llama.cpp rejects it. const requiresImatrix = quantLevels.some( (q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix, ); @@ -304,6 +381,11 @@ export function ExportPage() { const baseModelName = selectedModelData?.base_model ?? "—"; const isAdapter = !!selectedModelData?.peft_type; const isQuantized = !!selectedModelData?.is_quantized; + // isAdapter / isQuantized come from the checkpoint's metadata and are stale in "model" source + // mode (a direct base export), so treat both as false outside checkpoint mode to avoid wrongly + // gating the methods. + const effectiveIsAdapter = sourceMode === "checkpoint" && isAdapter; + const effectiveIsQuantized = sourceMode === "checkpoint" && isQuantized; const loraRank = selectedModelData?.lora_rank ?? null; const trainingMethodLabel = selectedModelData?.peft_type ? "LoRA / QLoRA" @@ -416,25 +498,30 @@ export function ExportPage() { setCheckpoint(null); }, [selectedModelIdx]); - // For a ?run= deep link, default to the run's main checkpoint. Declared after - // the reset effect above so it runs last and isn't clobbered back to null. + // Default to the newest checkpoint when none is chosen (checkpoints are sorted newest-first). + // Declared after the reset effect above so it runs last and isn't clobbered back to null. Covers + // both a ?run= deep link and a plain finetune opened without an explicit checkpoint pick. useEffect(() => { - if (appliedRunRef.current == null) return; - if (appliedRunRef.current !== selectedModelIdx) return; + if (sourceMode !== "checkpoint") return; if (checkpoint != null || checkpointsForModel.length === 0) return; setCheckpoint(checkpointsForModel[0].display_name); - }, [selectedModelIdx, checkpoint, checkpointsForModel]); + }, [sourceMode, selectedModelIdx, checkpoint, checkpointsForModel]); // Auto-reset export method if incompatible with the selected model type useEffect(() => { - if (!isAdapter && (exportMethod === "merged" || exportMethod === "lora")) { + // Only LoRA needs a real adapter; Merged and GGUF work for non-PEFT base models too. + if (!effectiveIsAdapter && exportMethod === "lora") { setExportMethod(null); } // Quantized non-PEFT models can't export to any format - if (!isAdapter && isQuantized && exportMethod !== null) { + if (!effectiveIsAdapter && effectiveIsQuantized && exportMethod !== null) { setExportMethod(null); } - }, [isAdapter, isQuantized, exportMethod]); + // The GGUF LoRA target only applies to an adapter checkpoint on a non-Mac host. + if ((!effectiveIsAdapter || isMacHost) && ggufTarget !== "model") { + setGgufTarget("model"); + } + }, [effectiveIsAdapter, effectiveIsQuantized, exportMethod, isMacHost, ggufTarget]); const handleSourceTabChange = useCallback((next: string) => { if (next === "checkpoint") { @@ -442,7 +529,7 @@ export function ExportPage() { } else if (next === "hf" || next === "local") { setSourceMode("model"); setModelSource(next); - setExportMethod("gguf"); + // Don't force GGUF: Local / HF sources can export Merged too; a stale LoRA pick auto-resets. } else { return; } @@ -508,10 +595,26 @@ export function ExportPage() { sourceMode, ]); const saveDirectory = customSaveDirectory?.trim() || defaultSaveDirectory; + // Each merged format uploads a full model to the repo root, so several to one repo would collide. + // GGUF method exporting an adapter checkpoint as a GGUF LoRA (vs full-model quants). Reuses the + // LoRA-adapter export path; no quant list needed. + const ggufAsLora = + exportMethod === "gguf" && + ggufTarget === "lora" && + effectiveIsAdapter && + !isMacHost; + + // Restrict a Hub merged export to a single format; multi-format stays available for local export. + const hubMultiFormat = + destination === "hub" && exportMethod === "merged" && selectedFormats.length > 1; + const canExport = !!( selectedExportSource && exportMethod && - (exportMethod !== "gguf" || quantLevels.length > 0) + !exportUnsupported && + !hubMultiFormat && + (exportMethod !== "gguf" || ggufAsLora || quantLevels.length > 0) && + (exportMethod !== "merged" || selectedFormats.length > 0) ); const applyHfSourceModel = useCallback((value: string) => { @@ -576,9 +679,14 @@ export function ExportPage() { const handleStart = useCallback(async () => { const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel; if (!source || !exportMethod) return; - // A GGUF export with no quant selected runs zero exports yet would still - // settle as success with no file; require at least one (mirrors canExport). - if (exportMethod === "gguf" && quantLevels.length === 0) return; + // No supported accelerator (or PyTorch/MLX missing): the backend would reject anyway; don't submit. + if (exportUnsupported) return; + // GGUF with no quant, or merged with no format, would run an unintended/empty export; require + // at least one (mirrors canExport, in case the panel's Start button bypasses the outer one). + if (exportMethod === "gguf" && !ggufAsLora && quantLevels.length === 0) return; + if (exportMethod === "merged" && selectedFormats.length === 0) return; + // A Hub merged push writes each format to the repo root; several would collide (mirrors canExport). + if (hubMultiFormat) return; const selectedCp = sourceMode === "checkpoint" ? checkpointsForModel.find((cp) => cp.display_name === checkpoint) @@ -591,8 +699,13 @@ export function ExportPage() { ? `${hfUsername}/${modelName}` : undefined; const token = pushToHub && hfToken ? hfToken : undefined; - const methodLabel = - EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod; + // The GGUF method with the LoRA target reuses the LoRA-adapter export path. + const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod; + const emitLoraGguf = + ggufAsLora || (effectiveMethod === "lora" && loraAsGguf && !isMacHost); + const methodLabel = ggufAsLora + ? "GGUF LoRA adapter" + : (EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod); const adapterExport = sourceMode === "checkpoint" && isAdapter; // Consent gate for an HF source's custom (auto_map) code, run before we hand @@ -624,11 +737,16 @@ export function ExportPage() { trustRemoteCode, approvedRemoteCodeFingerprint, loadToken: hfToken || null, - exportMethod, + exportMethod: effectiveMethod, isAdapter: adapterExport, quantLevels, useImatrix: effectiveImatrix, - mergedFormat, + mergedSelections: selectedFormats.map((v) => ({ + ...mergedFormatPayload(v), + label: MERGED_FORMATS.find((f) => f.value === v)?.label ?? v, + })), + loraGguf: emitLoraGguf, + loraGgufOuttype, saveDirectory, destination, repoId, @@ -639,8 +757,9 @@ export function ExportPage() { baseModelName: sourceBaseModelName, checkpointLabel: selectedExportSource, methodLabel, - method: exportMethod, + method: effectiveMethod, quantLevels, + mergedFormats: exportMethod === "merged" ? selectedFormats : [], destination, }, }); @@ -656,7 +775,13 @@ export function ExportPage() { isAdapter, quantLevels, effectiveImatrix, - mergedFormat, + selectedFormats, + hubMultiFormat, + ggufAsLora, + loraAsGguf, + isMacHost, + loraGgufOuttype, + exportUnsupported, destination, saveDirectory, hfUsername, @@ -1163,75 +1288,303 @@ export function ExportPage() {
+ {exportUnsupported && ( + + + Export unavailable + {exportUnsupportedMessage} + + )} + - {exportMethod === "merged" && isAdapter && ( -
-
Precision
-
- {MERGED_FORMATS.map((f) => ( - - ))} -
-
- {MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint} + {exportMethod === "merged" && !exportUnsupported && ( +
+
+
+
Precision
+ + — select one or more + +
+
+ {availableFormats + .filter((f) => f.common) + .map((f) => { + const active = selectedFormats.includes(f.value); + return ( + + ); + })} + + {availableFormats.some((f) => !f.common) && ( + + + + + + + Additional formats + + + {availableFormats + .filter((f) => !f.common) + .map((f) => ( + toggleFormat(f.value)} + onSelect={(e) => e.preventDefault()} + > + + + {f.label} + {f.needsCalibration ? " *" : ""} + + + {f.hint} + + + + ))} + + + )} +
+ + {selectedFormats.length > 0 && ( +
+ + {selectedFormats.length} selected:{" "} + {selectedFormats + .map( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.label ?? v, + ) + .join(", ")} + + {selectedFormats.length > 1 && ( + + )} +
+ )} + + {hubMultiFormat && ( +
+ Hub export supports one format at a time (each writes to + the repository root). Select a single format, or export + locally to produce several at once. +
+ )} + + {selectedFormats.some( + (v) => + MERGED_FORMATS.find((f) => f.value === v) + ?.needsCalibration, + ) && ( +
+ * calibrates on data (uses a small calibration set). +
+ )} + + {!hasNvidia && ( +
+ No NVIDIA GPU detected: compressed-tensors formats are + hidden. 16-bit and portable FP8/INT8 (torchao) still + work here and load in vLLM. +
+ )}
)} - {exportMethod === "gguf" && ( - <> - -
-
-
- Importance matrix (imatrix) + {exportMethod === "lora" && effectiveIsAdapter && !exportUnsupported && ( +
+
+
Adapter format
+
+ + +
+
+ {isMacHost + ? "GGUF LoRA is not available on macOS/MLX; exporting the safetensors adapter." + : loraAsGguf + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Standard PEFT adapter files. Pair with the base model at inference."} +
+
+ + {loraAsGguf && ( +
+
Output type
+ +
+ )} +
+ )} + + {exportMethod === "gguf" && !exportUnsupported && ( +
+ {effectiveIsAdapter && !isMacHost && ( +
+
Export target
+
+ +
- {requiresImatrix - ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." - : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} + {ggufTarget === "lora" + ? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate." + : "Merges the adapter into the base model, then quantizes the full model to GGUF."}
- -
- + )} + + {ggufAsLora ? ( +
+
Output type
+ +
+ ) : ( + <> + +
+
+
+ Importance matrix (imatrix) +
+
+ {requiresImatrix + ? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model." + : "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."} +
+
+ +
+ + )} +
)} {estimatedSize && (
diff --git a/studio/frontend/src/features/export/stores/export-runtime-store.ts b/studio/frontend/src/features/export/stores/export-runtime-store.ts index a87e7d93e9..6ee5c78994 100644 --- a/studio/frontend/src/features/export/stores/export-runtime-store.ts +++ b/studio/frontend/src/features/export/stores/export-runtime-store.ts @@ -5,7 +5,6 @@ import { create } from "zustand"; import { cancelExport, cleanupExport, - exportBase, exportGGUF, exportLoRA, exportMerged, @@ -119,6 +118,8 @@ export interface ExportRunSummary { methodLabel: string; method: ExportMethod; quantLevels: string[]; + /** Merged: the selected format values (for the summary "Formats" row and to reseed the picker). */ + mergedFormats: string[]; destination: ExportDestination; } @@ -140,8 +141,16 @@ export interface RunExportParams { quantLevels: string[]; /** GGUF: use an importance matrix (auto-download); required for the IQ quants. */ useImatrix?: boolean; - /** Merged: precision/format ("16-bit (FP16)" or a compressed-tensors option). */ - mergedFormat?: string; + /** Merged: precision formats, each exported to its own sibling directory. Defaults to 16-bit. + * `label` is the display name for the success banner's per-format output line. */ + mergedSelections?: { + formatType: string; + compressedMethod: string | null; + label: string; + }[]; + /** LoRA: also emit a GGUF LoRA adapter (llama.cpp `--lora`), and its output float type. */ + loraGguf?: boolean; + loraGgufOuttype?: string; saveDirectory: string; destination: ExportDestination; repoId?: string; @@ -172,7 +181,13 @@ interface ExportRuntimeState { * settling the run by polling /api/export/status instead. Logs keep streaming. */ reconnecting: boolean; startedAt: number | null; - result: { outputPath: string | null; destination: ExportDestination } | null; + /** `outputPath` is the first path (back-compat); `outputPaths` is one entry per written folder + * so a multi-format merged run can list every sibling directory it created. */ + result: { + outputPath: string | null; + outputPaths: { label: string; path: string }[]; + destination: ExportDestination; + } | null; error: string | null; cancelRequested: boolean; hasHydrated: boolean; @@ -317,6 +332,10 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success" as const, result: { outputPath: status.last_op_output_path ?? null, + // A run recovered from the backend only knows the last output path. + outputPaths: status.last_op_output_path + ? [{ label: "", path: status.last_op_output_path }] + : [], destination: state.result?.destination ?? "local", }, }; @@ -351,7 +370,9 @@ export const useExportRuntimeStore = create()((set, get) => const quantTotal = params.exportMethod === "gguf" ? Math.max(1, params.quantLevels.length) - : 1; + : params.exportMethod === "merged" + ? Math.max(1, params.mergedSelections?.length ?? 1) + : 1; set({ runId, @@ -431,67 +452,72 @@ export const useExportRuntimeStore = create()((set, get) => } if (!isCurrent()) return; - // 2. Run the export. Capture the resolved output_path for the success - // banner; multi-quant GGUF shares one directory, so keep the last. + // 2. Run the export. Collect every resolved output_path so the success + // banner can list each sibling directory a multi-format run created. set({ phase: "exporting" }); - let lastOutputPath: string | null = null; + const outputs: { label: string; path: string }[] = []; if (params.exportMethod === "merged") { - if (params.isAdapter) { + // Each selected format writes its own sibling directory (PEFT or non-PEFT base alike). + const selections = + params.mergedSelections && params.mergedSelections.length > 0 + ? params.mergedSelections + : [{ formatType: "16-bit (FP16)", compressedMethod: null, label: "16-bit" }]; + for (let i = 0; i < selections.length; i += 1) { + if (!isCurrent()) return; + set({ quantIndex: i }); + const sel = selections[i]; const { outputPath } = await runRecoverableOp(() => exportMerged({ save_directory: params.saveDirectory, - format_type: params.mergedFormat, + format_type: sel.formatType, + compressed_method: sel.compressedMethod, push_to_hub: pushToHub, repo_id: params.repoId, hf_token: params.token, private: params.privateRepo, }), ); - lastOutputPath = outputPath; - } else { - const { outputPath } = await runRecoverableOp(() => - exportBase({ - save_directory: params.saveDirectory, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - private: params.privateRepo, - base_model_id: params.baseModelId, - }), - ); - lastOutputPath = outputPath; - } - } else if (params.exportMethod === "gguf") { - for (let i = 0; i < params.quantLevels.length; i += 1) { - if (!isCurrent()) return; - set({ quantIndex: i }); - const quant = params.quantLevels[i]; - const { outputPath } = await runRecoverableOp(() => - exportGGUF({ - save_directory: params.saveDirectory, - quantization_method: quant, - push_to_hub: pushToHub, - repo_id: params.repoId, - hf_token: params.token, - imatrix: params.useImatrix, - }), - ); - lastOutputPath = outputPath ?? lastOutputPath; + if (outputPath) outputs.push({ label: sel.label, path: outputPath }); if (!isCurrent()) return; set({ quantIndex: i + 1 }); } + } else if (params.exportMethod === "gguf") { + // Send the whole quant list in ONE call: the model is merged once and every GGUF comes + // from that single merge (unsloth save_to_gguf loops internally). + const { outputPath } = await runRecoverableOp(() => + exportGGUF({ + save_directory: params.saveDirectory, + quantization_method: params.quantLevels, + push_to_hub: pushToHub, + repo_id: params.repoId, + hf_token: params.token, + imatrix: params.useImatrix, + }), + ); + if (outputPath) outputs.push({ label: "GGUF", path: outputPath }); + if (!isCurrent()) return; + set({ quantIndex: get().quantTotal }); } else if (params.exportMethod === "lora") { const { outputPath } = await runRecoverableOp(() => exportLoRA({ save_directory: params.saveDirectory, push_to_hub: pushToHub, repo_id: params.repoId, - hf_token: params.token, + // A local GGUF LoRA export still reloads a possibly-gated base config, so fall back to + // the load token when there is no hub-upload token (both are the same HF token). + hf_token: params.token ?? params.loadToken ?? null, private: params.privateRepo, + gguf: params.loraGguf ?? false, + gguf_outtype: params.loraGgufOuttype ?? "q8_0", }), ); - lastOutputPath = outputPath; + if (outputPath) { + outputs.push({ + label: params.loraGguf ? "GGUF LoRA adapter" : "LoRA adapter", + path: outputPath, + }); + } } if (!isCurrent()) return; @@ -499,7 +525,11 @@ export const useExportRuntimeStore = create()((set, get) => phase: "success", isExporting: false, reconnecting: false, - result: { outputPath: lastOutputPath, destination: params.destination }, + result: { + outputPath: outputs[0]?.path ?? null, + outputPaths: outputs, + destination: params.destination, + }, }); } catch (err) { if (!isCurrent()) return; diff --git a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx index 7895a89254..38464d36e4 100644 --- a/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx +++ b/studio/frontend/src/features/hub/catalog/hub-option-menu.tsx @@ -38,6 +38,7 @@ export function HubOptionMenu({ showChevron = true, title, triggerContent, + footer, }: { value: T; options: readonly HubOption[]; @@ -49,9 +50,12 @@ export function HubOptionMenu({ showChevron?: boolean; title?: string; triggerContent?: ReactNode; + /** Rendered under the options behind a separator; clicks keep the menu open. */ + footer?: ReactNode; }) { const [open, setOpen] = useState(false); - const [activeIndex, setActiveIndex] = useState(0); + // -1 = nothing highlighted (no hover, no keyboard nav yet). + const [activeIndex, setActiveIndex] = useState(-1); const triggerRef = useRef(null); const listboxRef = useRef(null); const idBase = useId(); @@ -63,9 +67,9 @@ export function HubOptionMenu({ }, [options, value]); const selected = options[selectedIndex]; const resolvedActiveIndex = - options.length === 0 + options.length === 0 || activeIndex < 0 ? -1 - : Math.min(Math.max(activeIndex, 0), options.length - 1); + : Math.min(activeIndex, options.length - 1); const activeOptionId = resolvedActiveIndex >= 0 ? `${idBase}-option-${resolvedActiveIndex}` : undefined; @@ -92,11 +96,13 @@ export function HubOptionMenu({ (nextOpen: boolean) => { setOpen(nextOpen); if (nextOpen) { - activateIndex(selectedIndex); + // Nothing highlighted until the user hovers or uses the keyboard; + // keyboard nav anchors on the selected option (handleContentKeyDown). + activateIndex(-1); requestAnimationFrame(() => listboxRef.current?.focus()); } }, - [activateIndex, selectedIndex], + [activateIndex], ); const handleContentKeyDown = useCallback( @@ -112,12 +118,21 @@ export function HubOptionMenu({ } if (event.key === "ArrowDown") { event.preventDefault(); - setActiveIndex((currentIndex + 1) % options.length); + // First arrow press highlights the selected option, then steps. + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex + 1) % options.length, + ); return; } if (event.key === "ArrowUp") { event.preventDefault(); - setActiveIndex((currentIndex - 1 + options.length) % options.length); + setActiveIndex( + resolvedActiveIndex < 0 + ? selectedIndex + : (currentIndex - 1 + options.length) % options.length, + ); return; } if (event.key === "Home") { @@ -197,6 +212,7 @@ export function HubOptionMenu({ aria-activedescendant={activeOptionId} tabIndex={0} onKeyDown={handleContentKeyDown} + onPointerLeave={() => activateIndex(-1)} className="outline-none" > {options.map((option, index) => { @@ -235,6 +251,12 @@ export function HubOptionMenu({ ); })}
+ {footer && ( + // -mt-3 cancels the surface's 16px flex gap down to 4px. No side + // padding: the footer label carries the same padding as the options + // so its checkbox lines up with the option text. +
{footer}
+ )} ); diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx index f5afce5c59..10d5a80e5e 100644 --- a/studio/frontend/src/features/hub/catalog/models-header.tsx +++ b/studio/frontend/src/features/hub/catalog/models-header.tsx @@ -15,6 +15,7 @@ import { PackageIcon, RamMemoryIcon, RemoveCircleIcon, + CpuIcon } from "@hugeicons/core-free-icons"; import type { IconSvgElement } from "@hugeicons/react"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -43,6 +44,7 @@ export function ModelsHeader({ isDataset, gpuLabel, ramLabel, + coreLabel, activeCheckpoint, activeGgufVariant, onTitleClick, @@ -53,6 +55,7 @@ export function ModelsHeader({ isDataset: boolean; gpuLabel: string; ramLabel: string; + coreLabel: string; activeCheckpoint: string | null; activeGgufVariant: string | null; onTitleClick: () => void; @@ -84,7 +87,8 @@ export function ModelsHeader({ value={String(localCount)} /> - + + {activeCheckpoint && (
diff --git a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx index 48f7fcffaa..7c08c9c486 100644 --- a/studio/frontend/src/features/hub/catalog/models-toolbar.tsx +++ b/studio/frontend/src/features/hub/catalog/models-toolbar.tsx @@ -1,6 +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 +import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { @@ -68,6 +69,8 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange, capabilityFilter, onCapabilityFilterChange, + fitOnDeviceOnly, + onFitOnDeviceOnlyChange, onManageLocalFolders, onOpenFineTune, }: { @@ -84,6 +87,9 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onFormatFilterChange: (value: ModelFormatFilter) => void; capabilityFilter: CapabilityFilter; onCapabilityFilterChange: (value: CapabilityFilter) => void; + /** Shared with the chat model selector: hide models over the device budget. */ + fitOnDeviceOnly: boolean; + onFitOnDeviceOnlyChange: (value: boolean) => void; onManageLocalFolders: () => void; /** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a * format-dropdown option rather than a standalone feed section. */ @@ -350,6 +356,33 @@ export const ModelsToolbar = memo(function ModelsToolbar({ onValueChange={onSortChange} ariaLabel="Sort models" className={cn(triggerBase, "w-[128px]")} + footer={ + isDataset ? undefined : ( + + + + + + Hides models larger than this device's memory budget. + Downloaded models stay visible. + + + ) + } /> )} diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index c696862283..c3920e18f6 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -5,6 +5,7 @@ import { loadRememberedLoadSettings, rememberedLoadSettingsKey, } from "@/components/assistant-ui/model-selector/remembered-load-settings"; +import { hfModelFitsDevice } from "@/components/assistant-ui/model-selector/recommended-fit"; import { useHubInventory } from "@/features/hub/inventory"; import { useDebouncedValue } from "@/hooks/use-debounced-value"; import { useGpuInfo } from "@/hooks/use-gpu-info"; @@ -327,6 +328,9 @@ export function ModelsPage() { const activeCheckpoint = checkpoint && !isExternalModelId(checkpoint) ? checkpoint : null; const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + // Shared with the chat model selector: list only models sized for this device. + const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly); + const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly); useEffect(() => { let cancelled = false; @@ -697,7 +701,12 @@ export function ModelsPage() { !isHiddenModelId(row.id) && matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) && matchesCapability(row.capabilities, deferredCapabilityFilter) && - (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)), + (!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)) && + // Models already on disk stay visible regardless of device fit, + // matching the chat model selector. + (!fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu)), ); }, [ discoverRows, @@ -705,6 +714,8 @@ export function ModelsPage() { effectiveDiscoverFormat, deferredCapabilityFilter, activeChannel, + fitOnDeviceOnly, + gpu, ]); const listRows = filteredDiscoverRows; @@ -724,8 +735,21 @@ export function ModelsPage() { effectiveLocalRows, ) .filter((row) => !isHiddenModelId(row.id)) - .filter((row) => matchesFormat(row.result.isGguf, "gguf")), - [hubFeed.trending.results, modelDiscoveryInventorySignature], + .filter((row) => matchesFormat(row.result.isGguf, "gguf")) + // Same fit filter as the main Discover list, so the feed carousel + // honors the toggle too. + .filter( + (row) => + !fitOnDeviceOnly || + row.isAvailableOnDevice || + hfModelFitsDevice(row.result, gpu), + ), + [ + hubFeed.trending.results, + modelDiscoveryInventorySignature, + fitOnDeviceOnly, + gpu, + ], ); const feedRows = useMemo(() => { if (!isFeedMode) return []; @@ -1061,11 +1085,15 @@ export function ModelsPage() { const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu); const gpuLabel = gpu.available - ? `${Math.floor(gpu.memoryTotalGb)} GB` + ? `${Math.round(gpu.memoryTotalGb)} GB` : "Unavailable"; const ramLabel = - gpu.systemRamAvailableGb > 0 - ? `${Math.floor(gpu.systemRamAvailableGb)} GB` + gpu.systemRamTotalGb > 0 + ? `${Math.round(gpu.systemRamTotalGb)} GB` + : "Unavailable"; + const coreLabel = + gpu.cpuCore > 0 && gpu.cpuThread > 0 + ? `${gpu.cpuCore}/${gpu.cpuThread}` : "Unavailable"; const openNewChat = useCallback(() => { @@ -1429,6 +1457,7 @@ export function ModelsPage() { isDataset={isDatasetMode} gpuLabel={gpuLabel} ramLabel={ramLabel} + coreLabel={coreLabel} activeCheckpoint={activeCheckpoint} activeGgufVariant={activeGgufVariant} onTitleClick={handleResetToDiscover} @@ -1448,6 +1477,8 @@ export function ModelsPage() { onFormatFilterChange={setFormatFilter} capabilityFilter={capabilityFilter} onCapabilityFilterChange={setCapabilityFilter} + fitOnDeviceOnly={fitOnDeviceOnly} + onFitOnDeviceOnlyChange={setFitOnDeviceOnly} onManageLocalFolders={handleManageLocalFolders} onOpenFineTune={() => handleOpenList("finetune")} /> diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts new file mode 100644 index 0000000000..8b6bc7ee7f --- /dev/null +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -0,0 +1,82 @@ +// 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 { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type EmbeddingModelSettings = { + embeddingModel: string; + defaultEmbeddingModel: string; + isCustom: boolean; +}; + +type ApiEmbeddingModelSettings = { + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + default_embedding_model: string; + // biome-ignore lint/style/useNamingConvention: API schema + is_custom: boolean; +}; + +/** 409 from the backend: the model could not be verified as an embedding model + * (wrong type, gated repo, or offline). Retry with force to save anyway. */ +export class EmbeddingModelVerificationError extends Error {} + +function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { + return { + embeddingModel: settings.embedding_model, + defaultEmbeddingModel: settings.default_embedding_model, + isCustom: settings.is_custom, + }; +} + +export async function loadEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load embedding model setting"), + ); + } + return fromApi(await res.json()); +} + +export async function updateEmbeddingModelSettings( + embeddingModel: string, + options?: { hfToken?: string; force?: boolean }, +): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + // biome-ignore lint/style/useNamingConvention: API schema + embedding_model: embeddingModel, + // biome-ignore lint/style/useNamingConvention: API schema + hf_token: options?.hfToken || null, + force: options?.force ?? false, + }), + }); + if (res.status === 409) { + throw new EmbeddingModelVerificationError( + await readFastApiError(res, "Could not verify the embedding model"), + ); + } + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to save embedding model"), + ); + } + return fromApi(await res.json()); +} + +export async function resetEmbeddingModelSettings(): Promise { + const res = await authFetch("/api/settings/embedding-model", { + method: "DELETE", + }); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to reset embedding model"), + ); + } + return fromApi(await res.json()); +} diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts new file mode 100644 index 0000000000..9e87922970 --- /dev/null +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -0,0 +1,73 @@ +// 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 the `unsloth start ` command for the API-keys panel. +// `unsloth start` reads UNSLOTH_STUDIO_URL (default 127.0.0.1:8888) and only +// auto-mints a key for a loopback server, so the bare command is correct only for +// the default local server. For a non-default port or tunnel/remote base, emit the +// URL (plus a key for non-loopback) so the copy targets what the UI shows. + +const DEFAULT_STUDIO_PORT = "8888"; +const DEFAULT_AGENT = "claude"; + +// URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is +// "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below. +function normalizeHost(host: string): string { + const lower = host.toLowerCase(); + return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; +} + +// The bare `unsloth start` probes exactly http://127.0.0.1:8888, so only that literal +// host earns the bare command. `localhost` can resolve to ::1 (and `::1` is never +// probed), so both keep an explicit UNSLOTH_STUDIO_URL -- harmless when they alias +// 127.0.0.1, correct when they don't. +function isDefaultLocalHost(host: string): boolean { + return host === "127.0.0.1"; +} + +// Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. +function isLoopbackHost(host: string): boolean { + if (host === "localhost" || host === "::1") return true; + const octets = host.split("."); + return ( + octets.length === 4 && + octets[0] === "127" && + octets.every((o) => /^\d{1,3}$/.test(o) && Number(o) <= 255) + ); +} + +export function buildAgentCommand( + base: string | null | undefined, + key: string | null | undefined, + os: "unix" | "windows", + agent: string = DEFAULT_AGENT, +): string { + const bare = `unsloth start ${agent}`; + + let url: URL | null = null; + try { + if (base) url = new URL(base); + } catch { + url = null; + } + // Unknown base: fall back to the bare default-local command. + if (!url) return bare; + + const host = normalizeHost(url.hostname); + const loopback = isLoopbackHost(host); + // Default local server (http://127.0.0.1/localhost:8888): bare command + // auto-discovers it. The CLI's bare default probes plain HTTP, so an HTTPS + // loopback on the same port must keep its explicit UNSLOTH_STUDIO_URL. + if (url.protocol === "http:" && isDefaultLocalHost(host) && url.port === DEFAULT_STUDIO_PORT) { + return bare; + } + + // Non-default server: set the URL; non-loopback also needs an explicit key. + let cmd = bare; + if (!loopback && key) cmd += ` --api-key ${key}`; + + const studioUrl = url.origin; + return os === "windows" + ? `$env:UNSLOTH_STUDIO_URL="${studioUrl}"; ${cmd}` + : `UNSLOTH_STUDIO_URL=${studioUrl} ${cmd}`; +} diff --git a/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx new file mode 100644 index 0000000000..b0e9a41b7d --- /dev/null +++ b/studio/frontend/src/features/settings/components/embedding-model-combobox.tsx @@ -0,0 +1,134 @@ +// 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 { + Combobox, + ComboboxContent, + ComboboxEmpty, + ComboboxInput, + ComboboxItem, + ComboboxList, +} from "@/components/ui/combobox"; +import { Spinner } from "@/components/ui/spinner"; +import type { PipelineType } from "@huggingface/hub"; +import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search"; +import { useDebouncedValue } from "@/hooks"; +import { type ReactElement, useMemo, useRef } from "react"; + +// HF pipeline filter for embedding models; matches the backend's +// is_embedding_model signals (sentence-similarity / feature-extraction). +const EMBEDDING_TASKS: readonly PipelineType[] = [ + "sentence-similarity", + "feature-extraction", +]; + +type EmbeddingModelComboboxProps = { + value: string; + /** Fires on typing, selection, and Enter with the current text. */ + onChange: (value: string) => void; + accessToken?: string; + disabled?: boolean; + placeholder?: string; + ariaLabel?: string; + className?: string; +}; + +export function EmbeddingModelCombobox({ + value, + onChange, + accessToken, + disabled, + placeholder, + ariaLabel, + className, +}: EmbeddingModelComboboxProps): ReactElement { + const selectingRef = useRef(false); + const anchorRef = useRef(null); + // Fully controlled: the parent updates value on every keystroke, so the + // prop itself is the search query. + const debouncedQuery = useDebouncedValue(value); + + const { results, isLoading } = useHubModelSearch(debouncedQuery, { + task: EMBEDDING_TASKS, + accessToken, + excludeGguf: true, + enabled: !disabled, + // Curated unsloth listing when empty (the global top-downloads page holds + // no unsloth mirrors to float); a typed query searches the whole Hub. + ownerScope: debouncedQuery.trim() ? "all" : "unsloth", + }); + + const items = useMemo(() => { + const ids = results.map((item) => item.id); + const selected = value.trim(); + if (selected && !ids.includes(selected)) { + ids.push(selected); + } + return ids; + }, [results, value]); + + return ( +
{ + if (event.key !== "Enter") return; + if (!(event.target instanceof HTMLInputElement)) return; + event.preventDefault(); + const typed = event.target.value.trim(); + if (typed) { + onChange(typed); + } else if (items.length > 0) { + onChange(items[0]); + } + }} + > + onChange(next ?? "")} + onInputValueChange={(next) => { + if (selectingRef.current) { + selectingRef.current = false; + return; + } + onChange(next); + }} + itemToStringValue={(item) => item} + autoHighlight={true} + > + + + {isLoading ? ( +
+ + Searching... +
+ ) : ( + No embedding models found + )} + + {(id: string) => ( + { + selectingRef.current = true; + }} + > + {id} + + )} + + + +
+ ); +} diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index aef2e7ffe6..c43dd0f219 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -32,45 +32,60 @@ import { loadOpenAIAutoSwitchSettings, updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; +import { buildAgentCommand } from "./agent-command"; -// API call type; OS axis applies to curl only (Python is OS-identical). type ExampleType = | "curl" | "python" + | "javascript" | "curlTools" | "pythonTools" + | "javascriptTools" | "curlAdvanced" - | "pythonAdvanced"; + | "pythonAdvanced" + | "javascriptAdvanced"; type Os = "unix" | "windows"; -// plain = bare call; tools = server-side tools; advanced = sampling + thinking + tools. type Variant = "plain" | "tools" | "advanced"; const TYPE_TABS: { id: ExampleType; label: string }[] = [ { id: "curl", label: "curl" }, { id: "python", label: "Python" }, + { id: "javascript", label: "JavaScript" }, { id: "curlTools", label: "curl + tools" }, { id: "pythonTools", label: "Python + tools" }, + { id: "javascriptTools", label: "JavaScript + tools" }, { id: "curlAdvanced", label: "curl + advanced" }, { id: "pythonAdvanced", label: "Python + advanced" }, + { id: "javascriptAdvanced", label: "JavaScript + advanced" }, ]; const TYPE_LABEL_KEY: Partial> = { curlTools: "settings.apiKeys.exampleCurlTools", pythonTools: "settings.apiKeys.examplePythonTools", + javascriptTools: "settings.apiKeys.exampleJavaScriptTools", curlAdvanced: "settings.apiKeys.exampleCurlAdvanced", pythonAdvanced: "settings.apiKeys.examplePythonAdvanced", + javascriptAdvanced: "settings.apiKeys.exampleJavaScriptAdvanced", }; const OS_AWARE: Record = { curl: true, python: false, + javascript: false, curlTools: true, pythonTools: false, + javascriptTools: false, curlAdvanced: true, pythonAdvanced: false, + javascriptAdvanced: false, }; const CURL_TYPES = new Set(["curl", "curlTools", "curlAdvanced"]); +const JAVASCRIPT_TYPES = new Set([ + "javascript", + "javascriptTools", + "javascriptAdvanced", +]); const PROMPT = "Can Unsloth Studio do API calling?"; // Auto-switch demo: a second call naming a different downloaded GGUF so the @@ -82,7 +97,6 @@ const SWITCH_MODEL = "your-other-downloaded-GGUF"; const SWITCH_PROMPT = "Now answer as a different model."; // web_search + python + terminal are the reliable built-in tools. const TOOLS = ["web_search", "python", "terminal"]; -// Sampling/thinking knobs for the "+ advanced" examples. const ADV = { temperature: 0.7, top_p: 0.8, @@ -93,37 +107,18 @@ const ADV = { } as const; const DOC_LINKS = [ - { - label: "Claude Code", - href: "https://unsloth.ai/docs/basics/claude-code", - }, - { - label: "Codex", - href: "https://unsloth.ai/docs/basics/codex", - }, - { - label: "OpenClaw", - href: "https://unsloth.ai/docs/integrations/openclaw", - }, - { - label: "OpenCode", - href: "https://unsloth.ai/docs/integrations/opencode", - }, - { - label: "Hermes Agent", - href: "https://unsloth.ai/docs/integrations/hermes-agent", - }, + { label: "Claude Code", href: "https://unsloth.ai/docs/basics/claude-code" }, + { label: "Codex", href: "https://unsloth.ai/docs/basics/codex" }, + { label: "OpenClaw", href: "https://unsloth.ai/docs/integrations/openclaw" }, + { label: "OpenCode", href: "https://unsloth.ai/docs/integrations/opencode" }, + { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; -// JSON-encode; also a valid Python literal, so odd model names never break output. const j = (s: string): string => JSON.stringify(s); -// Embed in a POSIX single-quoted string: close, escaped quote, reopen. const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); -// Embed in a PowerShell single-quoted string: '' is a literal quote. const psSingle = (s: string): string => s.replace(/'/g, "''"); const toolsJson = TOOLS.map(j).join(", "); -// Shared body fields (after model/messages, before stream) per variant. function bodyExtraLines(variant: Variant, indent: string): string[] { const lines: string[] = []; if (variant === "advanced") { @@ -152,7 +147,6 @@ function curlBodyPretty(model: string, variant: Variant): string { return `{\n${lines.join("\n")}\n }`; } -// One-line JSON for the Windows body file (PowerShell mangles inline quotes to curl.exe). function winBody(model: string, variant: Variant): string { const body: Record = { model, @@ -172,7 +166,7 @@ function winBody(model: string, variant: Variant): string { body.enabled_tools = TOOLS; } body.stream = true; - return JSON.stringify(body); + return JSON.stringify(body, null, 2); } // A leading comment (valid in both bash and PowerShell) noting the model field @@ -193,7 +187,6 @@ function curlUnix( -d '${shSingle(curlBodyPretty(model, variant))}'`; } -// Windows PowerShell: curl aliases to Invoke-WebRequest, so use curl.exe + body file. function curlWindows( base: string, key: string, @@ -233,7 +226,6 @@ function pythonSnippet( variant: Variant, autoSwitch: boolean, ): string { - // Standard OpenAI args are named; Unsloth extensions go through extra_body. const named = variant === "advanced" ? ` @@ -258,7 +250,6 @@ function pythonSnippet( ${extra.join("\n")} },` : ""; - // With tools, some chunks are tool-lifecycle events with no choices; guard it. const loop = variant !== "plain" ? `for chunk in response: @@ -281,6 +272,70 @@ response = client.chat.completions.create( ${loop}${autoSwitch ? pythonSwitchDemo() : ""}`; } +function javascriptSnippet( + base: string, + key: string, + model: string, + variant: Variant, + autoSwitch: boolean, +): string { + const options: string[] = []; + if (variant === "advanced") { + options.push(` temperature: ${ADV.temperature},`); + options.push(` top_p: ${ADV.top_p},`); + options.push(` max_tokens: ${ADV.max_tokens},`); + } + + // The JS SDK forwards unknown options into the request body, so these go at the + // top level (the Python SDK needs them under extra_body instead). + if (variant === "advanced") { + options.push(` top_k: ${ADV.top_k},`); + options.push(` min_p: ${ADV.min_p},`); + options.push(` repetition_penalty: ${ADV.repetition_penalty},`); + options.push(` enable_thinking: true,`); + } + if (variant !== "plain") { + options.push(` enable_tools: true,`); + options.push(` enabled_tools: [${toolsJson}],`); + } + + const trailingOptions = options.length ? `\n${options.join("\n")}` : ""; + + return `import OpenAI from "openai"; + +const client = new OpenAI({ + baseURL: ${j(`${base}/v1`)}, + apiKey: ${j(key)}, +}); + +const response = await client.chat.completions.create({ + model: ${j(model)}, + messages: [{ role: "user", content: ${j(PROMPT)} }],${trailingOptions} + stream: true, +}); + +for await (const chunk of response) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}${autoSwitch ? javascriptSwitchDemo() : ""}`; +} + +function javascriptSwitchDemo(): string { + return ` + +// "Switch model by request" is on: replace the model below with another GGUF you +// have downloaded and Studio loads it before serving. Unknown names keep serving +// the current model. +const switchResponse = await client.chat.completions.create({ + model: ${j(SWITCH_MODEL)}, + messages: [{ role: "user", content: ${j(SWITCH_PROMPT)} }], + stream: true, +}); + +for await (const chunk of switchResponse) { + process.stdout.write(chunk.choices?.[0]?.delta?.content || ""); +}`; +} + function buildSnippets( base: string, key: string, @@ -292,17 +347,24 @@ function buildSnippets( return { curl: curl(base, key, model, "plain", autoSwitch), python: pythonSnippet(base, key, model, "plain", autoSwitch), + javascript: javascriptSnippet(base, key, model, "plain", autoSwitch), curlTools: curl(base, key, model, "tools", autoSwitch), pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch), + javascriptTools: javascriptSnippet(base, key, model, "tools", autoSwitch), curlAdvanced: curl(base, key, model, "advanced", autoSwitch), pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch), + javascriptAdvanced: javascriptSnippet( + base, + key, + model, + "advanced", + autoSwitch, + ), }; } const KEY_PLACEHOLDER = "sk-unsloth-YOUR_KEY"; const MODEL_FALLBACK = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL"; - -// Default ON: when a tunnel exists, examples should show the public base_url. const USE_TUNNEL_KEY = "unsloth_api_use_tunnel"; function readUseTunnelPref(): boolean { @@ -319,11 +381,10 @@ function writeUseTunnelPref(value: boolean): void { try { window.localStorage.setItem(USE_TUNNEL_KEY, value ? "true" : "false"); } catch { - // Non-fatal: the toggle still applies for this session. + // Non-fatal } } -// Active local checkpoint as repo[:variant]; external/none falls back to a default. function useLoadedModelName(): string { const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint); const ggufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); @@ -338,7 +399,6 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } -// shiki highlighting via the app's shared code plugin + themes (same as chat). const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -352,7 +412,6 @@ function HighlightedCode({ code: string; language: string; }) { - // Fence so Streamdown's shiki plugin highlights it (no markdown inside a fence). const markdown = useMemo( () => `\`\`\`${language}\n${code}\n\`\`\``, [code, language], @@ -383,6 +442,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); + const [copiedAgent, setCopiedAgent] = useState(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); // null while loading; the same setting the General tab exposes (shared cache). const [autoSwitch, setAutoSwitch] = useState( @@ -390,7 +450,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); const [savingAutoSwitch, setSavingAutoSwitch] = useState(false); - // Tunnel may start after the first /api/health read; refresh so it surfaces here. useEffect(() => { void fetchDeviceType({ force: true }); }, []); @@ -410,10 +469,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { }, []); const model = useLoadedModelName(); - // Real key while revealed (before "Done"); otherwise a placeholder. const key = apiKey || KEY_PLACEHOLDER; - // Toggle on + tunnel up: public tunnel URL. Off: backend direct host:port - // (origin is only a last-resort fallback). const origin = typeof window !== "undefined" ? window.location.origin : ""; const base = useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); @@ -423,13 +479,20 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { () => buildSnippets(base, key, model, os, autoSwitchOn), [base, key, model, os, autoSwitchOn], ); + // Agent command must target the server the panel shows, not the :8888 default. + const agentCommand = useMemo( + () => buildAgentCommand(base, key, os), + [base, key, os], + ); const osAware = OS_AWARE[lang]; const shikiLang = CURL_TYPES.has(lang) ? os === "windows" ? "powershell" : "bash" - : "python"; + : JAVASCRIPT_TYPES.has(lang) + ? "javascript" + : "python"; const handleCopy = async () => { if (await copyToClipboard(snippets[lang])) { @@ -464,6 +527,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { } }; + const handleCopyAgent = async () => { + if (await copyToClipboard(agentCommand)) { + setCopiedAgent(true); + setTimeout(() => setCopiedAgent(false), 1800); + } + }; + return (

@@ -539,8 +609,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { )}

- {/* Always rendered (dimmed when off) so toggling never changes the - row height and shifts the code block below. */} - {/* key on the snippet so Streamdown remounts and re-highlights when - only a substring (e.g. the base URL) changes; its block memo - otherwise keeps the stale render. */}
+
+ + {t("settings.apiKeys.codingAgents")} + + + {t("settings.apiKeys.codingAgentsHint")} + +
+ + {agentCommand} + + +
+ + {t("settings.apiKeys.codingAgentsSwap")} + +
{t("settings.apiKeys.setupDocs")} {DOC_LINKS.map((link) => ( diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 0ee0dd0a7c..0d201edc0d 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -12,6 +12,7 @@ import { cn } from "@/lib/utils"; import { Cancel01Icon, CloudIcon, + CpuIcon, Globe02Icon, HelpCircleIcon, Message01Icon, @@ -33,6 +34,8 @@ import { ChatTab } from "./tabs/chat-tab"; import { ConnectionsTab } from "./tabs/connections-tab"; import { GeneralTab } from "./tabs/general-tab"; import { ProfileTab } from "./tabs/profile-tab"; +import { ResourcesTab } from "./tabs/resources-tab"; +import { FloatingMonitor } from "@/components/floating-monitor"; interface TabDef { id: SettingsTab; @@ -49,6 +52,11 @@ const TABS: TabDef[] = [ labelKey: "settings.tabs.appearance", icon: PaintBrush02Icon, }, + { + id: "resources", + labelKey: "settings.tabs.resources", + icon: CpuIcon, + }, { id: "chat", labelKey: "settings.tabs.chat", @@ -77,6 +85,8 @@ function renderTab(tab: SettingsTab) { return ; case "appearance": return ; + case "resources": + return ; case "chat": return ; case "connections": @@ -100,6 +110,7 @@ export function SettingsDialog() { general: null, profile: null, appearance: null, + resources: null, chat: null, connections: null, "api-keys": null, @@ -115,110 +126,113 @@ export function SettingsDialog() { }, [open, activeTab]); return ( - !o && closeDialog()}> - { - // Restore focus to the element that triggered openDialog(). Radix's - // FocusScope races our rAF-scheduled tab focus and loses the - // previous-focus reference, so restore it by hand. - if (opener && opener.isConnected) { - e.preventDefault(); - opener.focus({ preventScroll: true }); - } - }} - className={cn( - // Cap at 820px but shrink to the viewport so it doesn't clip on - // iPad-portrait widths (640-820px) where fixed `w-[820px]` overflows. - "settings-surface !max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden", - // Soft shadow, no outline ring. Pin --radius to the light value so - // corner rounding matches in dark mode. - "shadow-border rounded-xl ring-0 [--radius:1.1rem]", - "max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none", - )} - > - - {t("settings.dialog.title")} - - - {t("settings.dialog.description")} - -
- + {tab.badgeKey ? ( + + {t(tab.badgeKey)} + + ) : null} + + ); + })} + + -
- -
- {renderTab(activeTab)} -
-
-
- - +
+ +
+ {renderTab(activeTab)} +
+
+
+ + + + ); } diff --git a/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts new file mode 100644 index 0000000000..1af804a19e --- /dev/null +++ b/studio/frontend/src/features/settings/stores/monitor-overlay-store.ts @@ -0,0 +1,24 @@ +// 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 { create } from "zustand"; +import { persist } from "zustand/middleware"; + +interface MonitorOverlayState { + isOpen: boolean; + isMinimized: boolean; + setIsOpen: (open: boolean) => void; + toggleMinimized: () => void; +} + +export const useMonitorOverlayStore = create()( + persist( + (set) => ({ + isOpen: false, + isMinimized: false, + setIsOpen: (isOpen) => set({ isOpen }), + toggleMinimized: () => set((state) => ({ isMinimized: !state.isMinimized })), + }), + { name: "unsloth_monitor_overlay" } + ) +); \ No newline at end of file diff --git a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts index 7fab580f3c..234e92b3d0 100644 --- a/studio/frontend/src/features/settings/stores/settings-dialog-store.ts +++ b/studio/frontend/src/features/settings/stores/settings-dialog-store.ts @@ -7,6 +7,7 @@ export type SettingsTab = | "general" | "profile" | "appearance" + | "resources" | "chat" | "connections" | "api-keys" @@ -60,6 +61,7 @@ function loadInitialTab(): SettingsTab { "general", "profile", "appearance", + "resources", "chat", "connections", "api-keys", diff --git a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx index e83ecebfc8..f1e503f7a3 100644 --- a/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/api-keys-tab.tsx @@ -17,6 +17,7 @@ import { fetchApiKeys, revokeApiKey, type ApiKey } from "../api/api-keys"; import { ApiMonitorConsole } from "../components/api-monitor-console"; import { ApiKeyRow } from "../components/api-key-row"; import { CreateKeyForm } from "../components/create-key-form"; +import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; import { KeyRevealCard } from "../components/key-reveal-card"; import { UsageExamples } from "../components/usage-examples"; @@ -171,6 +172,8 @@ export function ApiKeysTab() { + + !o && setRevokeTarget(null)}> diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 247d5040fb..df684b7752 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -41,6 +41,13 @@ import { rotatePreviewLinks, updatePreviewSharing, } from "../api/preview-sharing"; +import { + type EmbeddingModelSettings, + EmbeddingModelVerificationError, + loadEmbeddingModelSettings, + resetEmbeddingModelSettings, + updateEmbeddingModelSettings, +} from "../api/embedding-model"; import { DEFAULT_UPLOAD_LIMIT_MB, type UploadLimitSettings, @@ -48,7 +55,7 @@ import { updateUploadLimitSettings, } from "../api/upload-limit"; import { ChangePasswordDialog } from "../components/change-password-dialog"; -import { ModelAutoSwitchSection } from "../components/model-auto-switch-section"; +import { EmbeddingModelCombobox } from "../components/embedding-model-combobox"; import { SettingsRow } from "../components/settings-row"; import { SettingsSection } from "../components/settings-section"; import { StudioVersionSection } from "../components/studio-version-section"; @@ -81,6 +88,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_load_on_selection", "unsloth_chat_expand_quantizations", "unsloth_chat_show_all_quantizations", + "unsloth_models_fit_on_device_only", // Chat presets "unsloth_chat_custom_presets", "unsloth_chat_active_preset", @@ -96,6 +104,7 @@ const PREFS_KEYS: string[] = [ "tour:studio:v1", // Update notifications "unsloth_show_llama_update_banner", + "unsloth_monitor_overlay", ]; // Set by resetAllPrefs so the unmount-commit effect skips writing back the @@ -163,6 +172,16 @@ export function GeneralTab() { const [revokePreviewOpen, setRevokePreviewOpen] = useState(false); const [isRevokingPreview, setIsRevokingPreview] = useState(false); const [modelsFolder, setModelsFolder] = useState(null); + const [embeddingModel, setEmbeddingModel] = + useState(null); + const [draftEmbeddingModel, setDraftEmbeddingModel] = useState(""); + const [embeddingModelError, setEmbeddingModelError] = useState( + null, + ); + // Set after a 409 (unverifiable model); offers "Save anyway". + const [embeddingModelNeedsForce, setEmbeddingModelNeedsForce] = + useState(false); + const [isSavingEmbeddingModel, setIsSavingEmbeddingModel] = useState(false); const draftRef = useRef(draftToken); useEffect(() => { @@ -257,6 +276,27 @@ export function GeneralTab() { }; }, [t]); + useEffect(() => { + let cancelled = false; + void loadEmbeddingModelSettings() + .then((settings) => { + if (cancelled) return; + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + }) + .catch((error) => { + if (cancelled) return; + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.loadError"), + ); + }); + return () => { + cancelled = true; + }; + }, [t]); + useEffect(() => { let cancelled = false; void loadModelsFolder() @@ -349,6 +389,58 @@ export function GeneralTab() { } }; + const saveEmbeddingModel = async (force: boolean) => { + const trimmed = draftEmbeddingModel.trim(); + if (!trimmed) { + setEmbeddingModelError(t("settings.general.rag.emptyError")); + return; + } + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + try { + const settings = await updateEmbeddingModelSettings(trimmed, { + hfToken: hfToken || undefined, + force, + }); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + setEmbeddingModelNeedsForce(false); + toast.success(t("settings.general.rag.saved"), { + description: t("settings.general.rag.reindexWarning"), + }); + } catch (error) { + if (error instanceof EmbeddingModelVerificationError) { + setEmbeddingModelNeedsForce(true); + } + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + + const resetEmbeddingModel = async () => { + setIsSavingEmbeddingModel(true); + setEmbeddingModelError(null); + setEmbeddingModelNeedsForce(false); + try { + const settings = await resetEmbeddingModelSettings(); + setEmbeddingModel(settings); + setDraftEmbeddingModel(settings.embeddingModel); + } catch (error) { + setEmbeddingModelError( + error instanceof Error + ? error.message + : t("settings.general.rag.saveError"), + ); + } finally { + setIsSavingEmbeddingModel(false); + } + }; + const saveUploadLimit = async () => { const parsed = Number(draftUploadLimit); if (!Number.isInteger(parsed)) { @@ -499,38 +591,6 @@ export function GeneralTab() { - - -
- void saveHelperPrecache(enabled)} - /> - {helperPrecache?.disabledByEnv ? ( - - {t("settings.general.helperLlm.disabledByEnv")} - - ) : helperPrecacheError ? ( - - {helperPrecacheError} - - ) : null} -
- - - - - @@ -567,6 +627,77 @@ export function GeneralTab() { + + +
+
+ { + setDraftEmbeddingModel(next); + setEmbeddingModelNeedsForce(false); + setEmbeddingModelError(null); + }} + accessToken={hfToken || undefined} + disabled={!embeddingModel} + placeholder={embeddingModel?.defaultEmbeddingModel ?? ""} + ariaLabel={t("settings.general.rag.embeddingModel")} + className="w-[220px]" + /> + +
+ {embeddingModelError ? ( + + {embeddingModelError} + + ) : null} +
+ {embeddingModelNeedsForce ? ( + + ) : null} + {embeddingModel?.isCustom ? ( + + ) : null} +
+ + {t("settings.general.rag.reindexWarning")} + +
+ + + )} + + +
+ void saveHelperPrecache(enabled)} + /> + {helperPrecache?.disabledByEnv ? ( + + {t("settings.general.helperLlm.disabledByEnv")} + + ) : helperPrecacheError ? ( + + {helperPrecacheError} + + ) : null} +
+ + + diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx new file mode 100644 index 0000000000..d5e19cc51c --- /dev/null +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -0,0 +1,477 @@ +// 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 { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { Switch } from "@/components/ui/switch"; +import { openModelsDir } from "@/features/native-intents"; +import { useSystemInfo, type GpuDevice } from "@/hooks/use-system"; +import { isTauri } from "@/lib/api-base"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { toast } from "@/lib/toast"; +import { cn } from "@/lib/utils"; +import { useT } from "@/i18n"; +import { useEffect, useMemo, useState } from "react"; +import { loadModelsFolder, type ModelsFolder } from "../api/models-folder"; +import { SettingsRow } from "../components/settings-row"; +import { SettingsSection } from "../components/settings-section"; +import { useMonitorOverlayStore } from "../stores/monitor-overlay-store"; +import { LayersIcon } from "lucide-react"; + +const POLL_MS = 3000; + +function isFiniteNumber(value: number | null | undefined): value is number { + return typeof value === "number" && Number.isFinite(value); +} + +function clampPercent(value: number | null | undefined): number { + if (!isFiniteNumber(value)) return 0; + return Math.max(0, Math.min(100, value)); +} + +function usageIndicatorClass(percent: number): string { + if (percent >= 90) return "bg-destructive"; + if (percent >= 70) return "bg-amber-500"; + return "bg-primary"; +} + +function usageTextClass(percent: number): string { + if (percent >= 90) return "text-destructive"; + if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + return "text-primary"; +} + +function formatGb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + const digits = safe >= 10 ? 1 : 2; + return `${safe.toFixed(digits)} GB`; +} + +function formatMb(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + return `${Math.round(safe).toLocaleString()} MB`; +} + +function formatPercent(value: number | null | undefined): string { + return `${Math.round(clampPercent(value))}%`; +} + +function formatFrequency(mhz: number | null | undefined): string | null { + if (!isFiniteNumber(mhz) || mhz <= 0) return null; + if (mhz >= 1000) return `${(mhz / 1000).toFixed(2)} GHz`; + return `${Math.round(mhz)} MHz`; +} + +function formatUptime(seconds: number | null | undefined): string { + if (!isFiniteNumber(seconds) || seconds <= 0) return "0m"; + const minutes = Math.floor(seconds / 60); + const hours = Math.floor(minutes / 60); + const days = Math.floor(hours / 24); + if (days > 0) return `${days}d ${hours % 24}h`; + if (hours > 0) return `${hours}h ${minutes % 60}m`; + return `${Math.max(1, minutes)}m`; +} + +function MetricTile({ + label, + value, + detail, + percent, +}: { + label: string; + value: string; + detail: string; + percent: number; +}) { + const safePercent = clampPercent(percent); + return ( +
+
+ + {label} + + + {formatPercent(safePercent)} + +
+
+
+ {value} +
+
+ {detail} +
+
+ +
+ ); +} + +function InfoRow({ + label, + value, + detail, +}: { + label: string; + value: string; + detail?: string; +}) { + return ( +
+ + {label} + + + {detail ? `${value} (${detail})` : value} + +
+ ); +} + +function deviceOrdinal(device: GpuDevice): number | undefined { + return device.visible_ordinal ?? device.index; +} + +export function ResourcesTab() { + const t = useT(); + const [liveUpdates, setLiveUpdates] = useState(true); + const { isOpen, setIsOpen } = useMonitorOverlayStore(); + const systemInfo = useSystemInfo({ + enabled: liveUpdates, + pollMs: liveUpdates ? POLL_MS : undefined, + }); + const [modelsFolder, setModelsFolder] = useState(null); + const [modelsFolderLoaded, setModelsFolderLoaded] = useState(false); + + useEffect(() => { + let cancelled = false; + void loadModelsFolder() + .then((folder) => { + if (cancelled) return; + setModelsFolder(folder); + setModelsFolderLoaded(true); + }) + .catch(() => { + if (cancelled) return; + setModelsFolderLoaded(true); + }); + return () => { + cancelled = true; + }; + }, []); + + const metrics = useMemo(() => { + const devices = systemInfo.gpu?.devices ?? []; + const ramTotal = systemInfo.memory?.total_gb ?? 0; + const ramAvailable = systemInfo.memory?.available_gb ?? 0; + const ramUsed = Math.max(0, ramTotal - ramAvailable); + const diskTotal = systemInfo.disk?.total_gb ?? 0; + const diskFree = systemInfo.disk?.free_gb ?? 0; + const diskUsed = Math.max(0, diskTotal - diskFree); + const vramTotal = devices.reduce( + (sum, device) => sum + (device.memory_total_gb ?? 0), + 0, + ); + const vramUsed = devices.reduce( + (sum, device) => sum + (device.vram_used_gb ?? 0), + 0, + ); + const vramFree = devices.reduce( + (sum, device) => + sum + + (device.vram_free_gb ?? + Math.max(0, (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0))), + 0, + ); + const vramPercent = vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0; + + return { + devices, + ramTotal, + ramUsed, + diskTotal, + diskFree, + diskUsed, + vramTotal, + vramUsed, + vramFree, + vramPercent, + }; + }, [systemInfo]); + + const handleModelsFolder = async () => { + const folder = modelsFolder; + if (!folder) return; + if (isTauri) { + try { + await openModelsDir(folder.path); + } catch (error) { + toast.error(t("settings.resources.storage.openError"), { + description: error instanceof Error ? error.message : undefined, + }); + } + return; + } + if (await copyToClipboard(folder.path)) { + toast.success(t("settings.resources.storage.copied")); + } else { + toast.error(t("settings.resources.storage.copyError")); + } + }; + + const cpuCoresLabel = + systemInfo.cpu?.logical_count && systemInfo.cpu?.physical_count + ? t("settings.resources.liveMonitor.cpuCores", { + logical: systemInfo.cpu.logical_count, + physical: systemInfo.cpu.physical_count, + }) + : t("settings.resources.environment.unknown"); + const cpuFrequencyLabel = formatFrequency(systemInfo.cpu?.frequency_mhz); + const hasGpu = + (systemInfo.gpu?.available ?? false) && metrics.devices.length > 0; + const backendLabel = ( + systemInfo.gpu?.backend ?? systemInfo.device_backend ?? "cpu" + ).toUpperCase(); + const modelsFolderPath = modelsFolder + ? modelsFolder.path + : modelsFolderLoaded + ? t("settings.resources.environment.unknown") + : t("common.loading"); + + return ( +
+
+
+

+ {t("settings.resources.title")} +

+

+ {t("settings.resources.description")} +

+
+
+ + +
+ {t("settings.resources.liveUpdates")} + +
+
+
+ + +
+ + + + +
+ + + + {hasGpu ? ( + metrics.devices.map((device, index) => { + const ordinal = deviceOrdinal(device); + const total = device.memory_total_gb ?? 0; + const used = device.vram_used_gb ?? 0; + const free = device.vram_free_gb ?? Math.max(0, total - used); + const percent = + device.vram_utilization_pct ?? + (total > 0 ? (used / total) * 100 : null); + const safePercent = clampPercent(percent); + return ( +
+
+
+
+ {device.name ?? + t("settings.resources.gpu.unknownDevice")} +
+
+ {ordinal === undefined + ? backendLabel + : `${t("settings.resources.gpu.deviceWithIndex", { + index: ordinal, + })}, ${backendLabel}`} +
+
+
+ + {formatPercent(safePercent)}{" "} + {t("settings.resources.gpu.vramUtilization")} + +
+
+
+ + {t("settings.resources.gpu.used", { + value: formatGb(used), + })} + + + {t("settings.resources.gpu.free", { + value: formatGb(free), + })} + + + {t("settings.resources.gpu.total", { + value: formatGb(total), + })} + +
+ +
+ ); + }) + ) : ( +
+ {t("settings.resources.gpu.noGpu")} +
+ )} + + + + + +
+ + {modelsFolderPath} + + +
+ + + + + + + + + + + +
+ ); +} diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 6aaca86e0b..abab35db93 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -29,6 +29,7 @@ import { import { getTrainingMethodLabel } from "@/features/training/lib/training-methods"; import type { TrainingViewData } from "@/features/training"; import { useGpuUtilization } from "@/hooks"; +import type { GpuUtilization } from "@/hooks/use-gpu-utilization"; import { cn } from "@/lib/utils"; import { ChartAverageIcon, @@ -42,7 +43,7 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Link, useNavigate } from "@tanstack/react-router"; -import { type ReactElement, type ReactNode, useState } from "react"; +import { type ReactElement, type ReactNode, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { ChartSettingsSheet } from "./charts/chart-settings-sheet"; import { @@ -123,18 +124,17 @@ export function ProgressSection({ const [stopDialogOpen, setStopDialogOpen] = useState(false); const [stopRequestedLocal, setStopRequestedLocal] = useState(false); - // Auto-resets when training stops; no useEffect needed const stopRequested = data.isTrainingRunning && stopRequestedLocal; const pct = data.totalSteps > 0 ? Math.min( - 100, - Math.max( - 0, - Math.round((data.currentStep / data.totalSteps) * 100), - ), - ) + 100, + Math.max( + 0, + Math.round((data.currentStep / data.totalSteps) * 100), + ), + ) : Math.round(data.progressPercent); const elapsed = data.elapsedSeconds; @@ -214,16 +214,16 @@ export function ProgressSection({ }, ...(data.trainingMethod !== "full" ? [ - { - section: "LoRA", - rows: [ - configRow(t("studio.progress.rank"), cfgLoraRank), - configRow(t("studio.progress.alpha"), cfgLoraAlpha), - configRow(t("studio.progress.dropout"), cfgLoraDropout), - configRow(t("studio.progress.variant"), cfgLoraVariant), - ], - }, - ] + { + section: "LoRA", + rows: [ + configRow(t("studio.progress.rank"), cfgLoraRank), + configRow(t("studio.progress.alpha"), cfgLoraAlpha), + configRow(t("studio.progress.dropout"), cfgLoraDropout), + configRow(t("studio.progress.variant"), cfgLoraVariant), + ], + }, + ] : []), ]; @@ -350,8 +350,8 @@ export function ProgressSection({ {stepsPerSecond == null ? t("studio.progress.noStepsPerSecond") : t("studio.progress.stepsPerSecond", { - value: stepsPerSecond.toFixed(2), - })} + value: stepsPerSecond.toFixed(2), + })} {data.currentNumTokens != null && ( {t("studio.progress.tokens", { value: data.currentNumTokens })} @@ -373,14 +373,50 @@ function LiveGpuPanel({ isTrainingRunning: boolean; }): ReactElement { const t = useT(); - const gpu = useGpuUtilization(isTrainingRunning); + const [selectedGpu, setSelectedGpu] = useState(0); + const gpuData = useGpuUtilization(isTrainingRunning); + const gpus: GpuUtilization[] = + Array.isArray(gpuData?.devices) && gpuData.devices.length > 0 + ? gpuData.devices + : gpuData && Object.keys(gpuData).length > 0 + ? [gpuData] + : []; + + useEffect(() => { + if (selectedGpu > 0 && selectedGpu >= gpus.length) { + setSelectedGpu(0); + } + }, [gpus.length, selectedGpu]); + + const gpuCount = gpus.length; + const currentGpu: Partial = gpus[selectedGpu] || gpus[0] || {}; return (
-
-

- {t("studio.progress.gpuMonitor")} -

+
+
+

+ {t("studio.progress.gpuMonitor")} +

+ {gpuCount > 1 && ( + + )} +
{t("studio.progress.live")} @@ -388,51 +424,44 @@ function LiveGpuPanel({
- } + icon={} value={ - gpu.gpu_utilization_pct != null - ? `${gpu.gpu_utilization_pct}%` + currentGpu.gpu_utilization_pct != null + ? `${currentGpu.gpu_utilization_pct}%` : "--" } - pct={gpu.gpu_utilization_pct ?? 0} + pct={currentGpu.gpu_utilization_pct ?? 0} /> - } + icon={} value={ - gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--" + currentGpu.temperature_c != null ? `${currentGpu.temperature_c}°C` : "--" } - pct={gpu.temperature_c ?? 0} + pct={currentGpu.temperature_c ?? 0} max={100} /> } value={ - gpu.vram_used_gb != null && gpu.vram_total_gb != null - ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB` + currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null + ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB` : "--" } - pct={gpu.vram_utilization_pct ?? 0} + pct={currentGpu.vram_utilization_pct ?? 0} /> } value={ - gpu.power_draw_w != null - ? gpu.power_limit_w != null - ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W` - : `${gpu.power_draw_w} W` + currentGpu.power_draw_w != null + ? currentGpu.power_limit_w != null + ? `${currentGpu.power_draw_w} / ${currentGpu.power_limit_w} W` + : `${currentGpu.power_draw_w} W` : "--" } - pct={gpu.power_utilization_pct ?? 0} + pct={currentGpu.power_utilization_pct ?? 0} />
@@ -560,7 +589,10 @@ function TrainingHeaderActions({ {stopRequested ? t("studio.training.stopping") : t("studio.training.stopAction")} - + {t("studio.training.stopTitle")} diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index d5c923d57d..5b2fa43ae7 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -1,6 +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 + export { useDebouncedValue } from "./use-debounced-value"; export { useGpuInfo } from "./use-gpu-info"; export { useGpuUtilization } from "./use-gpu-utilization"; @@ -9,3 +10,4 @@ export { useHfDatasetSplits } from "./use-hf-dataset-splits"; export { useHfTokenValidation } from "./use-hf-token-validation"; export { useTauriBackend } from "./use-tauri-backend"; export { useCollapseScrollLock } from "./use-collapse-scroll-lock"; +export { useSystemInfo } from "./use-system"; diff --git a/studio/frontend/src/hooks/use-gpu-info.ts b/studio/frontend/src/hooks/use-gpu-info.ts index eb4d89abd8..1e313acdf3 100644 --- a/studio/frontend/src/hooks/use-gpu-info.ts +++ b/studio/frontend/src/hooks/use-gpu-info.ts @@ -3,19 +3,26 @@ import { authFetch } from "@/features/auth"; import { useEffect, useState } from "react"; +import type { SystemInfoResponse } from "./use-system"; export interface GpuInfo { available: boolean; name: string; memoryTotalGb: number; + cpuCore: number; + cpuThread: number; systemRamAvailableGb: number; + systemRamTotalGb: number } const DEFAULT_GPU: GpuInfo = { available: false, name: "Unknown", memoryTotalGb: 0, + cpuCore: 0, + cpuThread: 0, systemRamAvailableGb: 0, + systemRamTotalGb: 0 }; // Module-level cache so multiple components share one fetch. @@ -30,24 +37,30 @@ async function fetchGpuOnce(): Promise { try { const res = await authFetch("/api/system"); if (!res.ok) throw new Error(`HTTP ${res.status}`); - const data = await res.json(); - const ramAvailableGb = data?.memory?.available_gb ?? 0; + + const data = await res.json() as SystemInfoResponse; const gpuData = data?.gpu; - if (!gpuData?.available || !gpuData.devices?.length) { - // No discrete GPU (e.g. Mac): still surface system RAM so memory math - // (unified memory) has a budget to work with. - const info: GpuInfo = { ...DEFAULT_GPU, systemRamAvailableGb: ramAvailableGb }; - cachedGpu = info; - return info; - } - const devices = gpuData.devices as Array<{ name?: string; memory_total_gb?: number }>; - const totalGb = devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0); - const info: GpuInfo = { - available: true, - name: devices[0]?.name ?? "Unknown", - memoryTotalGb: totalGb, - systemRamAvailableGb: ramAvailableGb, + + // CPU/RAM exist even on hosts without a GPU, so populate them on every path. + // No discrete GPU (e.g. Mac): still surface system RAM so memory math + // (unified memory) has a budget to work with. + const base = { + cpuCore: data?.cpu?.physical_count ?? 0, + cpuThread: data?.cpu?.logical_count ?? 0, + systemRamAvailableGb: data?.memory?.available_gb ?? 0, + systemRamTotalGb: data?.memory?.total_gb ?? 0, }; + + const devices = gpuData?.devices ?? []; + const info: GpuInfo = + gpuData?.available && devices.length + ? { + ...base, + available: true, + name: devices[0]?.name ?? "Unknown", + memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0), + } + : { ...DEFAULT_GPU, ...base }; cachedGpu = info; return info; } catch { @@ -78,4 +91,4 @@ export function useGpuInfo(): GpuInfo { }, []); return gpu; -} +} \ No newline at end of file diff --git a/studio/frontend/src/hooks/use-gpu-utilization.ts b/studio/frontend/src/hooks/use-gpu-utilization.ts index 1a9f6102dd..be16647a4d 100644 --- a/studio/frontend/src/hooks/use-gpu-utilization.ts +++ b/studio/frontend/src/hooks/use-gpu-utilization.ts @@ -7,6 +7,9 @@ import { useEffect, useRef, useState } from "react"; export interface GpuUtilization { available: boolean; backend: string | null; + devices?: GpuUtilization[]; + index?: number; + visible_ordinal?: number; gpu_utilization_pct: number | null; temperature_c: number | null; vram_used_gb: number | null; @@ -57,11 +60,10 @@ export function useGpuUtilization( const json = (await res.json()) as GpuUtilization; if (!cancelled) setData(json); } catch { - // Silently ignore — next poll will retry + // Retry on the next poll. } } - // Fetch immediately, then set up interval void poll(); timerRef.current = setInterval(() => void poll(), intervalMs); diff --git a/studio/frontend/src/hooks/use-hardware-info.ts b/studio/frontend/src/hooks/use-hardware-info.ts index 36e3c532f3..4d63d4d6af 100644 --- a/studio/frontend/src/hooks/use-hardware-info.ts +++ b/studio/frontend/src/hooks/use-hardware-info.ts @@ -25,6 +25,13 @@ export interface HardwareInfo { transformers: string | null; unsloth: string | null; llamaCpp: string | null; + // Whether export can run here (true only on a supported accelerator), with a torch-aware + // reason. `null` until the authoritative response lands, so callers don't briefly enable + // export; `loaded` flips true once a real (non-error) response arrives. + exportSupported: boolean | null; + exportUnsupportedReason: string | null; + exportUnsupportedMessage: string | null; + loaded: boolean; } const DEFAULT: HardwareInfo = { @@ -38,6 +45,10 @@ const DEFAULT: HardwareInfo = { transformers: null, unsloth: null, llamaCpp: null, + exportSupported: null, + exportUnsupportedReason: null, + exportUnsupportedMessage: null, + loaded: false, }; // Module-level cache so multiple components share one fetch. @@ -87,6 +98,10 @@ async function fetchOnce(): Promise { transformers: data?.versions?.transformers ?? null, unsloth: data?.versions?.unsloth ?? null, llamaCpp: data?.llama_cpp ?? null, + exportSupported: data?.export_supported ?? null, + exportUnsupportedReason: data?.export_unsupported_reason ?? null, + exportUnsupportedMessage: data?.export_unsupported_message ?? null, + loaded: true, }; if (generation === cacheGeneration) { cached = info; diff --git a/studio/frontend/src/hooks/use-system.ts b/studio/frontend/src/hooks/use-system.ts new file mode 100644 index 0000000000..a135cce86e --- /dev/null +++ b/studio/frontend/src/hooks/use-system.ts @@ -0,0 +1,130 @@ +// 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 { authFetch } from "@/features/auth"; +import { useEffect, useState } from "react"; + +export interface GpuDevice { + index?: number; + index_kind?: string; + visible_ordinal?: number; + name?: string; + memory_total_gb?: number; + vram_used_gb?: number; + vram_free_gb?: number; + vram_utilization_pct?: number | null; +} + +export interface SystemInfoResponse { + platform: string; + python_version: string; + device_backend: "cuda" | "rocm" | "cpu" | "mlx" | "xpu"; + uptime_seconds: number | null; + cpu: { + logical_count: number; + physical_count: number; + usage_percent: number; + frequency_mhz: number | null; + }; + memory: { + total_gb: number; + available_gb: number; + percent_used: number; + process_used_mb: number; + }; + disk: { + total_gb: number; + free_gb: number; + percent_used: number; + }; + gpu: { + available: boolean; + backend?: string; + backend_cuda_visible_devices?: string | null; + parent_visible_gpu_ids?: number[]; + index_kind?: string; + devices: GpuDevice[]; + }; + ml_packages: { + torch?: string; + transformers?: string; + }; +} + +let cachedSystem: SystemInfoResponse | null = null; +let systemFetchPromise: Promise | null = null; + +const DEFAULT_SYSTEM: SystemInfoResponse = { + platform: "Unknown", + python_version: "Unknown", + device_backend: "cpu", + uptime_seconds: 0, + cpu: { logical_count: 0, physical_count: 0, usage_percent: 0, frequency_mhz: null }, + memory: { total_gb: 0, available_gb: 0, percent_used: 0, process_used_mb: 0 }, + disk: { total_gb: 0, free_gb: 0, percent_used: 0 }, + gpu: { available: false, devices: [] }, + ml_packages: {} +}; + +async function fetchSystemOnce({ + force = false, +}: { force?: boolean } = {}): Promise { + if (systemFetchPromise) return systemFetchPromise; + if (!force && cachedSystem) return cachedSystem; + + systemFetchPromise = (async () => { + try { + const res = await authFetch("/api/system"); + if (!res.ok) throw new Error(`HTTP ${res.status}`); + const data = await res.json(); + + cachedSystem = data as SystemInfoResponse; + return cachedSystem; + } catch { + cachedSystem = null; + return DEFAULT_SYSTEM; + } finally { + systemFetchPromise = null; + } + })(); + + return systemFetchPromise; +} + +interface UseSystemInfoOptions { + pollMs?: number; + enabled?: boolean; +} + +export function useSystemInfo({ + pollMs, + enabled = true, +}: UseSystemInfoOptions = {}): SystemInfoResponse { + const [systemInfo, setSystemInfo] = useState(cachedSystem ?? DEFAULT_SYSTEM); + + useEffect(() => { + if (!enabled) return; + + let cancelled = false; + let timeoutId: number | null = null; + + const update = (force: boolean) => { + void fetchSystemOnce({ force }) + .then((info) => { + if (!cancelled) setSystemInfo(info); + }) + .finally(() => { + if (cancelled || !pollMs) return; + timeoutId = window.setTimeout(() => update(true), pollMs); + }); + }; + + update(Boolean(pollMs)); + return () => { + cancelled = true; + if (timeoutId !== null) window.clearTimeout(timeoutId); + }; + }, [enabled, pollMs]); + + return systemInfo; +} diff --git a/studio/frontend/src/i18n/AGENTS.md b/studio/frontend/src/i18n/AGENTS.md index 42d964acca..35c025cd0b 100644 --- a/studio/frontend/src/i18n/AGENTS.md +++ b/studio/frontend/src/i18n/AGENTS.md @@ -2,10 +2,11 @@ - `locales/en.ts` is the complete baseline message file. - Non-English locale files may be partial. Missing keys must fall back to English at runtime. -- Use BCP 47 locale tags for new languages, for example `zh-CN`, `ja-JP`, and `ko-KR`. +- Use BCP 47 locale tags for new languages, for example `zh-CN`, `pt-BR`, `ja-JP`, and `ko-KR`. - Do not change fallback logic to hide missing translations. - Do not add automatic DOM translation, MutationObserver text replacement, or runtime guess-based translation. - Preserve interpolation variables exactly, for example `{count}`, `{model}`, and `{provider}`. - Keep product and technical names unchanged unless there is an established localized name, for example `Unsloth Studio`, `LoRA`, `GGUF`, and `Hugging Face`. - Keep translation changes small and reviewable. Prefer separate commits for runtime changes, UI migration, and locale text. - When adding user-facing Studio UI text, add the English message key first and add non-English overrides only when the translation is clear. +- Run `npx tsx src/i18n/check-parity.ts` before committing to ensure there are no shape mismatches or placeholder discrepancies in the non-English overlays. \ No newline at end of file diff --git a/studio/frontend/src/i18n/check-parity.ts b/studio/frontend/src/i18n/check-parity.ts index 8c027f9ce3..e2b66f5cdc 100644 --- a/studio/frontend/src/i18n/check-parity.ts +++ b/studio/frontend/src/i18n/check-parity.ts @@ -3,13 +3,14 @@ // Parity check between en.ts and every non-English locale. // - Locale files may be partial; missing keys must fall back to English. -// - All zh-CN keys must exist in en (no extras). +// - All non-English keys must exist in en (no extras). // - Placeholder set must match per leaf between en and the overlay. // // Run: npx tsx src/i18n/check-parity.ts import { en } from "./locales/en.ts"; import { zhCN } from "./locales/zh-CN.ts"; +import { ptBR } from "./locales/pt-br.ts"; import { ja } from "./locales/ja.ts"; type Tree = { readonly [k: string]: string | Tree }; @@ -90,6 +91,7 @@ function checkExtras( const overlays: Record = { "zh-CN": zhCN as unknown as Tree, + "pt-BR": ptBR as unknown as Tree, "ja": ja as unknown as Tree, }; let anyError = false; @@ -112,4 +114,4 @@ for (const [locale, overlay] of Object.entries(overlays)) { } if (anyError) process.exit(1); -console.log("\nAll locale overlays pass parity."); +console.log("\nAll locale overlays pass parity."); \ No newline at end of file diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 8507fe5174..d597a919d9 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -93,6 +93,7 @@ export const en = { general: "General", profile: "Profile", appearance: "Appearance", + resources: "System", chat: "Chat", connections: "Connections", apiKeys: "API", @@ -194,6 +195,20 @@ export const en = { maxUploadSize: "Training dataset upload cap", maxUploadSizeDescription: "Default is {defaultSize} MB.", }, + rag: { + sectionTitle: "Documents & RAG", + embeddingModel: "Embedding model", + embeddingModelDescription: + "Hugging Face model or local path used to index and search your documents. Default is {defaultModel}.", + reindexWarning: + "Only affects newly indexed documents. Re-upload existing ones after changing the model.", + emptyError: "Enter a Hugging Face model id or local path.", + loadError: "Failed to load the embedding model setting.", + saveError: "Failed to save the embedding model.", + saved: "Embedding model saved.", + saveAnyway: "Save anyway", + resetAction: "Reset to default", + }, storage: { sectionTitle: "Storage", modelsFolder: "Models folder", @@ -262,6 +277,58 @@ export const en = { "Keep the sidebar expanded instead of collapsing to icons.", }, }, + resources: { + title: "System", + description: "Monitor this Studio server's hardware and storage.", + liveUpdates: "Live updates", + floatingWindow: "Floating window", + disableOverlay: "Disable overlay", + liveMonitor: { + title: "Live monitor", + cpu: "CPU", + ram: "RAM", + disk: "Disk", + vram: "VRAM", + cpuCores: "{logical} logical / {physical} physical cores", + currentLoad: "Current load", + free: "{value} free", + noGpu: "No visible GPU", + }, + gpu: { + title: "GPU devices", + noGpu: "No visible GPU detected. CPU-only resources are shown above.", + unknownDevice: "Unknown GPU", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} used", + free: "{value} free", + total: "{value} total", + }, + storage: { + title: "Storage", + systemDisk: "System disk", + diskUsage: "{used} used / {total}", + diskFree: "{free} free", + modelsFolder: "Models folder", + modelsFolderDescription: "Where downloaded models are stored.", + openAction: "Open", + copyAction: "Copy path", + copied: "Path copied", + openError: "Couldn't open the folder", + copyError: "Couldn't copy the path", + }, + environment: { + title: "Environment", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Uptime", + processMemory: "Process memory", + notInstalled: "Not installed", + unknown: "Unknown", + }, + }, chat: { title: "Chat", description: "Manage chat history stored on this device.", @@ -360,8 +427,10 @@ export const en = { usageTools: "Tools", exampleCurlTools: "curl + tools", examplePythonTools: "Python + tools", + exampleJavaScriptTools: "JavaScript + tools", exampleCurlAdvanced: "curl + advanced", examplePythonAdvanced: "Python + advanced", + exampleJavaScriptAdvanced: "JavaScript + advanced", osUnix: "Linux / macOS / WSL", osWindows: "Windows", secureHttps: "Secure HTTPS", @@ -372,6 +441,10 @@ export const en = { copy: "Copy", copied: "Copied", setupDocs: "Setup docs:", + codingAgents: "Coding agents", + codingAgentsHint: + "Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.", + codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts new file mode 100644 index 0000000000..c261e4ed0c --- /dev/null +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -0,0 +1,934 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export const ptBR = { + common: { + cancel: "Cancelar", + close: "Fechar", + delete: "Excluir", + done: "Concluído", + error: "Erro", + export: "Exportar", + help: "Ajuda", + loading: "Carregando...", + new: "Novo", + rename: "Renomear", + save: "Salvar", + saving: "Salvando...", + search: "Buscar", + shutdown: "Desligar", + }, + shell: { + beta: "BETA", + brand: "unsloth", + product: "Unsloth Studio", + accountMenu: "Menu de conta {name}", + updateAvailable: "Atualização disponível", + aria: { + home: "Início do Unsloth", + closeSidebar: "Fechar barra lateral", + openSidebar: "Abrir barra lateral", + chatOptions: "Opções de chat", + runOptions: "Opções de execução", + }, + navigation: { + newChat: "Novo Chat", + returnToChat: "Retornar ao Chat", + compare: "Comparar", + search: "Buscar", + hub: "Hub", + train: "Treinar", + recipes: "Receitas", + export: "Exportar", + recents: "Recentes", + settings: "Configurações", + api: "API", + lightMode: "Modo Claro", + darkMode: "Modo Escuro", + guidedTour: "Tour Guiado", + help: "Ajuda", + logOut: "Sair", + shutdown: "Desligar", + }, + notFound: { + title: "Página não encontrada", + description: "{path} não existe.", + backToChat: "Voltar para o chat", + }, + dialog: { + deleteChat: { + title: "Excluir chat", + description: 'Tem certeza de que deseja excluir este chat "{name}"?', + }, + deleteRun: { + title: "Excluir execução de treino", + description: 'Tem certeza de que deseja excluir esta execução "{name}"?', + }, + renameChat: { + title: "Renomear chat", + placeholder: "Título do chat", + }, + renameRun: { + title: "Renomear execução", + placeholder: "Nome da execução", + }, + }, + toast: { + cannotDeleteRunningRun: "Não é possível excluir uma execução de treino em andamento", + failedToDeleteChat: "Falha ao excluir o chat", + failedToDeleteRun: "Falha ao excluir a execução", + failedToRenameChat: "Falha ao renomear o chat", + failedToRenameRun: "Falha ao renomear a execução", + }, + }, + settings: { + title: "Configurações", + dialog: { + title: "Configurações", + description: "Gerencie suas preferências do Unsloth.", + closeAriaLabel: "Fechar configurações", + }, + tabs: { + general: "Geral", + profile: "Perfil", + appearance: "Aparência", + resources: "Sistema", + chat: "Chat", + connections: "Conexões", + apiKeys: "API", + about: "Sobre", + }, + general: { + title: "Geral", + description: "Preferências globais do Unsloth.", + account: "Conta", + huggingFaceToken: "Token do Hugging Face", + huggingFaceTokenDescription: + "Usado para carregar modelos restritos e enviar artefatos.", + tokenSaved: "Token salvo", + hideToken: "Ocultar token", + showToken: "Mostrar token", + password: "Senha", + passwordDescription: "Altere a senha desta conta do Studio.", + passwordDialog: { + trigger: "Alterar senha", + title: "Alterar senha", + description: + "Insira sua senha atual e escolha uma nova (no mínimo {minLength} caracteres).", + currentPassword: "Senha atual", + newPassword: "Nova senha", + confirmPassword: "Confirmar nova senha", + currentTooShort: + "A senha atual deve ter no mínimo {minLength} caracteres.", + newTooShort: "A nova senha deve ter no mínimo {minLength} caracteres.", + mismatch: "As senhas não coincidem.", + samePassword: + "A nova senha deve ser diferente da senha atual.", + update: "Atualizar senha", + updating: "Atualizando...", + updated: "Senha atualizada.", + updateFailed: "Falha ao atualizar a senha.", + }, + chatDefaults: "Padrões do chat", + autoTitleNewChats: "Gerar título automático para novos chats", + autoTitleNewChatsDescription: + "Gera um título curto a partir da primeira mensagem.", + helperLlm: { + sectionTitle: "LLM Auxiliar", + preloadOnStartup: "Pré-carregar LLM Auxiliar na inicialização", + preloadOnStartupDescription: + "Baixa o modelo auxiliar do Assistente de IA em segundo plano ao iniciar. Desativado por padrão; o Assistente de IA ainda pode buscá-lo sob demanda.", + disabledByEnv: + "Desativado por UNSLOTH_HELPER_MODEL_DISABLE no ambiente de backend.", + loadError: "Falha ao carregar as configurações do LLM Auxiliar.", + saveError: "Falha ao salvar as configurações do LLM Auxiliar.", + }, + notifications: { + sectionTitle: "Notificações", + showLlamaUpdates: "Notificações de atualização do llama.cpp", + showLlamaUpdatesDescription: + "Notifica quando uma nova versão do llama.cpp estiver disponível. Desative se você apenas realiza treinos.", + }, + gettingStarted: "Primeiros passos", + startOnboarding: "Iniciar integração", + startOnboardingDescription: + "Reabre o assistente de configuração sem alterar sua conta.", + startOnboardingAction: "Iniciar integração", + uploads: { + sectionTitle: "Uploads", + maxUploadSize: "Limite de upload do dataset de treino", + maxUploadSizeDescription: + "O padrão é {defaultSize} MB.", + }, + storage: { + sectionTitle: "Armazenamento", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: + "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + resetPreferences: { + sectionTitle: "Zona de perigo", + label: "Redefinir todas as preferências locais", + description: + "Limpa apenas as preferências locais. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + action: "Redefinir preferências", + confirmTitle: "Redefinir todas as preferências locais?", + confirmDescription: + "Limpa as preferências locais e recarrega o Unsloth. Chats, acesso à API e configurações salvas no banco de dados são mantidos.", + confirmAction: "Redefinir e recarregar", + }, + }, + profile: { + title: "Perfil", + description: "Como seu perfil aparece no Unsloth.", + changePicture: "Alterar foto de perfil", + displayName: "Nome de exibição", + nickname: "Como o Unsloth deve chamar você?", + nicknamePlaceholder: "Apelido", + nicknameSaved: "Nome preferido salvo", + avatarShape: "Formato da foto de perfil", + avatarShapeCircle: "Círculo", + avatarShapeRounded: "Arredondado", + chooseSloth: "Ou escolha uma preguiça", + nameSaved: "Nome de perfil salvo", + namePersistErrorTitle: "Não foi possível persistir o nome de perfil", + namePersistErrorDescription: + "Nome atualizado para esta sessão, mas pode não persistir após recarregar.", + photoUpdated: "Foto de perfil atualizada", + photoPersistErrorTitle: "Não foi possível persistir a foto de perfil", + photoPersistErrorDescription: + "Foto atualizada para esta sessão, mas pode não persistir após recarregar.", + photoUpdateErrorTitle: "Não foi possível atualizar a foto de perfil", + imageUseError: "Não foi possível usar esta imagem.", + }, + appearance: { + title: "Aparência", + description: "Como o Unsloth Studio se parece neste dispositivo.", + theme: { + title: "Tema", + label: "Esquema de cores", + description: "Claro, escuro ou seguir o sistema.", + system: "Sistema", + light: "Claro", + dark: "Escuro", + }, + language: { + title: "Idioma", + label: "Idioma de exibição", + description: "O idioma utilizado pelo Unsloth.", + }, + layout: { + title: "Layout", + compactSidebar: "Fixar barra lateral por padrão", + compactSidebarDescription: + "Mantém a barra lateral expandida em vez de recolhê-la em ícones.", + }, + }, + resources: { + title: "Sistema", + description: "Monitore o hardware e o armazenamento deste servidor Studio.", + liveUpdates: "Atualizações ao vivo", + floatingWindow: "Janela flutuante", + disableOverlay: "Desativar sobreposição", + liveMonitor: { + title: "Monitor ao vivo", + cpu: "CPU", + ram: "RAM", + disk: "Disco", + vram: "VRAM", + cpuCores: "{logical} lógicos / {physical} físicos", + currentLoad: "Carga atual", + free: "{value} livres", + noGpu: "Nenhuma GPU visível", + }, + gpu: { + title: "Dispositivos GPU", + noGpu: "Nenhuma GPU visível detectada. Os recursos somente CPU aparecem acima.", + unknownDevice: "GPU desconhecida", + deviceWithIndex: "GPU {index}", + vramUtilization: "VRAM", + used: "{value} usados", + free: "{value} livres", + total: "{value} total", + }, + storage: { + title: "Armazenamento", + systemDisk: "Disco do sistema", + diskUsage: "{used} usados / {total}", + diskFree: "{free} livres", + modelsFolder: "Pasta de modelos", + modelsFolderDescription: "Onde os modelos baixados são armazenados.", + openAction: "Abrir", + copyAction: "Copiar caminho", + copied: "Caminho copiado", + openError: "Não foi possível abrir a pasta", + copyError: "Não foi possível copiar o caminho", + }, + environment: { + title: "Ambiente", + backend: "Backend", + python: "Python", + torch: "Torch", + transformers: "Transformers", + uptime: "Tempo ativo", + processMemory: "Memória do processo", + notInstalled: "Não instalado", + unknown: "Desconhecido", + }, + }, + chat: { + title: "Chat", + description: "Gerencie o histórico de chat armazenado neste dispositivo.", + modelDisclaimer: "Mostrar aviso do modelo", + modelDisclaimerDescription: + 'Mostra "LLMs podem cometer erros" abaixo da caixa de chat.', + artifacts: { + title: "Canvas", + collapseHtmlBlocks: "Recolher blocos HTML", + collapseHtmlBlocksDescription: + "O modo Canvas recolhe o HTML completo automaticamente. Ative isso para também recolher documentos HTML delimitados quando o Canvas estiver desativado.", + allowNetworkAccess: "Permitir acesso à rede no canvas", + allowNetworkAccessDescription: + "Permite que as pré-visualizações do canvas carreguem scripts, estilos, fontes, mídia e recursos de rede de CDNs. Mantenha desativado para pré-visualizações totalmente offline.", + }, + data: "Dados", + exportHistory: "Exportar histórico de chat", + exportHistoryDescription: + "Baixe todos os chats e mensagens em formato JSON.", + exportAction: "Exportar", + exportingAction: "Exportando...", + exportConversations: "Exportar Recentes e Projetos", + exportConversationsDescription: + "Baixe os Recentes ou Recentes mais chats de projetos como JSONL bruto, CSV ou ShareGPT JSONL, combinados ou por chat.", + exportConversationsAction: "Exportar", + exportScopeRecents: "Recentes", + exportScopeAll: "Recentes + Projetos", + exportCombinedSuffix: "(combinado)", + exportPerChatSuffix: "(por chat)", + importChats: "Importar chats", + importChatsDescription: + "Importe um arquivo exportado em JSONL, NDJSON ou CSV para os Recentes.", + importChatsAction: "Importar", + importNoConversations: "Nenhuma conversa encontrada no arquivo.", + importedOneChat: "Importada 1 conversa para os Recentes.", + importedChatCount: "Importadas {count} conversas para os Recentes.", + importFailed: "Falha na importação.", + clearHistory: "Limpar histórico de chat", + clearHistoryDescription: "Exclui o histórico de chat deste dispositivo.", + clearAction: "Limpar", + clearAllChats: "Limpar todos os chats", + clearAllChatsDescription: "Exclui permanentemente todos os chats deste dispositivo.", + noChatsToClear: "Nenhum chat para limpar.", + clearOneChatDescription: + "Exclui permanentemente o único chat deste dispositivo.", + clearChatCountDescription: + "Exclui permanentemente todos os {count} chats deste dispositivo.", + clearChatsAction: "Limpar chats", + clearOneChatTitle: "Limpar 1 chat?", + clearChatsTitle: "Limpar {count} chats?", + clearChatsConfirmDescription: + "Exclui permanentemente todos os chats deste dispositivo. Esta ação não pode ser desfeita.", + clearingAction: "Limpando...", + clearOneChatAction: "Limpar 1 chat", + clearChatCountAction: "Limpar {count} chats", + clearedAllChats: "Todos os chats foram limpos", + clearedOneChat: "1 chat foi limpo", + clearedChatCount: "{count} chats foram limpos", + someChatsCouldNotBeCleared: "Não foi possível limpar alguns chats", + chatsClearedRemainOne: + "{clearedCount} chats limpos; 1 chat restante. Por favor, tente novamente.", + chatsClearedRemain: + "{clearedCount} chats limpos; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemain: + "1 chat limpo; {remainingCount} chats restantes. Por favor, tente novamente.", + oneChatClearedRemainOne: "1 chat limpo; 1 chat restante. Por favor, tente novamente.", + storageClearFailedOne: + "Falha ao limpar o armazenamento; 1 chat pode ter restado. Por favor, tente novamente.", + storageClearFailed: + "Falha ao limpar o armazenamento; {count} chats podem ter restado. Por favor, tente novamente.", + failedToClearChats: "Falha ao limpar os chats", + }, + connections: { + title: "Conexões", + description: "Gerencie provedores e conexões externas.", + }, + apiKeys: { + title: "API", + description: + "Acesse o Unsloth por meio da API compatível com OpenAI.", + readDocs: "Leia a documentação da API", + noAccess: "Nenhum acesso à API ainda.", + newBadge: "Novo", + accessTokens: "Tokens de acesso", + loadError: "Não foi possível carregar o acesso à API.", + createError: "Não foi possível criar o token de acesso.", + revokeError: "Não foi possível revogar o token de acesso.", + never: "Nunca", + tokenNamePlaceholder: "Nome do token (ex: producao)", + newAccessTokenName: "Nome do novo token de acesso", + createToken: "Criar token", + creating: "Criando...", + newTokenCreated: "Novo token de acesso criado", + accessTokenCopied: "Token de acesso copiado", + copyAccessToken: "Copiar token de acesso", + copyNow: "Copie agora - isto não será exibido novamente.", + usageExamples: "Exemplos de uso", + usageTools: "Ferramentas", + exampleCurlTools: "curl + ferramentas", + examplePythonTools: "Python + ferramentas", + exampleJavaScriptTools: "JavaScript + ferramentas", + exampleCurlAdvanced: "curl + avançado", + examplePythonAdvanced: "Python + avançado", + exampleJavaScriptAdvanced: "JavaScript + avançado", + osUnix: "Linux / macOS / WSL", + osWindows: "Windows", + secureHttps: "HTTPS Seguro", + secureHttpsHint: + "A porta 0.0.0.0 ainda está acessível globalmente. Para segurança total, inicie o Unsloth Studio com --secure para expor apenas este link HTTPS.", + copyTunnelUrl: "Copiar URL do túnel", + copySnippet: "Copiar trecho de código", + copy: "Copiar", + copied: "Copiado", + setupDocs: "Docs de configuração:", + relativeNever: "nunca", + relativeJustNow: "agora mesmo", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + relativeMonthsAgo: "há {count} meses", + relativeYearsAgo: "há {count} anos", + expired: "expirado", + today: "hoje", + inDays: "em {count}d", + created: "Criado {value}", + used: "Usado {value}", + expires: "Expira {value}", + actionsFor: "Ações para {name}", + copyPrefix: "Copiar prefixo", + revokeToken: "Revogar token", + revokeTitle: 'Revogar token de acesso "{name}"?', + revokeDescription: + "Aplicativos que usam este token perderão o acesso imediatamente. Esta ação não pode ser desfeita.", + revokeAction: 'Revogar "{name}"', + revoking: "Revogando...", + }, + about: { + title: "Sobre", + description: + "Documentação, notas de lançamento, feedback e informações da build.", + studioVersion: "Versão do Unsloth", + packageVersion: "Versão do Pacote", + llamaCppVersion: "Versão do llama.cpp", + hardware: "Hardware", + gpu: "GPU", + cuda: "CUDA", + rocm: "ROCm", + updates: "Atualização", + help: "Ajuda", + documentation: "Documentação", + releaseNotes: "Notas de lançamento", + whatsNew: "O que há de novo", + feedback: "Feedback", + reportIssue: "Reportar um problema", + license: { + sectionTitle: "Licença", + studioLabel: "Unsloth Studio", + studioLicense: "AGPL-3.0", + studioDescription: + "Código aberto sob a licença GNU AGPL v3.0.", + libraryLabel: "Unsloth Core", + libraryLicense: "Apache-2.0", + libraryDescription: "Licenciado sob Apache 2.0.", + }, + dangerZone: "Zona de perigo", + shutDownStudio: "Desligar Unsloth Studio", + shutDownStudioDescription: + "Interrompe o servidor Unsloth e encerra sua sessão.", + shutDown: "Desligar", + update: { + title: "Atualizar Unsloth Studio", + commandText: "Texto de {label}", + copied: "Copiado", + copyCommand: "Copiar comando", + commandCopied: "{label} copiado", + copyNamedCommand: "Copiar {label}", + checkingInstall: "Verificando como o Unsloth foi instalado...", + installIntro: "Para instalar ou atualizar o Unsloth:", + localUpdateHeading: "Atualização local", + installCommandUnix: "Comando de instalação para macOS/Linux", + installCommandWindows: "Comando de instalação para Windows", + localInstallDetected: + "Instalação local detectada. Atualize a partir do seu repositório original para evitar substituí-lo pelo PyPI.", + pullThenUpdate: "Puxe as últimas alterações (git pull) e depois execute o instalador local:", + gitPullCommand: "comando git pull", + localInstallerCommand: "comando do instalador local", + sourceInstallDetected: + "Instalação do pacote por código-fonte ou VCS detectada. Reinstale a partir do caminho local original ou URL do Git.", + repoCheckoutFallback: + "Se você ainda tiver o repositório baixado, execute o instalador local a partir dele:", + restartAfterUpdate: "Reinicie o Unsloth após a atualização.", + desktopManaged: + "O aplicativo de desktop mantém seu backend integrado atualizado e avisará quando uma nova versão estiver disponível.", + unknownInstall: + "Não foi possível detectar como o Unsloth foi instalado. Para instalações via instalador ou PyPI, use os comandos acima.", + localCheckout: + "Para instalações de repositório local, execute o instalador local a partir desse diretório:", + docs: "Docs de instalação:", + docsInstall: "Instalação", + docsUpdating: "Atualização", + docsMac: "Mac", + docsWindows: "Windows", + }, + }, + }, + studio: { + routeTitle: "Treinar", + title: "Estúdio de Fine-tuning", + subtitles: { + configure: "Configure e inicie o treinamento", + trainingInProgress: "Treinamento em andamento", + viewPastRuns: "Visualizar execuções de treino anteriores", + viewingPastRun: "Visualizando execução anterior", + }, + tabs: { + configure: "Configurar", + currentRun: "Execução Atual", + history: "Histórico", + }, + loadingRuntime: "Carregando ambiente de execução de treino...", + backToHistory: "Voltar ao histórico", + sections: { + model: "Modelo", + dataset: "Dataset", + params: "Parâmetros", + training: "Treinamento", + charts: "Gráficos", + progress: "Progresso do Treinamento", + }, + configure: { + title: "Configurar", + description: "Escolha um modelo, dataset e configurações de treinamento.", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + trainingConfig: "Configuração de Treino", + }, + model: { + title: "Modelo", + description: "Selecione o modelo base e o método de treinamento", + fasterTrainingBadge: "Treinamento 2x Mais Rápido", + baseModel: "Modelo base", + localModel: "Modelo Local", + localModelTooltip: + "Caminho para um modelo baixado localmente ou um repositório HF customizado.", + scanningLocalAndCachedModels: "Escaneando modelos locais e em cache...", + scanning: "Escaneando...", + scanningLocalModels: "Escaneando modelos locais...", + noLocalModelsFound: "Nenhum modelo local encontrado", + noLocalModelsFoundManual: "Nenhum modelo local encontrado. Insira o caminho manualmente.", + failedToLoadLocalModels: "Falha ao carregar modelos locais", + hfCache: "Cache do HF", + customFolders: "Pastas Customizadas", + localDir: "Diretório local", + huggingFaceModel: "Modelo do Hugging Face", + huggingFaceModelTooltip: + "Busque modelos no Hugging Face ou escolha da nossa lista recomendada.", + searchModels: "Buscar modelos...", + searching: "Buscando...", + noModelsFound: "Nenhum modelo encontrado", + needsVram: "Precisa de ~{vram}GB de VRAM (GPU: {gpu}GB)", + tightVram: "~{vram}GB de VRAM (limite na {gpu}GB)", + vramEstimate: "~{vram}GB de VRAM", + method: "Método", + methodTooltip: + "O QLoRA usa quantização de 4 bits para menor uso de VRAM. O LoRA usa 16 bits. O Full atualiza todos os pesos. O CPT (Continued Pretraining) treina em texto bruto para adaptar o modelo a um novo domínio sem formatação de chat.", + readMore: "Leia mais", + fullFineTune: "Fine-tune Completo (Full)", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + huggingFaceTokenOptional: "Token do Hugging Face (Opcional)", + continuedPretraining: "Pré-treinamento Contínuo (CPT)", + localModels: "Modelos locais", + localModelsFound: "{count} modelos locais/em cache encontrados", + loadingLocalModels: "Carregando modelos locais...", + }, + dataset: { + title: "Dataset", + description: "Selecione ou envie os dados de treinamento", + source: "Origem do dataset", + chooseDataset: "Escolher dataset", + chooseDatasetTooltip: + "Use as abas do pop-up para alternar entre o Hugging Face e as saídas de receitas locais.", + localTab: "Local", + searchHuggingFaceDatasets: "Buscar datasets no Hugging Face...", + searchLocalDatasets: "Buscar datasets locais...", + searching: "Buscando...", + noDatasetsFound: "Nenhum dataset encontrado", + loadingLocalDatasets: "Carregando datasets locais...", + failedToLoadLocalDatasets: "Falha ao carregar datasets locais.", + noLocalDatasetsYet: "Nenhum dataset local ainda.", + noLocalDatasetsMatchSearch: "Nenhum dataset local corresponde à busca.", + openDataRecipes: "Abrir Receitas de Dados", + browsingSource: "Navegando em {browsing}. A seleção atual permanece {current}.", + localDatasets: "Datasets locais", + localDataset: "Dataset local", + localDatasetRows: " / {count} linhas", + huggingFaceDataset: "Dataset do Hugging Face", + localDatasetMetadata: "Metadados do dataset local", + dataRecipeOutput: "Saída da Receita de Dados.", + rows: "Linhas", + columns: "Colunas", + batches: "Lotes", + updated: "Atualizado", + evalDataset: "Dataset de validação (Eval)", + uploading: "Enviando...", + upload: "Upload", + uploadEvalFile: "Enviar arquivo de validação", + evalDatasetDescription: + "Opcional. Se não for fornecido, uma pequena parte será dividida a partir dos dados de treinamento.", + advanced: "Avançado", + targetFormat: "Formato de Destino", + targetFormatTooltip: + "Formato dos seus dados de treinamento. A detecção automática funciona para a maioria dos datasets.", + auto: "Auto", + rawText: "Texto Bruto", + trainSplitStart: "Início da Divisão de Treino", + trainSplitStartTooltip: + "Treine apenas em um subconjunto da sua divisão de treino especificando um índice de linha inicial (inclusivo, baseado em 0). Deixe em branco para começar da primeira linha.", + trainSplitEnd: "Fim da Divisão de Treino", + trainSplitEndTooltip: + "Último índice de linha a ser incluído da divisão de treino (inclusivo, baseado em 0). Por exemplo, defina o Início como 0 e o Fim como 99 para treinar nas primeiras 100 linhas. Deixe em branco para usar todas as linhas restantes.", + endPlaceholder: "Fim", + clear: "Limpar", + dropFileOrClick: "Solte 1 arquivo aqui ou clique para fazer upload", + viewDataset: "Visualizar dataset", + uploadFailed: "Falha no envio", + unknownError: "Erro desconhecido", + unsupportedFileType: "Tipo de arquivo não suportado", + uploadOneFileType: "Envie um arquivo do tipo {types}.", + datasetUploaded: "Dataset enviado", + evalDatasetUploaded: "Dataset de validação enviado", + uploadOneFileAtATime: "Envie um arquivo por vez", + uploadSingleFileDescription: + "O upload do dataset de treinamento aceita apenas um único arquivo.", + checkingToken: "Verificando token...", + getOrUpdateToken: "Obter ou atualizar token", + preview: "Pré-visualizar dataset", + split: "Divisão (Split)", + subset: "Subconjunto (Subset)", + s3: { + title: "Configuração do S3", + description: "Carregue datasets em .parquet, .json, .jsonl ou .csv do Amazon S3", + bucket: "Nome do Bucket", + bucketPlaceholder: "meu-bucket-de-dados-de-treino", + region: "Região da AWS", + regionPlaceholder: "us-east-1", + prefix: "Prefixo do Caminho", + prefixPlaceholder: "datasets/whisper/", + prefixTooltip: "Caminho opcional dentro do bucket para os arquivos do seu dataset", + accessKeyId: "ID da Chave de Acesso", + accessKeyIdPlaceholder: "AKIAIOSFODNN7EXAMPLE", + secretAccessKey: "Chave de Acesso Secreta", + secretAccessKeyPlaceholder: "Sua chave de acesso secreta da AWS", + useIamRole: "Usar Função IAM", + useIamRoleTooltip: "Usa credenciais de função IAM em vez de chaves de acesso (recomendado para EC2/SageMaker)", + testConnection: "Testar Conexão", + connectionSuccess: "Conectado com sucesso ao bucket S3", + connectionFailed: "Falha ao conectar ao bucket S3", + comingSoon: "Integração com S3 em breve", + comingSoonDescription: "O carregamento de datasets do S3 requer o boto3. Este recurso está em desenvolvimento.", + }, + }, + params: { + title: "Parâmetros", + description: "Configure os hiperparâmetros de treinamento", + loraSettings: "Configurações do LoRA", + trainingHyperparameters: "Hiperparâmetros de Treinamento", + maxSteps: "Passos Máximos (Max Steps)", + epochs: "Épocas (Epochs)", + useMaxSteps: "Usar Passos Máximos", + useEpochs: "Usar Épocas", + maxStepsTooltip: "Sobrescreve o total de passos do otimizador.", + epochsTooltip: "Número de passagens completas pelo dataset.", + epochsDescription: "Cada época é uma passagem completa pelo seu dataset.", + maxStepsDescription: + "Limita o treinamento a um número fixo de passos do otimizador.", + contextLength: "Comprimento do Contexto", + contextLengthTooltip: "Número máximo de tokens por amostra de treinamento.", + customContextLength: "Insira um valor personalizado", + contextLengthDescription: "Comprimento máximo de sequência para amostras de treino", + learningRate: "Taxa de Aprendizado (Learning Rate)", + learningRateTooltip: + "Tamanho do passo para atualizações de peso. Valores menores treinam mais lentamente, mas com mais estabilidade.", + learningRateDescription: + "Recomendado: 2e-4 para LoRA, 5e-5 para CPT, 2e-5 para fine-tune completo", + embeddingLearningRate: "Taxa de Aprendizado do Embedding", + embeddingLearningRateTooltip: + "Usado apenas quando o CPT está treinando embed_tokens. Os embeddings são mais fáceis de desestabilizar do que os pesos LoRA, por isso geralmente precisam de um LR menor. Deixe em branco para usar lr/10; a faixa típica de funcionamento é de 2x a 10x menor que o LR principal. Aumente apenas se a adaptação de vocabulário ou de tokens de domínio estiver muito lenta.", + embeddingLearningRateDescription: + "Deixe em branco para usar lr/10 (recomendado). A faixa típica é de 2x a 10x menor que a taxa de aprendizado principal.", + rank: "Rank", + rankTooltip: + "Dimensão das matrizes de baixo rank. Maior = mais capacidade.", + alpha: "Alpha", + alphaTooltip: "Fator de escala para atualizações LoRA. Geralmente o dobro do rank.", + dropout: "Dropout", + dropoutTooltip: + "Probabilidade de dropout para as camadas LoRA para reduzir o overfitting.", + visionLayers: "Camadas de visão", + languageLayers: "Camadas de linguagem", + attentionModules: "Módulos de atenção", + mlpModules: "Módulos MLP", + targetModules: "Módulos de Destino", + enableLora: "Ativar LoRA", + trainWithLora: "Treinar com LoRA", + stableRank: "Stable Rank", + memoryEfficient: "Eficiente em Memória", + optimization: "Otimização", + schedule: "Cronograma", + memory: "Memória", + optimizer: "Otimizador", + optimizerTooltip: + "Algoritmo de otimização. Variantes de 8 bits reduzem o uso de memória. Fused é recomendado para modelos de visão.", + lrScheduler: "Agendador de LR", + lrSchedulerTooltip: + "Como a taxa de aprendizado muda ao longo do treino. Linear decai de forma constante; cosine decai em curva.", + optimizerOptions: { + adamw8bit: "AdamW 8-bit", + pagedAdamw8bit: "Paged AdamW 8-bit", + adamwBnb8bit: "AdamW BNB 8-bit", + pagedAdamw32bit: "Paged AdamW 32-bit", + adamwTorch: "AdamW (PyTorch)", + adamwTorchFused: "AdamW (PyTorch Fused)", + }, + lrSchedulerOptions: { + linear: "Linear", + cosine: "Cosine", + }, + batchSize: "Tamanho do Lote (Batch Size)", + batchSizeTooltip: "Amostras processadas por passo. Maior consome mais VRAM.", + gradAccum: "Acúmulo de Gradiente", + gradAccumTooltip: "Simula tamanhos de lote maiores sem gastar VRAM extra.", + weightDecay: "Decaimento de Peso", + weightDecayTooltip: "Regularização L2 para evitar overfitting.", + warmupSteps: "Passos de Aquecimento (Warmup)", + warmupStepsTooltip: + "Aumenta gradualmente a LR no início do treino para garantir estabilidade.", + scheduleEpochsTooltip: + "Número de passagens completas pelo dataset. Defina 0 para rodar por passos máximos.", + saveSteps: "Passos para Salvar", + saveStepsTooltip: "Salva um checkpoint a cada N passos. 0 para desativar.", + evalSteps: "Passos de Validação", + evalStepsTooltip: + "Fração dos passos totais de treino entre as validações (0-1). Defina como 0 para desativar. Ex: 0.01 = valida a cada 1% dos passos.", + seed: "Seed", + seedTooltip: "Semente aleatória para reprodutibilidade.", + gradCheckpoint: "Grad Checkpoint", + gradCheckpointTooltip: + "Troca processamento por memória recalculando as ativações.", + none: "Nenhum", + standard: "Padrão", + enablePacking: "Ativar empacotamento (packing)", + assistantCompletionsOnly: "Apenas respostas do assistente", + readMore: "Leia mais", + }, + training: { + title: "Treinamento", + description: "Monitore e controle o treinamento", + chartNoDataTitle: "Nenhum dado de treinamento ainda", + chartNoDataDescription: "Inicie o treinamento para ver o progresso da loss", + startTraining: "Iniciar Treinamento", + starting: "Iniciando...", + loadingModel: "Carregando modelo...", + checkingDataset: "Verificando dataset...", + configLabel: "Configuração de Treino", + upload: "Upload", + uploadConfigTooltip: "Carregar uma configuração YAML salva", + save: "Salvar", + saveConfigTooltip: "Baixar configuração atual como YAML", + reset: "Redefinir", + resetConfigTooltip: "Redefinir para os padrões do modelo", + configLoaded: "Configuração carregada", + failedToLoadConfig: "Falha ao carregar a configuração", + invalidYamlFile: "Arquivo YAML inválido", + failedToReadFile: "Falha ao ler o arquivo", + parametersReset: "Parâmetros redefinidos para os padrões do modelo", + audioIncompatible: + "Este modelo não suporta áudio. Mude para um modelo compatível com áudio ou escolha um dataset sem áudio.", + visionIncompatible: + "O modelo de texto não é compatível com um dataset multimodal. Mude para um modelo de visão ou escolha um dataset apenas de texto.", + cancelTitle: "Cancelar Treinamento", + cancelDescription: "Deseja cancelar a execução de treinamento atual?", + continueAction: "Continuar Treinamento", + cancelAction: "Cancelar Treinamento", + stopTitle: "Interromper Treinamento", + stopDescription: "Escolha como você deseja interromper a execução de treinamento atual.", + stopAction: "Interromper", + stopping: "Interrompendo...", + stopAndSave: "Interromper e Salvar", + compareInChat: "Comparar no Chat", + exportModel: "Exportar Modelo", + milestone: "Marco", + halfwayDone: "Metade concluída. O treinamento passou de 50%.", + doneNextStep: + "Treinamento concluído. Próximo passo: comparar as saídas do modelo base vs fine-tuned.", + }, + history: { + title: "Histórico", + emptyTitle: "Nenhuma execução de treino ainda", + emptyDescription: + "Nenhuma execução de treino ainda. Inicie sua primeira execução na aba Configurar.", + loadError: "Falha ao carregar as execuções de treino", + deleteError: "Falha ao excluir a execução de treino. Por favor, tente novamente.", + retry: "Tentar novamente", + loadMore: "Carregar mais", + loading: "Carregando...", + loadingRun: "Carregando execução de treino...", + runNotFound: "Execução não encontrada", + deleteTitle: "Excluir execução de treino?", + deleteDescription: + "Isso excluirá permanentemente esta execução de treino e todas as suas métricas. Esta ação não pode ser desfeita.", + runCount: "{count} execuções", + oneRun: "1 execução", + resume: "Retomar", + resumeTraining: "Retomar treinamento", + resuming: "Retomando...", + deleteRun: "Excluir execução", + loss: "Loss", + steps: "Passos", + lossTrendSparkline: "Minigráfico de tendência da loss", + relativeJustNow: "agora mesmo", + relativeMinutesAgo: "há {count}m", + relativeHoursAgo: "há {count}h", + relativeDaysAgo: "há {count}d", + status: { + completed: "Concluído", + stopped: "Interrompido", + error: "Erro", + running: "Em andamento", + continued: "Continuado", + }, + message: { + completed: "Treinamento concluído", + stopped: "Treinamento interrompido", + running: "Treinamento em andamento", + errored: "Treinamento com erro", + }, + }, + charts: { + settings: "Configurações do Gráfico", + settingsDescription: + "Ajuste a apresentação do gráfico enquanto o treinamento continua rodando.", + openSettings: "Abrir configurações do gráfico", + viewWindow: "Janela de visualização", + viewWindowDescription: "Mostra apenas os passos mais recentes ou o histórico completo.", + window: "Janela", + all: "Tudo", + trainingLoss: "Loss de Treinamento", + trainingLossDescription: "Controle as sobreposições e a suavização EMA.", + smoothing: "Suavização", + smoothingDescription: "Mova para a direita para mais suavização. `0` = bruto.", + showRawLoss: "Mostrar loss bruta", + showSmoothedLoss: "Mostrar loss suavizada", + showAverageLine: "Mostrar linha média", + scaleAndCleanup: "Escala e limpeza", + linear: "Linear", + log: "Log", + noClip: "Sem corte", + clipP99: "Cortar p99", + clipP95: "Cortar p95", + lossAxis: "Eixo da loss", + gradientNormAxis: "Eixo da norma do gradiente", + learningRateAxis: "Eixo da taxa de aprendizado", + resetDefaults: "Redefinir padrões", + loss: "Loss", + smoothed: "Suavizado", + evalLoss: "Loss de Validação", + learningRate: "Taxa de Aprendizado", + lr: "LR", + gradNorm: "Norma do Grad.", + gradientNorm: "Norma do Gradiente", + step: "Passo {step}", + averageValue: "média {value}", + waitingForFirstEvaluationStep: "Aguardando o primeiro passo de validação...", + evaluationNotConfigured: "Validação não configurada", + evalChartWillAppear: "O gráfico aparecerá assim que o eval_steps for alcançado", + setEvalDatasetAndSteps: + "Defina o dataset de validação e eval_steps para acompanhar a loss de validação", + }, + progress: { + title: "Progresso do Treinamento", + liveMetrics: "Métricas de treino em tempo real", + exportGguf: "Exportar para GGUF", + openConfig: "Abrir configuração de treino", + configLabel: "Configuração de Treino", + hyperparams: "Hiperparâmetros", + epochs: "Épocas", + batchSize: "Tamanho do lote", + learningRate: "Taxa de aprendizado", + optimizer: "Otimizador", + maxSteps: "Passos máximos", + contextLength: "Comprimento do contexto", + warmupSteps: "Passos de warmup", + rank: "Rank", + alpha: "Alpha", + dropout: "Dropout", + variant: "Variante", + epoch: "Época {value}", + percentComplete: "{percent}% completo", + stepProgress: "Passo {current} / {total}", + loss: "Loss", + lr: "LR", + gradNorm: "Norma do Grad.", + model: "Modelo", + method: "Método", + elapsed: "Decorrido: {value}", + eta: "ETA: {value}", + stepsPerSecond: "{value} passos/s", + noStepsPerSecond: "-- passos/s", + tokens: "Tokens: {value}", + gpuMonitor: "Monitor da GPU", + live: "Ao vivo", + utilization: "Utilização", + temperature: "Temperatura", + vram: "VRAM", + power: "Energia", + phase: { + idle: "Ocioso", + downloadingModel: "Baixando modelo", + downloadingDataset: "Baixando dataset", + loadingModel: "Carregando modelo", + loadingDataset: "Carregando dataset", + configuring: "Configurando", + training: "Treinando", + completed: "Concluído", + error: "Erro", + stopped: "Interrompido", + }, + }, + trainingStart: { + ready: "Pronto", + downloading: "Baixando", + preparing: "Preparando", + left: "restam {eta}", + downloaded: "{size} baixados", + terminalStart: "> treinamento do unsloth iniciado...", + preparingResources: "> Preparando modelo e dataset...", + gettingReady: "> Estamos deixando tudo pronto para a sua execução...", + waitingForFirstStep: "> {message} | aguardando o primeiro passo... ({step})", + resumingTraining: "Retomando treinamento...", + startingTraining: "iniciando treinamento...", + dataset: "Dataset", + datasetStreaming: "Dataset: streaming (sem download completo)", + modelWeights: "Pesos do modelo", + }, + tour: { + guidedTour: "Tour Guiado", + }, + }, +} as const; diff --git a/studio/frontend/src/i18n/messages.ts b/studio/frontend/src/i18n/messages.ts index ef58ca23fe..23db2ce326 100644 --- a/studio/frontend/src/i18n/messages.ts +++ b/studio/frontend/src/i18n/messages.ts @@ -4,19 +4,26 @@ import { getLocale } from "./locale-store"; import { en } from "./locales/en"; import { zhCN } from "./locales/zh-CN"; +import { ptBR } from "./locales/pt-br"; import { ja } from "./locales/ja"; import type { InterpolationValues, MessageKey } from "./types"; export const LOCALES = { en: { label: "English", nativeLabel: "English" }, "zh-CN": { label: "Chinese (Simplified)", nativeLabel: "简体中文" }, - ja: { label: "Japanese", nativeLabel: "日本語" }, + "pt-BR": { label: "Portuguese (Brazil)", nativeLabel: "Português (Brasil)" }, + "ja": { label: "Japanese", nativeLabel: "日本語" }, } as const; export type Locale = keyof typeof LOCALES; export type TranslationKey = MessageKey; -export const messages = { en, "zh-CN": zhCN, ja } as const; +export const messages = { + en, + "zh-CN": zhCN, + "pt-BR": ptBR, + ja +} as const; const PLACEHOLDER_PATTERN = /\{([a-zA-Z0-9_]+)\}/g; @@ -75,4 +82,4 @@ export function isSupportedLocale(value: unknown): value is Locale { typeof value === "string" && Object.prototype.hasOwnProperty.call(LOCALES, value) ); -} +} \ No newline at end of file diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 9060b57542..439d3ffe7b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -181,6 +181,75 @@ def _probe_installed_torch_version() -> str | None: return lines[-1] if lines else None +def _installed_torch_is_windows_rocm() -> bool: + """Return True when the target venv currently has a Windows ROCm torch build. + + This is a belt-and-suspenders guard for the torchao override step: if the + earlier ROCm install path failed to set _rocm_windows_torch_installed but the + venv already contains a ROCm torch wheel, still skip torchao because it + crashes on import on Windows ROCm. + """ + if not IS_WINDOWS: + return False + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import sys, torch; " + "hip = getattr(getattr(torch, 'version', None), 'hip', None) or ''; " + "ver = getattr(torch, '__version__', '').lower(); " + "sys.stdout.write('yes' if (hip or 'rocm' in ver or 'rocmsdk' in ver) else '')" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 90, + **_windows_hidden_subprocess_kwargs(), + ) + except (OSError, subprocess.TimeoutExpired): + return False + lines = [line.strip() for line in (probe.stdout or "").splitlines() if line.strip()] + return probe.returncode == 0 and bool(lines and lines[-1] == "yes") + + +# constraints.txt caps new anyio resolutions at <4.14 (#6483), but an install +# from before the cap existed can already be stuck at 4.14+, which later +# constrained installs won't touch since it already satisfies mcp/fastmcp. +_ANYIO_BAD_FLOOR = (4, 14) + + +def _installed_anyio_version() -> tuple[int, int] | None: + try: + from importlib.metadata import version as _pkg_version + raw = _pkg_version("anyio") + except Exception: + return None + parts = raw.split(".") + try: + major = int(parts[0]) + minor = int(re.sub(r"[^0-9].*", "", parts[1])) if len(parts) > 1 else 0 + except (IndexError, ValueError): + return None + return (major, minor) + + +def _repair_bad_anyio() -> None: + installed = _installed_anyio_version() + if installed is None or installed < _ANYIO_BAD_FLOOR: + return + _safe_print(_dim(f" anyio {installed[0]}.{installed[1]} found -- reinstalling anyio<4.14...")) + pip_install( + "Repairing anyio version", + "--no-cache-dir", + "--force-reinstall", + "anyio<4.14.0", + constrain = False, + ) + + # AMD Windows ROCm wheels (repo.amd.com/rocm/whl/{arch_family}/). # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped/mirror installs. _ROCM_WINDOWS_INDEX_BASE = ( @@ -1970,7 +2039,7 @@ def install_python_stack() -> int: package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") # --local overlays a local repo checkout after updating deps. local_repo = os.environ.get("STUDIO_LOCAL_REPO", "") - base_total = 10 if IS_WINDOWS else 11 + base_total = 11 if IS_WINDOWS else 12 # +1 for the anyio repair check (step 8b) if IS_MACOS: base_total -= 1 # triton step is skipped on macOS if not IS_MACOS and not NO_TORCH: @@ -2221,7 +2290,7 @@ def install_python_stack() -> int: # (no working build; see below). if NO_TORCH: _progress("dependency overrides (skipped, no torch)") - elif _rocm_windows_torch_installed: + elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm(): # No working Windows ROCm torchao build: it imports an absent c10d backend # and crashes transformers.quantizers. Studio stubs it at runtime, so # installing it only ships a package that crashes on import -- skip it. @@ -2284,6 +2353,10 @@ def install_python_stack() -> int: req = REQ_ROOT / "studio.txt", ) + # 8b. anyio repair (#6483) + _progress("anyio check") + _repair_bad_anyio() + # 9. Data-designer dependencies _progress("data designer deps") pip_install( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 38398eb5f5..bb1e88cc4e 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2645,6 +2645,28 @@ if ($env:SKIP_STUDIO_BASE -ne "1" -and $env:STUDIO_LOCAL_INSTALL -ne "1") { if ($InstalledVer -and $LatestVer -and ($InstalledVer -eq $LatestVer)) { step "python" "$_PkgName $InstalledVer is up to date" $SkipPythonDeps = $true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PkgName itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + $_anyioBad = $false + try { + & python -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>$null + if ($LASTEXITCODE -eq 0) { $_anyioBad = $true } + } catch {} + if ($_anyioBad) { + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." "Cyan" + $SkipPythonDeps = $false + } # ...but not if an AMD GPU is present and installed PyTorch is CPU-only # (host predates ROCm-wheel support, or GPU added later): the fast "up to # date" path would leave the user on CPU torch with Train/Export disabled. @@ -3158,7 +3180,86 @@ if ($LlamaPr) { $SkipPrebuiltInstall = $true } -if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { +$LocalLlamaCppLinked = $false +$LocalLlamaCppSrc = $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR +if ($LocalLlamaCppSrc) { + if (-not (Test-Path -LiteralPath $LocalLlamaCppSrc -PathType Container)) { + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $LocalLlamaCppSrc" "Red" + exit 1 + } + $ResolvedLocal = (Resolve-Path -LiteralPath $LocalLlamaCppSrc).Path + # Reusing a local dir disables both the prebuilt download and the source + # build, so a runnable llama-server.exe must already be present. Accept any + # layout LlamaCppBackend._layout_candidates() resolves (root-level, build\bin, + # or build\bin\Release) so the flag never rejects a tree Studio could run. + $LocalLlamaServerFound = $false + foreach ($_cand in @( + (Join-Path $ResolvedLocal "llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\llama-server.exe"), + (Join-Path $ResolvedLocal "build\bin\Release\llama-server.exe"))) { + if (Test-Path -LiteralPath $_cand) { $LocalLlamaServerFound = $true; break } + } + if ($ResolvedLocal -eq $LlamaCppDir) { + # Points at the canonical install location itself: never delete-then-link + # onto itself. Reuse an existing build here (skip prebuilt + source) so the + # staged prebuilt installer can't replace a build the user asked to reuse; + # if nothing is built yet, fall through to the normal install. + if ($LocalLlamaServerFound) { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" "Yellow" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } else { + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" "Yellow" + } + } else { + # Fail clearly rather than junction an unbuilt or wrong-platform checkout + # and leave Studio with no usable binary. + if (-not $LocalLlamaServerFound) { + step "llama.cpp" "no llama-server.exe under $ResolvedLocal (looked for .\llama-server.exe, .\build\bin and .\build\bin\Release) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "Red" + exit 1 + } + # If the target is already a junction/symlink (e.g. a previous + # --with-llama-cpp-dir run), delete only the link via DirectoryInfo.Delete(). + # Remove-Item -Recurse -Force on a reparse point can traverse the link and + # wipe the user's real llama.cpp directory on PowerShell 5.1. Dropping the + # stale link here also keeps the custom-home ownership check below idempotent. + # Use Get-Item -Force (not Test-Path): a *broken* junction whose target was + # moved/deleted makes Test-Path return false, which would leave the dangling + # link in place and make mklink below fail; Get-Item still resolves it so we + # can remove it and relink to a new valid directory. + $existing = Get-Item -LiteralPath $LlamaCppDir -Force -ErrorAction SilentlyContinue + if ($existing -and ($existing.Attributes -band [System.IO.FileAttributes]::ReparsePoint)) { + $existing.Delete() + } + if ($StudioHomeIsCustom) { + Assert-StudioOwnedOrAbsent -Path $LlamaCppDir -Label "llama.cpp install" + } + if (Test-Path -LiteralPath $LlamaCppDir) { + Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue + # A locked/in-use tree can silently survive removal (SilentlyContinue + # masks it). Don't then junction/copy over a half-present dir; mirror the + # prebuilt path's active-process handling and stop with a clear message. + if (Test-Path -LiteralPath $LlamaCppDir) { + step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" + substep "Close Studio or other llama.cpp users and retry" "Yellow" + exit 3 + } + } + cmd /c "mklink /J `"$LlamaCppDir`" `"$ResolvedLocal`"" 2>&1 | Out-Null + if ($LASTEXITCODE -ne 0) { + substep "Could not create directory junction; copying instead..." "Yellow" + Copy-Item -Recurse -LiteralPath $ResolvedLocal -Destination $LlamaCppDir + } + Write-Host "" + step "llama.cpp" "linked local directory: $ResolvedLocal" + $LocalLlamaCppLinked = $true + $NeedLlamaSourceBuild = $false + } +} + +if ($LocalLlamaCppLinked) { + # local directory linked above; skip prebuilt install +} elseif ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") { Write-Host "" substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow" $NeedLlamaSourceBuild = $true @@ -3368,7 +3469,8 @@ if (Test-Path -LiteralPath $LlamaServerBin) { # Install build tools now (last resort) rather than eagerly in Phase 1, so the # prebuilt path stays fast. Same condition as the if/elseif chain below: a source -# build runs only when needed and no usable binary is already present. +# build runs only when needed and no usable binary is already present. A linked +# local dir sets $NeedLlamaSourceBuild = $false, so this no-ops for that path. $WillBuildLlamaFromSource = $NeedLlamaSourceBuild -and ` -not ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") if ($WillBuildLlamaFromSource) { @@ -3377,7 +3479,13 @@ if ($WillBuildLlamaFromSource) { $HasCmakeForBuild = $null -ne (Get-Command cmake -ErrorAction SilentlyContinue) } -if (-not $NeedLlamaSourceBuild) { +if ($LocalLlamaCppLinked) { + # Local dir linked above -- honor the flag's contract: skip BOTH the prebuilt + # download and the source build. Falling through here would run CMake inside + # the user's checkout (via the junction) when it lacks build\bin\Release\llama-server.exe. + Write-Host "" + step "llama.cpp" "linked (skipping build)" +} elseif (-not $NeedLlamaSourceBuild) { Write-Host "" step "llama.cpp" "prebuilt (validated)" } elseif ((Test-Path -LiteralPath $LlamaServerBin) -and -not $NeedRebuild -and $RequestedLlamaTag -ne "master") { diff --git a/studio/setup.sh b/studio/setup.sh index 43ec9fe7b9..6a74cd2296 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -947,6 +947,23 @@ print(version(sys.argv[1])) if [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ] && [ "$INSTALLED_VER" = "$LATEST_VER" ]; then step "python" "$_PKG_NAME $INSTALLED_VER is up to date" _SKIP_PYTHON_DEPS=true + # A pre-#6483-fix install can be stuck on anyio>=4.14 even though + # $_PKG_NAME itself is current; the fast path above would otherwise + # never reach install_python_stack's anyio repair (#6797). + if "$VENV_DIR/bin/python" -c " +import re, sys +from importlib.metadata import version, PackageNotFoundError +try: + parts = version('anyio').split('.') + major = int(parts[0]) + minor = int(re.sub(r'[^0-9].*', '', parts[1])) if len(parts) > 1 else 0 +except (PackageNotFoundError, ValueError, IndexError): + sys.exit(1) +sys.exit(0 if (major, minor) >= (4, 14) else 1) +" 2>/dev/null; then + substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." + _SKIP_PYTHON_DEPS=false + fi elif [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ]; then substep "$_PKG_NAME $INSTALLED_VER -> $LATEST_VER available, updating..." elif [ -z "$LATEST_VER" ]; then @@ -1254,7 +1271,90 @@ fi verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)" -if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then +# GGUF export's check_llama_cpp() looks for a llama-quantize shim at the root of +# the install dir, but a source build keeps the binary under build/bin/. Mirror +# the source-build-reuse step and create the shim when the reused tree has one +# but no root shim yet. Best-effort: the tree may be read-only (shared/CI cache), +# and under `set -e` a failed ln would otherwise abort an good reuse. +_link_local_llama_quantize_shim() { + if [ -x "$1/build/bin/llama-quantize" ] && [ ! -e "$1/llama-quantize" ]; then + ln -sf build/bin/llama-quantize "$1/llama-quantize" 2>/dev/null || \ + substep "could not create llama-quantize shim in linked dir (read-only?); GGUF export may be unavailable" + fi +} + +# Accept any layout LlamaCppBackend._layout_candidates() resolves so the flag +# never rejects a tree Studio could actually run: a root-level llama-server (a +# `make` build or a flat-extracted release) or the CMake build/bin/llama-server. +_has_local_llama_server() { + [ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ] +} + +_LOCAL_LLAMA_CPP_LINKED=false +if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then + if [ ! -d "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" ]; then + step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $UNSLOTH_LOCAL_LLAMA_CPP_DIR" "$C_ERR" + exit 1 + fi + _RESOLVED_LOCAL="$(CDPATH= cd -P -- "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" && pwd -P)" + # Canonicalize the install path the same way before comparing: _RESOLVED_LOCAL + # is fully resolved, but LLAMA_CPP_DIR is textual ($UNSLOTH_HOME/llama.cpp). If + # $HOME (or UNSLOTH_HOME) contains a symlink, the two never match even when the + # user pointed the flag at the canonical install itself -- and the rm -rf below + # would then wipe the very tree they asked to reuse. Resolve via the parent so + # this works whether or not the leaf currently exists. + _CANON_LLAMA_CPP_DIR="$LLAMA_CPP_DIR" + _LLAMA_CPP_PARENT="$(dirname "$LLAMA_CPP_DIR")" + if [ -d "$_LLAMA_CPP_PARENT" ]; then + _CANON_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_LLAMA_CPP_PARENT" && pwd -P)/$(basename "$LLAMA_CPP_DIR")" + fi + if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then + # Points at the canonical install location itself: never delete-then-link + # it onto itself. If a usable build is already there, reuse it and skip + # both the prebuilt download and the source build -- the prebuilt installer + # uses os.replace() and would otherwise clobber an existing source build at + # this path. If nothing is built there yet, fall through to the normal + # install so it gets built in place exactly as it would without the flag. + if _has_local_llama_server "$LLAMA_CPP_DIR"; then + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR is the canonical install location and already holds a build; reusing it" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + else + substep "UNSLOTH_LOCAL_LLAMA_CPP_DIR points to the canonical install location with nothing built there yet; running the normal install" + fi + else + # Reusing disables BOTH the prebuilt download and the source build, so the + # linked tree must already contain a runnable llama-server in one of the + # layouts the backend resolves (root-level or build/bin/). Fail clearly + # rather than link an unbuilt or wrong-platform checkout and leave Studio + # with no usable binary. + if ! _has_local_llama_server "$_RESOLVED_LOCAL"; then + step "llama.cpp" "no llama-server under $_RESOLVED_LOCAL (looked for ./llama-server and ./build/bin/llama-server) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "$C_ERR" + exit 1 + fi + # A stale link from a previous --with-llama-cpp-dir run isn't Studio-owned + # content; drop it before the ownership check so re-runs stay idempotent + # for a custom UNSLOTH_STUDIO_HOME (the assert would otherwise follow the + # link into the user's dir and reject it as unowned). + [ -L "$LLAMA_CPP_DIR" ] && rm -f "$LLAMA_CPP_DIR" + if [ "$_STUDIO_HOME_IS_CUSTOM" = true ]; then + _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" + fi + rm -rf "$LLAMA_CPP_DIR" + ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR" + _link_local_llama_quantize_shim "$LLAMA_CPP_DIR" + step "llama.cpp" "linked local directory: $_RESOLVED_LOCAL" + _LOCAL_LLAMA_CPP_LINKED=true + _NEED_LLAMA_SOURCE_BUILD=false + _SKIP_PREBUILT_INSTALL=true + fi +fi + +if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then + : # local directory linked above; skip prebuilt install +elif [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN" _NEED_LLAMA_SOURCE_BUILD=true elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then diff --git a/tests/python/test_grpo_ddp_model_config.py b/tests/python/test_grpo_ddp_model_config.py new file mode 100644 index 0000000000..5af31f65b8 --- /dev/null +++ b/tests/python/test_grpo_ddp_model_config.py @@ -0,0 +1,31 @@ +"""GRPO logit-scaling helpers must read config through DDP wrappers.""" + +from __future__ import annotations + +import os + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)) +SOURCE_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py") + + +def _read_source() -> str: + with open(SOURCE_PATH, "r") as fh: + return fh.read() + + +def test_grpo_logit_scaling_uses_model_config_helper(): + src = _read_source() + # Helper exists and unwraps DDP/Accelerate wrappers via `.module`. + assert "def _unsloth_get_model_config(model):" in src + assert 'getattr(model.module, "config", None)' in src + # Softcapping takes the model and tolerates a missing config. + assert "logit_softcapping = _unsloth_get_final_logit_softcapping(model)" in src + assert "if config is None:" in src.split("def _unsloth_get_final_logit_softcapping")[1] + # Logit scale/divide read through the unwrapped config, not bare model.config. + assert 'getattr(model_config, "logit_scale", 0)' in src + assert 'getattr(model_config, "logits_scaling", 0)' in src + assert src.count("model_config = _unsloth_get_model_config(model)") >= 2 + # Helper source is injected into the compiled GRPO trainer. + assert "inspect.getsource(_unsloth_get_model_config)" in src + # No direct model.config access remains in the RL logit path. + assert "model.config" not in src diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py new file mode 100644 index 0000000000..2c33f86af1 --- /dev/null +++ b/tests/python/test_mlx_public_trainer_api.py @@ -0,0 +1,1206 @@ +"""Tests for the MLX public trainer compatibility surface.""" + +from __future__ import annotations + +import builtins +import importlib +import importlib.util +import platform +import sys +import types +import warnings + +import pytest + +_MLX_SKIP_REASON = "MLX public trainer API is only active on the MLX backend" + + +def _import_mlx_unsloth(): + """Import unsloth and skip when the current platform is not using MLX.""" + # Skip before importing unsloth so non-MLX hosts missing optional GPU deps + # (e.g. bitsandbytes) skip cleanly instead of erroring at collection. + if not ( + platform.system() == "Darwin" + and platform.machine() == "arm64" + and importlib.util.find_spec("mlx") is not None + ): + pytest.skip(_MLX_SKIP_REASON) + unsloth = importlib.import_module("unsloth") + if getattr(unsloth, "DEVICE_TYPE", None) != "mlx": + pytest.skip(_MLX_SKIP_REASON) + return unsloth + + +class _DummyModel: + """Small model stub that satisfies MLXTrainer constructor probes.""" + + def trainable_parameters(self): + """Return no trainable parameters for constructor-only tests.""" + return {} + + +class _DummyVLMModel(_DummyModel): + """Small VLM model stub for MLX vision trainer constructor probes.""" + + _is_vlm_model = True + + +def test_mlx_exports_unsloth_trainer_api(): + """MLX imports should expose the public Unsloth trainer API.""" + unsloth = _import_mlx_unsloth() + from unsloth import ( + RawTextDataLoader, + TextPreprocessor, + UnslothTrainer, + UnslothTrainingArguments, + clear_gpu_memory, + get_gpu_memory_stats, + ) + + assert RawTextDataLoader is unsloth.RawTextDataLoader + assert TextPreprocessor is unsloth.TextPreprocessor + assert UnslothTrainer is unsloth.UnslothTrainer + assert UnslothTrainingArguments is unsloth.UnslothTrainingArguments + assert get_gpu_memory_stats is unsloth.get_gpu_memory_stats + assert clear_gpu_memory is unsloth.clear_gpu_memory + assert issubclass(UnslothTrainer, unsloth.MLXTrainer) + assert issubclass(UnslothTrainingArguments, unsloth.MLXTrainingConfig) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_non_mlx_exports_public_trainer_api_when_available(): + """GPU/ROCm imports should keep exporting the public Unsloth trainer API.""" + try: + unsloth = importlib.import_module("unsloth") + except ImportError as exc: + # Non-MLX import pulls the optional GPU stack (numpy/torch/unsloth-zoo, + # bitsandbytes/triton, and _gpu_init can re-raise missing deps as + # ImportError). Skip when any of it is unavailable rather than failing + # collection on CPU/ROCm/XPU review hosts. + pytest.skip(f"non-MLX import dependency unavailable: {exc}") + if getattr(unsloth, "DEVICE_TYPE", None) == "mlx": + pytest.skip("non-MLX export smoke test only runs on GPU/ROCm backends") + + assert unsloth.UnslothTrainer is not None + assert unsloth.UnslothTrainingArguments is not None + assert callable(unsloth.get_gpu_memory_stats) + assert callable(unsloth.clear_gpu_memory) + assert importlib.util.find_spec("unsloth.memory") is None + + +def test_mlx_training_arguments_accept_trl_style_kwargs(): + """TRL/SFTConfig-style kwargs should normalize without breaking MLX config.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "bf16.*dataset_kwargs"): + args = unsloth.UnslothTrainingArguments( + max_length = 123, + max_steps = 10, + warmup_ratio = 0.2, + remove_unused_columns = False, + dataset_kwargs = {"skip_prepare_dataset": True}, + bf16 = True, + ) + + assert args.max_seq_length == 123 + assert args.warmup_steps == 2 + assert args.remove_unused_columns is False + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.bf16 is True + assert args.warmup_ratio == 0.2 + assert args._unsloth_mlx_max_seq_length_explicit is False + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): + """Implemented and falsey inert compatibility kwargs should stay quiet.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + args = unsloth.UnslothTrainingArguments( + warmup_ratio = 0.2, + max_steps = 10, + padding_free = False, + remove_unused_columns = False, + assistant_only_loss = False, + completion_only_loss = False, + ) + + assert args.warmup_steps == 2 + assert args.padding_free is False + assert args.remove_unused_columns is False + assert args.completion_only_loss is False + assert caught == [] + + +def test_mlx_training_arguments_prefer_canonical_max_seq_length(): + """Canonical MLX config fields should win over compatibility aliases.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_seq_length = 456, max_length = 123) + dict_args = unsloth.UnslothTrainingArguments( + {"max_length": 123, "max_seq_length": 456}, + ) + + assert args.max_seq_length == 456 + assert args.max_length == 456 + assert args._unsloth_mlx_max_length_value == 456 + assert dict_args.max_seq_length == 456 + assert dict_args.max_length == 456 + assert dict_args._unsloth_mlx_max_length_value == 456 + assert args._unsloth_mlx_max_seq_length_explicit is True + assert dict_args._unsloth_mlx_max_seq_length_explicit is True + + +def test_mlx_training_arguments_preserve_explicit_positive_warmup_steps(): + """Explicit warmup_steps should take precedence over warmup_ratio.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments( + max_steps = 10, + warmup_steps = 5, + warmup_ratio = 0.1, + ) + + assert args.warmup_steps == 5 + assert args._unsloth_mlx_warmup_steps_explicit is True + + +def test_mlx_clear_gpu_memory_uses_metal_fallback(monkeypatch): + """Older MLX releases expose cache clearing under mx.metal.clear_cache.""" + unsloth = _import_mlx_unsloth() + import mlx.core as mx + + called = [] + metal = getattr(mx, "metal", None) or type("Metal", (), {})() + monkeypatch.delattr(mx, "clear_cache", raising = False) + monkeypatch.setattr(mx, "metal", metal, raising = False) + monkeypatch.setattr(metal, "clear_cache", lambda: called.append("metal"), raising = False) + + unsloth.clear_gpu_memory() + + assert called == ["metal"] + + +def test_mlx_training_arguments_preserve_explicit_epoch_training(): + """Epoch-based configs should not inherit the MLX max_steps default.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(num_train_epochs = 1, warmup_ratio = 0.1) + default_args = unsloth.UnslothTrainingArguments() + + assert args.num_train_epochs == 1 + assert args.max_steps == -1 + assert args.warmup_ratio == 0.1 + assert args._unsloth_mlx_warmup_steps_explicit is False + assert default_args.max_steps == unsloth.MLXTrainingConfig.max_steps + + +def test_mlx_training_arguments_keep_mlx_dataset_order_default(): + """Training arguments alone should not override MLX's native data order.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments(max_steps = 1) + explicit_default = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ) + + assert args.dataset_order == "default" + assert args._unsloth_mlx_dataset_order_explicit is False + assert args._unsloth_mlx_max_seq_length_explicit is False + assert explicit_default.dataset_order == "default" + assert explicit_default._unsloth_mlx_dataset_order_explicit is True + + +def test_mlx_training_arguments_warn_on_meaningful_inert_kwargs(): + """Unsupported TrainingArguments knobs should not be silently ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "push_to_hub.*save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "steps", + push_to_hub = True, + padding_free = False, + ) + + assert args.save_strategy == "steps" + assert args.push_to_hub is True + assert args.padding_free is False + + +def test_mlx_training_arguments_reject_unknown_kwargs(): + """Unknown SFTConfig flags should fail instead of becoming inert attributes.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth.UnslothTrainingArguments(assistant_only_loss = True) + + completion_args = unsloth.UnslothTrainingArguments(completion_only_loss = True) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_reject_unsupported_object_flags(): + """Object-style SFTConfig flags should not be silently dropped.""" + unsloth = _import_mlx_unsloth() + + class ArgsObject: + max_steps = 1 + assistant_only_loss = True + + with pytest.raises(NotImplementedError, match = "assistant_only_loss"): + unsloth._coerce_mlx_training_args(ArgsObject()) + + class CompletionArgsObject: + max_steps = 1 + completion_only_loss = True + + completion_args = unsloth._coerce_mlx_training_args(CompletionArgsObject()) + assert completion_args.completion_only_loss is True + + +def test_mlx_training_arguments_accept_output_dir_positional(): + """A single positional output_dir should match TrainingArguments behavior.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("custom-outputs", max_steps = 3) + + assert args.output_dir == "custom-outputs" + assert args.max_steps == 3 + + +def test_mlx_training_arguments_normalize_optim_and_object_aliases(): + """Common notebook optimizer names and object aliases should normalize.""" + unsloth = _import_mlx_unsloth() + + class Scheduler: + value = "cosine" + + class ArgsObject: + optim = "adamw_8bit" + eval_steps = None + lr_scheduler_type = Scheduler() + max_length = 321 + max_steps = 10 + num_train_epochs = 3.0 + save_steps = 500 + save_strategy = "no" + warmup_ratio = 0.1 + warmup_steps = 0 + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth._coerce_mlx_training_args(ArgsObject()) + + assert args.optim == "adamw" + assert args.eval_steps == 0 + assert args.lr_scheduler_type == "cosine" + assert args.max_seq_length == 321 + assert args.num_train_epochs == 3 + assert type(args.num_train_epochs) is int + assert args.save_steps == 0 + assert args.warmup_steps == 1 + assert args._unsloth_mlx_warmup_steps_explicit is False + + +def test_mlx_training_arguments_accept_supported_notebook_kwargs(): + """Supported SFT notebooks should be able to pass their current args.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns( + RuntimeWarning, + match = "bf16.*dataset_kwargs.*gradient_checkpointing_kwargs.*save_strategy", + ): + args = unsloth.UnslothTrainingArguments( + bf16 = True, + dataset_kwargs = {"skip_prepare_dataset": True}, + dataset_num_proc = 4, + dataset_text_field = "text", + embedding_learning_rate = 5e-5, + fp16 = False, + gradient_accumulation_steps = 8, + gradient_checkpointing = True, + gradient_checkpointing_kwargs = {"use_reentrant": False}, + learning_rate = 1e-4, + logging_steps = 2, + lr_scheduler_type = "cosine", + max_grad_norm = 0.3, + max_length = 1024, + max_steps = 10, + num_train_epochs = 1, + optim = "paged_adamw_8bit", + output_dir = "outputs", + padding_free = False, + per_device_train_batch_size = 1, + remove_unused_columns = False, + report_to = "none", + save_strategy = "steps", + seed = 123, + warmup_ratio = 0.1, + weight_decay = 0.01, + ) + + assert args.dataset_num_proc == 4 + assert args.dataset_text_field == "text" + assert args.embedding_learning_rate == 5e-5 + assert args.gradient_accumulation_steps == 8 + assert args.gradient_checkpointing is True + assert args.learning_rate == 1e-4 + assert args.logging_steps == 2 + assert args.lr_scheduler_type == "cosine" + assert args.max_grad_norm == 0.3 + assert args.max_seq_length == 1024 + assert args.max_steps == 10 + assert args.num_train_epochs == 1 + assert args.optim == "adamw" + assert args.output_dir == "outputs" + assert args.per_device_train_batch_size == 1 + assert args.report_to == "none" + assert args.seed == 123 + assert args.warmup_ratio == 0.1 + assert args.warmup_steps == 1 + assert args.weight_decay == 0.01 + assert args.dataset_kwargs == {"skip_prepare_dataset": True} + assert args.gradient_checkpointing_kwargs == {"use_reentrant": False} + assert args.save_strategy == "steps" + + +def test_mlx_training_arguments_honor_direct_no_save_strategy(): + """Direct kwargs should map save_strategy=no to save_steps=0.""" + unsloth = _import_mlx_unsloth() + + with pytest.warns(RuntimeWarning, match = "save_strategy"): + args = unsloth.UnslothTrainingArguments( + save_strategy = "no", + save_steps = 500, + ) + + assert args.save_steps == 0 + + +def test_mlx_trainer_accepts_common_sft_kwargs(): + """UnslothTrainer should accept common SFTTrainer kwargs on MLX.""" + unsloth = _import_mlx_unsloth() + + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + dataset_num_proc = 8, + max_length = 456, + optim = "adamw_bnb_8bit", + processing_class = object(), + ) + + assert trainer.args.max_steps == 1 + assert trainer.args.dataset_num_proc == 8 + assert trainer.args.max_seq_length == 456 + assert trainer.args.max_grad_norm == 1.0 + assert trainer.args.optim == "adamw" + assert trainer.args.dataset_order == "torch_randperm" + assert trainer._unsloth_mlx_ignored_trainer_kwargs == {} + assert caught == [] + + +def test_mlx_trainer_preserves_explicit_dataset_order(): + """UnslothTrainer should only set torch_randperm when order is implicit.""" + unsloth = _import_mlx_unsloth() + + explicit_default = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "default", + ), + ) + explicit_sequential = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + dataset_order = "sequential", + ), + ) + implicit_with_override = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + dataset_num_proc = 4, + ) + implicit_streaming = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, streaming = True), + ) + explicit_no_clip = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + max_grad_norm = 0.0, + ), + ) + + assert explicit_default.args.dataset_order == "default" + assert explicit_sequential.args.dataset_order == "sequential" + assert implicit_with_override.args.dataset_order == "torch_randperm" + assert implicit_streaming.args.dataset_order == "default" + assert implicit_with_override.args.max_grad_norm == 1.0 + assert explicit_no_clip.args.max_grad_norm == 0.0 + + +def test_mlx_trainer_uses_model_context_length_when_implicit(): + """UnslothTrainer should mirror CUDA's max_length bridge precedence.""" + unsloth = _import_mlx_unsloth() + model = _DummyModel() + model.max_seq_length = 321 + max_length_model = _DummyModel() + max_length_model.max_seq_length = 321 + none_model = _DummyModel() + none_model.max_seq_length = 321 + explicit_seq_model = _DummyModel() + explicit_seq_model.max_seq_length = 321 + clamped_seq_model = _DummyModel() + clamped_seq_model.max_seq_length = 321 + model_max_length = _DummyModel() + model_max_length.max_length = 777 + metadata_model = _DummyModel() + metadata_model.config = type("Config", (), {"max_position_embeddings": 888})() + metadata_tokenizer = type("Tokenizer", (), {"model_max_length": 999})() + explicit_max_length_no_model = _DummyModel() + trainer_override_model = _DummyModel() + trainer_override_model.max_seq_length = 321 + config_override_model = _DummyModel() + config_override_model.max_seq_length = 432 + + implicit = unsloth.UnslothTrainer( + model = model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + max_length_args = unsloth.UnslothTrainer( + model = max_length_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + none_args = unsloth.UnslothTrainer( + model = none_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = None), + ) + explicit_seq = unsloth.UnslothTrainer( + model = explicit_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 123), + ) + clamped_seq = unsloth.UnslothTrainer( + model = clamped_seq_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_seq_length = 654), + ) + model_max_length_only = unsloth.UnslothTrainer( + model = model_max_length, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + metadata_ignored = unsloth.UnslothTrainer( + model = metadata_model, + tokenizer = metadata_tokenizer, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + ) + explicit_max_length = unsloth.UnslothTrainer( + model = explicit_max_length_no_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1, max_length = 123), + ) + trainer_override = unsloth.UnslothTrainer( + model = trainer_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments(max_steps = 1), + max_seq_length = 654, + ) + config_with_override = unsloth.UnslothTrainer( + model = config_override_model, + tokenizer = None, + train_dataset = [], + args = unsloth.MLXTrainingConfig(max_steps = 1), + dataset_num_proc = 4, + ) + + assert implicit.args.max_seq_length == 321 + assert implicit.args.max_length == 321 + assert max_length_args.args.max_seq_length == 321 + assert max_length_args.args.max_length == 321 + assert none_args.args.max_seq_length == 321 + assert none_args.args.max_length == 321 + assert explicit_seq.args.max_seq_length == 123 + assert explicit_seq.args.max_length == 123 + assert clamped_seq.args.max_seq_length == 321 + assert clamped_seq.args.max_length == 321 + assert model_max_length_only.args.max_seq_length == 777 + assert model_max_length_only.args.max_length == 777 + assert metadata_ignored.args.max_seq_length == 1024 + assert metadata_ignored.args.max_length == 1024 + assert explicit_max_length.args.max_seq_length == 123 + assert explicit_max_length.args.max_length == 123 + assert trainer_override.args.max_seq_length == 654 + assert trainer_override.args.max_length == 654 + assert config_with_override.args.max_seq_length == 432 + assert config_with_override.args.max_length == 432 + + +def test_mlx_trainer_processing_class_overrides_explicit_none_tokenizer(): + """TRL passes tokenizer=None while processing_class carries the tokenizer.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = processor, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_vision_collator_processor_overrides_processing_class(): + """Vision notebooks pass the tokenizer as processing_class and processor in collator.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + + class Processor: + pass + + processor = Processor() + processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_preserves_explicit_processor_over_vision_collator(): + """Explicit processor kwargs should stay authoritative over collator metadata.""" + unsloth = _import_mlx_unsloth() + tokenizer = object() + explicit_processor = object() + + class Processor: + pass + + collator_processor = Processor() + collator_processor.tokenizer = tokenizer + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), collator_processor) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + processor = explicit_processor, + processing_class = tokenizer, + data_collator = collator, + ) + + assert trainer.processor is explicit_processor + assert trainer.tokenizer is tokenizer + + +def test_mlx_trainer_forwards_vision_collator_positional_defaults(): + """Vision collator CUDA-style positionals should route into MLX args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator( + _DummyVLMModel(), + object(), + 2048, + None, + "max", + -100, + False, + None, + None, + True, + None, + False, + ) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = {"max_steps": 1}, + data_collator = collator, + ) + + assert trainer.args.max_seq_length == 2048 + assert trainer.args.image_size == "max" + assert trainer.args.completion_only_loss is False + + +def test_mlx_vision_collator_default_does_not_override_explicit_args(): + """Implicit collator defaults should not override explicit trainer args.""" + unsloth = _import_mlx_unsloth() + collator = unsloth.UnslothVisionDataCollator(_DummyVLMModel(), object()) + + trainer = unsloth.UnslothTrainer( + model = _DummyVLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = False, + ), + data_collator = collator, + ) + + assert trainer.args.completion_only_loss is False + + +def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): + """Unsupported kwargs that change training semantics should fail on MLX.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "peft_config"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + peft_config = object(), + ) + + +def test_mlx_trainer_rejects_metrics_and_callbacks(): + """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "callbacks"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [object()], + ) + with pytest.raises(NotImplementedError, match = "compute_metrics"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + compute_metrics = lambda *_: None, + ) + + +def test_mlx_trainer_rejects_custom_data_collator(): + """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "data_collator"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + data_collator = object(), + ) + + +def test_mlx_trainer_rejects_text_completion_only_loss(): + """Text MLX training should not silently ignore completion_only_loss=True.""" + unsloth = _import_mlx_unsloth() + + with pytest.raises(NotImplementedError, match = "completion_only_loss=True"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + +def test_mlx_trainer_allows_vlm_completion_only_loss(): + """VLM MLX training supports completion_only_loss during collation.""" + unsloth = _import_mlx_unsloth() + + class VLMModel(_DummyModel): + _is_vlm_model = True + + trainer = unsloth.UnslothTrainer( + model = VLMModel(), + tokenizer = None, + train_dataset = [], + args = unsloth.UnslothTrainingArguments( + max_steps = 1, + completion_only_loss = True, + ), + ) + + assert trainer.args.completion_only_loss is True + + +def test_mlx_trainer_accepts_trl_style_positional_args(): + """TRL-style positional `(model, args, ...)` should not be read as tokenizer.""" + unsloth = _import_mlx_unsloth() + + args = unsloth.UnslothTrainingArguments("trl-outputs", max_steps = 2) + trainer = unsloth.UnslothTrainer( + _DummyModel(), + args, + train_dataset = [], + tokenizer = None, + ) + + assert trainer.args is args + assert trainer.args.output_dir == "trl-outputs" + assert trainer.train_dataset == [] + + +def test_mlx_trainer_accepts_trl_none_placeholder_positionals(): + """Explicit TRL default placeholders should preserve later positional args.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + processing_class = object() + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + None, + processing_class, + ) + + assert getattr(trainer.train_dataset, "_dataset", trainer.train_dataset) is dataset + assert getattr(trainer, "_mlx_train_dataset_for_batches", dataset) is dataset + assert trainer.tokenizer is processing_class + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_none_placeholder_positionals(): + """Short TRL placeholder calls should keep the fourth arg as train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_accepts_short_trl_placeholders_with_keyword_dataset(): + """Short TRL placeholders should not conflict with keyword train_dataset.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + None, + train_dataset = dataset, + ) + + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + assert trainer.args.max_steps == 60 + + +def test_mlx_trainer_preserves_mlx_positional_schema_with_none_tokenizer(): + """MLX-style `(model, tokenizer, train_dataset, ...)` should still work.""" + unsloth = _import_mlx_unsloth() + dataset = [{"text": "hello"}] + + trainer = unsloth.UnslothTrainer( + _DummyModel(), + None, + dataset, + None, + ) + + assert trainer.tokenizer is None + assert trainer.train_dataset is dataset + assert trainer.eval_dataset is None + + +def test_mlx_compatibility_shims_are_installed(): + """Old notebook imports should resolve to the MLX public API after unsloth import.""" + unsloth = _import_mlx_unsloth() + + trl = importlib.import_module("trl") + trainer_module = importlib.import_module("unsloth.trainer") + chat_templates = importlib.import_module("unsloth.chat_templates") + dataset_utils = importlib.import_module("unsloth_zoo.dataset_utils") + + assert importlib.util.find_spec("trl") is not None + assert importlib.util.find_spec("unsloth.trainer") is not None + assert unsloth.trainer is trainer_module + assert unsloth.chat_templates is chat_templates + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trainer_module.UnslothTrainer is unsloth.UnslothTrainer + assert trainer_module.UnslothVisionDataCollator is unsloth.UnslothVisionDataCollator + assert chat_templates.train_on_responses_only is dataset_utils.train_on_responses_only + assert callable(unsloth.train_on_responses_only) + + +def test_mlx_trl_shim_preserves_existing_trl_module(monkeypatch): + """The MLX TRL shim should patch, not replace, an already-loaded TRL module.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.existing_marker = object() + trl.ExistingExport = object() + trl.__all__ = ["ExistingExport", "BrokenExport"] + + def _raise_for_broken_export(name): + if name == "BrokenExport": + raise RuntimeError("optional dependency missing") + raise AttributeError(name) + + trl.__getattr__ = _raise_for_broken_export + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + assert sys.modules["trl"] is trl + assert trl.__path__ == ["real-trainer-package"] + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + assert "ExistingExport" in trl.__all__ + assert "BrokenExport" not in trl.__all__ + assert "SFTTrainer" in trl.__all__ + assert "SFTConfig" in trl.__all__ + + +def test_mlx_trl_shim_installs_real_trl_or_stub(monkeypatch): + """The MLX TRL shim should prefer real TRL and stub only if unavailable.""" + unsloth = _import_mlx_unsloth() + monkeypatch.delitem(sys.modules, "trl", raising = False) + real_trl_available = importlib.util.find_spec("trl") is not None + + unsloth._install_mlx_trl_sft_shim() + trl = importlib.import_module("trl") + + if real_trl_available: + assert trl.__version__ != "0.0.0+unsloth-mlx" + else: + assert trl.__version__ == "0.0.0+unsloth-mlx" + assert trl.SFTTrainer is unsloth.UnslothTrainer + assert issubclass(trl.SFTConfig, unsloth.UnslothTrainingArguments) + assert trl.__UNSLOTH_MLX_COMPAT__ is True + + +def test_mlx_trl_star_import_exports_public_shims(): + """Existing `from trl import *` callers should receive MLX SFT shims.""" + unsloth = _import_mlx_unsloth() + namespace = {} + + exec("from trl import *", namespace) + + assert namespace["SFTTrainer"] is unsloth.UnslothTrainer + assert issubclass(namespace["SFTConfig"], unsloth.UnslothTrainingArguments) + + +def test_mlx_rl_trainers_stub_with_clear_error(monkeypatch): + """GRPO/DPO/ORPO trainers have no MLX path, so the shim retargets the ones trl + exposes to a clear NotImplementedError instead of a confusing CUDA crash, and + never invents trainers trl does not have.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + + class _RealTrainer: + def __init__(self, *args, **kwargs): + raise AssertionError("the real torch/CUDA trainer must not run on MLX") + + trl.GRPOTrainer = _RealTrainer + trl.DPOTrainer = _RealTrainer + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + for name in ("GRPOTrainer", "DPOTrainer"): + assert getattr(trl, name) is not _RealTrainer + with pytest.raises(NotImplementedError) as exc: + getattr(trl, name)(model = None, args = None) + assert "MLX" in str(exc.value) and name in str(exc.value) + # trainers trl never exposed must not be invented + assert not hasattr(trl, "PPOTrainer") + # idempotent: a second install keeps the same stub + stub = trl.GRPOTrainer + unsloth._install_mlx_trl_sft_shim() + assert trl.GRPOTrainer is stub + + +def test_mlx_rl_trainer_stub_is_lazy_import_safe(monkeypatch): + """Stubbing unsupported trl trainers must not resolve them: trl lazy-imports + pull torch, so on a torch-free MLX install a getattr probe would crash + `import unsloth`. The shim reads __all__/vars metadata and never triggers + trl's __getattr__ for a trainer it is about to replace.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "GRPOTrainer", "DPOTrainer"] + resolved = [] + + def _lazy_getattr(name): + resolved.append(name) + raise ImportError(f"lazy import of {name} would pull torch") + + trl.__getattr__ = _lazy_getattr + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() # must not raise despite the lazy trl + + # trainers declared in __all__ are stubbed WITHOUT ever resolving the real one + assert resolved == [] + for name in ("GRPOTrainer", "DPOTrainer"): + with pytest.raises(NotImplementedError): + getattr(trl, name)(model = None) + + +def test_mlx_stubs_trl_trainers_outside_fixed_set(monkeypatch): + """Any non-SFT trainer trl exports (e.g. a newer RLOOTrainer not in the fixed + list) must be stubbed too, so no torch trainer slips through on MLX.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + trl.__all__ = ["SFTTrainer", "SFTConfig", "RLOOTrainer"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + with pytest.raises(NotImplementedError) as exc: + trl.RLOOTrainer(model = None) + assert "MLX" in str(exc.value) and "RLOOTrainer" in str(exc.value) + # SFT stays usable; only non-SFT trainers are stubbed + assert trl.SFTTrainer is unsloth.UnslothTrainer + + +def test_mlx_preserve_dataset_order_is_accepted(): + """preserve_dataset_order=True must be accepted (it is a real MLX config field), + not rejected as an unknown/unsupported argument.""" + unsloth = _import_mlx_unsloth() + args = unsloth.UnslothTrainingArguments( + output_dir = "mlx-out", + max_steps = 10, + preserve_dataset_order = True, + ) + assert getattr(args, "preserve_dataset_order", False) is True + + +def test_mlx_sftconfig_alias_keeps_trl_epoch_default(monkeypatch): + """`trl.SFTConfig` (aliased on MLX) keeps TRL's default training length: with + no explicit max_steps/num_train_epochs it runs TRL's 3 epochs, not the native + MLX 60-step default. An explicit length is authoritative and untouched.""" + unsloth = _import_mlx_unsloth() + trl = types.ModuleType("trl") + trl.__path__ = ["real-trainer-package"] + monkeypatch.setitem(sys.modules, "trl", trl) + + unsloth._install_mlx_trl_sft_shim() + + # no explicit length -> TRL epoch default (3 epochs, step cap disabled) + cfg = trl.SFTConfig(output_dir = "mlx-out") + assert cfg.num_train_epochs == 3 + assert cfg.max_steps == -1 + # explicit step / epoch counts stay exactly as written + assert trl.SFTConfig(output_dir = "mlx-out", max_steps = 17).max_steps == 17 + assert trl.SFTConfig(output_dir = "mlx-out", num_train_epochs = 2).num_train_epochs == 2 + + +def test_mlx_vision_collator_is_constructor_compatible(): + """Vision notebooks should be able to instantiate the collator placeholder.""" + unsloth = _import_mlx_unsloth() + + collator = unsloth.UnslothVisionDataCollator("model", "processor", flag = True) + + assert collator.model == "model" + assert collator.processor == "processor" + assert collator.kwargs == {"completion_only_loss": True, "flag": True} + + +def test_mlx_train_on_responses_only_returns_shared_mask_function(): + """The MLX public shim should expose the shared response-mask helper.""" + unsloth = _import_mlx_unsloth() + + class Tokenizer: + def __call__( + self, + text, + add_special_tokens = False, + ): + return types.SimpleNamespace( + input_ids = { + "": [1], + "": [2], + }[text] + ) + + def convert_tokens_to_ids(self, token): + return token + + mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + ) + masked = mask_fn( + { + "input_ids": [[1, 10, 2, 20, 21, 1, 11]], + } + ) + + assert masked == {"labels": [[-100, -100, -100, 20, 21, -100, -100]]} + + last_mask_fn = unsloth.train_on_responses_only( + None, + instruction_part = "", + response_part = "", + tokenizer = Tokenizer(), + return_function = True, + last_response_only = True, + ) + last_masked = last_mask_fn( + { + "input_ids": [[1, 10, 2, 20, 1, 11, 2, 30]], + } + ) + + assert last_masked == {"labels": [[-100, -100, -100, -100, -100, -100, -100, 30]]} + + +def test_mlx_get_chat_template_uses_light_tokenizer_patch(monkeypatch): + """MLX notebooks should not import CUDA-heavy tokenizer/save helpers.""" + _import_mlx_unsloth() + from unsloth.chat_templates import get_chat_template + import unsloth_zoo.tokenizer_utils as tokenizer_utils + + class Tokenizer: + is_fast = True + padding_side = "right" + eos_token = "" + bos_token = "" + unk_token = "" + pad_token = "" + added_tokens_decoder = {} + + def fake_patch_tokenizer(model, tokenizer): + return model, tokenizer + + real_import = builtins.__import__ + + def guarded_import(name, *args, **kwargs): + if name.startswith("unsloth.models") or name.startswith("unsloth.save"): + raise AssertionError(f"unexpected CUDA-heavy import: {name}") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(tokenizer_utils, "patch_tokenizer", fake_patch_tokenizer) + monkeypatch.setattr(builtins, "__import__", guarded_import) + + tokenizer = get_chat_template( + Tokenizer(), + chat_template = ("{{ messages }}", ""), + ) + + assert tokenizer.chat_template == "{{ messages }}" + assert tokenizer.padding_side == "right" + + +def test_mlx_gpu_memory_stats_helper_shape(): + """The portable memory helper should return CUDA-shaped values.""" + unsloth = _import_mlx_unsloth() + + stats, used, total = unsloth.get_gpu_memory_stats() + + assert isinstance(stats.name, str) + assert hasattr(stats, "total_memory") + assert isinstance(used, float) + assert total > 0 + + +def test_mlx_torch_cuda_compatibility_shim(): + """Existing CUDA memory and move calls should run on MLX.""" + unsloth = _import_mlx_unsloth() + torch = pytest.importorskip("torch") + from transformers.tokenization_utils_base import BatchEncoding + + stats, used, total = unsloth.get_gpu_memory_stats() + cuda_stats = torch.cuda.get_device_properties(0) + + assert cuda_stats.name == stats.name + assert cuda_stats.total_memory == stats.total_memory + assert torch.cuda.get_device_name(0) == stats.name + assert torch.cuda.max_memory_reserved() == int(used * 1024 * 1024 * 1024) + assert torch.cuda.max_memory_allocated() == torch.cuda.max_memory_reserved() + # current (non-max) APIs report live active memory, not the peak high-water + # mark, and never exceed it. + assert 0 <= torch.cuda.memory_reserved() <= torch.cuda.max_memory_reserved() + assert torch.cuda.memory_allocated() == torch.cuda.memory_reserved() + assert torch.cuda.device_count() == 1 + assert torch.cuda.current_device() == 0 + assert torch.cuda.get_device_capability() == (0, 0) + assert total > 0 + + free_bytes, total_bytes = torch.cuda.mem_get_info() + assert total_bytes == int(total * 1024 * 1024 * 1024) + assert 0 <= free_bytes <= total_bytes + + torch.cuda.empty_cache() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + torch.cuda.set_device(0) + + tensor = torch.tensor([1, 2, 3]) + assert tensor.to("cuda") is tensor + assert tensor.cuda() is tensor + assert tensor.to(device = "cuda") is tensor + assert tensor.to("cuda", dtype = torch.float32).dtype == torch.float32 + + batch = BatchEncoding({"input_ids": tensor}) + assert batch.to("cuda") is batch + assert batch.to(device = "cuda") is batch diff --git a/tests/sh/test_with_llama_cpp_dir_flag.sh b/tests/sh/test_with_llama_cpp_dir_flag.sh new file mode 100644 index 0000000000..cee158bf40 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_flag.sh @@ -0,0 +1,172 @@ +#!/bin/bash +# Static analysis: the --with-llama-cpp-dir flag must be wired consistently +# across both installers (install.sh / install.ps1) and both setup scripts +# (studio/setup.sh / studio/setup.ps1). +# +# The flag lets a user point the installer at a local llama.cpp directory so it +# skips BOTH the prebuilt download (Phase 3) and the source build (Phase 4), +# linking the local dir into the canonical install location instead. The path +# crosses the installer->setup boundary via the UNSLOTH_LOCAL_LLAMA_CPP_DIR env +# var. These checks pin that contract so a future refactor of either side can't +# silently break it (e.g. installer parses the flag but setup never reads the +# env var, or setup links the dir but still runs the build). +# +# This is a shape/wiring test, not a behavioral one: it greps the committed +# scripts. It needs no Python, no GPU, no network. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +INSTALL_PS1="$SCRIPT_DIR/../../install.ps1" +SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh" +SETUP_PS1="$SCRIPT_DIR/../../studio/setup.ps1" +ENV_VAR="UNSLOTH_LOCAL_LLAMA_CPP_DIR" +PASS=0 +FAIL=0 + +assert_contains() { + _label="$1"; _file="$2"; _needle="$3" + if grep -qF -- "$_needle" "$_file"; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected to find '$_needle' in $(basename "$_file"))" + FAIL=$((FAIL + 1)) + fi +} + +# Count of distinct lines matching a regex, used to assert a guard appears +# in more than one place (e.g. env var forwarded on both setup invocations). +assert_min_count() { + _label="$1"; _file="$2"; _pattern="$3"; _min="$4" + _n=$(grep -cE -- "$_pattern" "$_file" || true) + if [ "$_n" -ge "$_min" ]; then + echo " PASS: $_label (found $_n, need >= $_min)" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label (found $_n in $(basename "$_file"), need >= $_min)" + FAIL=$((FAIL + 1)) + fi +} + +echo "" +echo "=== install.sh: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.sh: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_SH" "--with-llama-cpp-dir" +assert_contains \ + "install.sh: validates the path exists before forwarding" \ + "$INSTALL_SH" 'if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then' +# The path must be forwarded to setup.sh on BOTH the local and the +# non-local setup invocations, else --local users (the documented path) +# would silently lose the flag. +assert_min_count \ + "install.sh: forwards $ENV_VAR on both setup invocations" \ + "$INSTALL_SH" "$ENV_VAR=\"\\\$_WITH_LLAMA_CPP_DIR\"" 2 + +echo "" +echo "=== install.ps1: parses --with-llama-cpp-dir and forwards the env var ===" + +assert_contains \ + "install.ps1: accepts --with-llama-cpp-dir flag" \ + "$INSTALL_PS1" '"--with-llama-cpp-dir"' +assert_contains \ + "install.ps1: errors when flag is given with no path argument" \ + "$INSTALL_PS1" "--with-llama-cpp-dir requires a path argument" +assert_contains \ + "install.ps1: validates the path exists before forwarding" \ + "$INSTALL_PS1" "--with-llama-cpp-dir path does not exist" +assert_contains \ + "install.ps1: exports $ENV_VAR for setup.ps1" \ + "$INSTALL_PS1" "\$env:$ENV_VAR =" +# The exported env var must be cleaned up so a later setup invocation in the +# same shell session doesn't inherit a stale local-dir link. +assert_contains \ + "install.ps1: clears $ENV_VAR after the setup run" \ + "$INSTALL_PS1" "Remove-Item Env:$ENV_VAR" + +echo "" +echo "=== studio/setup.sh: reads the env var, links, and skips download+build ===" + +assert_contains \ + "setup.sh: reads $ENV_VAR" \ + "$SETUP_SH" "$ENV_VAR" +assert_contains \ + "setup.sh: symlinks the local dir into the canonical install location" \ + "$SETUP_SH" 'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"' +assert_contains \ + "setup.sh: disables the source build when the local dir is linked" \ + "$SETUP_SH" "_NEED_LLAMA_SOURCE_BUILD=false" +assert_contains \ + "setup.sh: skips the prebuilt download when the local dir is linked" \ + "$SETUP_SH" "_SKIP_PREBUILT_INSTALL=true" +# The link branch must short-circuit the FORCE_COMPILE / prebuilt chain rather +# than fall through into it. +assert_contains \ + "setup.sh: link branch gates the prebuilt/compile chain" \ + "$SETUP_SH" 'if [ "$_LOCAL_LLAMA_CPP_LINKED" = true ]; then' + +echo "" +echo "=== studio/setup.ps1: reads the env var, junctions, and skips download+build ===" + +assert_contains \ + "setup.ps1: reads $ENV_VAR" \ + "$SETUP_PS1" "\$env:$ENV_VAR" +assert_contains \ + "setup.ps1: creates a directory junction into the canonical location" \ + "$SETUP_PS1" "mklink /J" +assert_contains \ + "setup.ps1: falls back to a copy when the junction can't be created" \ + "$SETUP_PS1" "Copy-Item -Recurse -LiteralPath \$ResolvedLocal -Destination \$LlamaCppDir" +assert_contains \ + "setup.ps1: disables the source build when the local dir is linked" \ + "$SETUP_PS1" '$NeedLlamaSourceBuild = $false' +# The link branch must gate the prebuilt-install chain (the elseif on +# FORCE_COMPILE), and the linked-dir case must short-circuit the build chain +# so neither a prebuilt download nor a source build runs against it. +assert_contains \ + "setup.ps1: link branch gates the prebuilt/compile chain" \ + "$SETUP_PS1" 'if ($LocalLlamaCppLinked) {' +assert_contains \ + "setup.ps1: linked-dir case short-circuits the build chain" \ + "$SETUP_PS1" 'step "llama.cpp" "linked (skipping build)"' + +echo "" +echo "=== both setup scripts: validate against every layout the backend resolves ===" + +# The linked tree is accepted only if it already holds a runnable llama-server, +# but the check must match LlamaCppBackend._layout_candidates() (root-level +# first, then build/bin, then build/bin/Release on Windows). A narrower check +# would reject a make/flat-release tree the backend could run. +assert_contains \ + "setup.sh: accepts root-level or build/bin llama-server layouts" \ + "$SETUP_SH" '[ -x "$1/llama-server" ] || [ -x "$1/build/bin/llama-server" ]' +assert_contains \ + "setup.ps1: accepts the build\\bin (non-Release) llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "build\bin\llama-server.exe"' +assert_contains \ + "setup.ps1: accepts the root-level llama-server.exe layout" \ + "$SETUP_PS1" 'Join-Path $ResolvedLocal "llama-server.exe"' + +echo "" +echo "=== both setup scripts: a local dir pointing at the canonical path is a no-op ===" + +# Guard against the self-link footgun: if the user passes the canonical install +# dir itself, neither script should delete-then-link it onto itself. +assert_contains \ + "setup.sh: ignores a local dir equal to the canonical install location" \ + "$SETUP_SH" 'if [ "$_RESOLVED_LOCAL" = "$_CANON_LLAMA_CPP_DIR" ]; then' +assert_contains \ + "setup.ps1: ignores a local dir equal to the canonical install location" \ + "$SETUP_PS1" 'if ($ResolvedLocal -eq $LlamaCppDir) {' + +echo "" +echo "=== Results ===" +echo " PASS: $PASS" +echo " FAIL: $FAIL" +if [ "$FAIL" -gt 0 ]; then + echo "FAILED" + exit 1 +fi +echo "ALL PASSED" diff --git a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh new file mode 100644 index 0000000000..fb09e56c51 --- /dev/null +++ b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# Behavioral test for the --with-llama-cpp-dir linking block in studio/setup.sh. +# The companion test_with_llama_cpp_dir_flag.sh is a static wiring check; this one +# actually RUNS the real link logic (extracted from setup.sh by content anchors, +# not line numbers) against hermetic fake dirs and asserts the outcomes Lee asked +# for: an external built dir gets linked, neither the prebuilt download nor the +# source build is armed, an unbuilt dir is rejected, a relink doesn't destroy the +# target, and pointing at the canonical path is a no-op. POSIX symlinks here; +# the Windows junction path is covered by the backend test suite. +set -u +HERE="$(CDPATH= cd -P -- "$(dirname "$0")" && pwd -P)" +SETUP="$HERE/../../studio/setup.sh" +fails=0 +check() { # name expected actual + if [ "$2" = "$3" ]; then printf ' PASS %s\n' "$1" + else printf ' FAIL %s : expected [%s] got [%s]\n' "$1" "$2" "$3"; fails=$((fails+1)); fi +} + +# Extract the two helpers + the whole `UNSLOTH_LOCAL_LLAMA_CPP_DIR` if-block. +# Starts at the quantize-shim helper, ends at the first column-0 `fi` after the +# `if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR..` guard (inner ifs are indented). +block="$(awk ' + /^_link_local_llama_quantize_shim\(\) \{/ {grab=1} + grab {print} + /^if \[ -n "\$\{UNSLOTH_LOCAL_LLAMA_CPP_DIR/ {inif=1} + inif && /^fi$/ {exit} +' "$SETUP")" + +# Self-validate the extraction so a future setup.sh refactor fails loudly here. +case "$block" in *'ln -sfn "$_RESOLVED_LOCAL" "$LLAMA_CPP_DIR"'*) : ;; + *) echo "FAIL: link block extraction broke (no ln -sfn)"; exit 1 ;; esac +case "$block" in *'_has_local_llama_server'*) : ;; + *) echo "FAIL: link block extraction broke (no _has_local_llama_server)"; exit 1 ;; esac + +# Stub setup.sh's logging + ownership helpers, seed the vars the block reads, +# then run the extracted block and print the resulting state. +PREAMBLE=' +set -u +step() { :; }; substep() { :; }; verbose_substep() { :; } +_assert_studio_owned_or_absent() { :; } +C_ERR="" +_STUDIO_HOME_IS_CUSTOM=false +_NEED_LLAMA_SOURCE_BUILD=UNSET +_SKIP_PREBUILT_INSTALL=UNSET +' +EPILOGUE=' +echo "LINKED=$_LOCAL_LLAMA_CPP_LINKED" +echo "NEED_BUILD=$_NEED_LLAMA_SOURCE_BUILD" +echo "SKIP_PREBUILT=$_SKIP_PREBUILT_INSTALL" +if [ -L "$LLAMA_CPP_DIR" ]; then echo "ISLINK=1"; echo "TARGET=$(readlink "$LLAMA_CPP_DIR")"; else echo "ISLINK=0"; fi +' +SNIP="$PREAMBLE"$'\n'"$block"$'\n'"$EPILOGUE" + +# run_link -> prints state lines; RC in $RC +run_link() { + OUT="$(env -i PATH="$PATH" HOME="$T" \ + UNSLOTH_LOCAL_LLAMA_CPP_DIR="$1" LLAMA_CPP_DIR="$2" \ + bash -c "$SNIP" 2>/dev/null)" + RC=$? +} +val() { printf '%s\n' "$OUT" | grep "^$1=" | head -1 | cut -d= -f2-; } + +T="$(mktemp -d)" +trap 'rm -rf "$T"' EXIT + +# Some environments (Windows git-bash without native symlinks) make `ln -s` copy +# instead of link. The symlink-identity assertions (ISLINK / readlink target) +# only run where real symlinks exist; the link/skip/no-data-loss assertions run +# everywhere, including CI (Linux), where the link path is the real one. +ln -s "$T" "$T/.symprobe" 2>/dev/null +if [ -L "$T/.symprobe" ]; then SYMLINKS=1; else SYMLINKS=0; fi +rm -rf "$T/.symprobe" + +# A built external tree (CMake layout) + a flat/`make` tree (root-level binary). +# The fake binary is a shebang script so the `-x` test in _has_local_llama_server +# holds on both Linux (chmod +x) and Windows git-bash (MSYS treats #!-files as +# executable), without needing a real platform binary. +mk_exe() { printf '#!/bin/sh\necho fake\n' > "$1"; chmod +x "$1"; } +mk_built() { mkdir -p "$1/build/bin"; mk_exe "$1/build/bin/llama-server"; } +mk_flat() { mkdir -p "$1"; mk_exe "$1/llama-server"; } + +# 1. External CMake build -> linked, and BOTH install paths disarmed. +EXT1="$T/ext_cmake"; mk_built "$EXT1"; : > "$EXT1/keep.txt" +CANON1="$T/home1/llama.cpp"; mkdir -p "$(dirname "$CANON1")" +run_link "$EXT1" "$CANON1" +check "cmake build: linked" "true" "$(val LINKED)" +check "cmake build: source build off" "false" "$(val NEED_BUILD)" +check "cmake build: prebuilt skipped" "true" "$(val SKIP_PREBUILT)" +if [ "$SYMLINKS" = 1 ]; then + check "cmake build: canonical is a symlink" "1" "$(val ISLINK)" + check "cmake build: link points at external" "$(CDPATH= cd -P -- "$EXT1" && pwd -P)" "$(val TARGET)" +else + printf ' SKIP cmake build: symlink-identity (no real symlinks here)\n' +fi + +# 2. Flat / make tree (root-level llama-server, no build/bin) -> still linked +# (the new layout-candidate acceptance; the old check rejected this). +EXT2="$T/ext_flat"; mk_flat "$EXT2" +CANON2="$T/home2/llama.cpp"; mkdir -p "$(dirname "$CANON2")" +run_link "$EXT2" "$CANON2" +check "flat build: linked (root-level llama-server accepted)" "true" "$(val LINKED)" + +# 3. Unbuilt tree -> rejected (non-zero exit, no link created). +EXT3="$T/ext_empty"; mkdir -p "$EXT3" +CANON3="$T/home3/llama.cpp"; mkdir -p "$(dirname "$CANON3")" +run_link "$EXT3" "$CANON3" +check "unbuilt tree: rejected (exit != 0)" "yes" "$([ "$RC" -ne 0 ] && echo yes || echo no)" +check "unbuilt tree: no link left behind" "no" "$([ -L "$CANON3" ] && echo yes || echo no)" + +# 4. Relink over a stale link must NOT destroy the (new) target's contents. +OLD="$T/ext_old"; mk_built "$OLD" +NEW="$T/ext_new"; mk_built "$NEW"; : > "$NEW/precious.txt" +CANON4="$T/home4/llama.cpp"; mkdir -p "$(dirname "$CANON4")" +ln -sfn "$OLD" "$CANON4" # simulate a prior --with-llama-cpp-dir run +run_link "$NEW" "$CANON4" +if [ "$SYMLINKS" = 1 ]; then + check "relink: now points at the new external" "$(CDPATH= cd -P -- "$NEW" && pwd -P)" "$(val TARGET)" +fi +check "relink: new target's contents preserved" "yes" "$([ -f "$NEW/precious.txt" ] && echo yes || echo no)" +check "relink: old target's contents preserved" "yes" "$([ -f "$OLD/build/bin/llama-server" ] && echo yes || echo no)" + +# 5. Pointing at the canonical path itself is a no-op reuse: linked, not turned +# into a self-referential symlink, contents untouched. +CANON5="$T/home5/llama.cpp"; mk_built "$CANON5"; : > "$CANON5/keep.txt" +run_link "$CANON5" "$CANON5" +check "canonical no-op: linked" "true" "$(val LINKED)" +check "canonical no-op: not made a symlink" "0" "$(val ISLINK)" +check "canonical no-op: contents preserved" "yes" "$([ -f "$CANON5/keep.txt" ] && echo yes || echo no)" + +echo "" +if [ "$fails" -ne 0 ]; then echo "$fails check(s) failed"; exit 1; fi +echo "All checks passed" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 4c50fa1c8e..c8f2053946 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1473,12 +1473,14 @@ class TestHardwareAmdBranching: assert "from . import amd" in source def test_hardware_branches_on_is_rocm_for_utilization(self): - """get_gpu_utilization dispatches to amd.py via _smi_query when IS_ROCM.""" + """get_gpu_utilization dispatches visible metrics through amd.py on ROCm.""" hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" source = hw_path.read_text(encoding = "utf-8") func_start = source.find("def get_gpu_utilization") func_body = source[func_start : source.find("\ndef ", func_start + 1)] - assert '_smi_query("get_primary_gpu_utilization"' in func_body + assert "_smi_query(" in func_body + assert '"get_visible_gpu_utilization"' in func_body + assert "_reconcile_rocm_unified_memory" in func_body smi = source[ source.find("def _smi_query") : source.find("\ndef ", source.find("def _smi_query") + 1) ] @@ -2294,6 +2296,70 @@ class TestRocmTorchInstalledEnvVar: mock_bnb.assert_not_called() +class TestWindowsRocmTorchaoGuard: + """Verify the torchao skip can detect an installed Windows ROCm torch build.""" + + def test_installed_torch_is_windows_rocm_accepts_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "yes" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is True + + def test_installed_torch_is_windows_rocm_rejects_non_rocm_probe(self): + rv = MagicMock() + rv.returncode = 0 + rv.stdout = "" + with ( + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod.subprocess, "run", return_value = rv), + ): + assert stack_mod._installed_torch_is_windows_rocm() is False + + def test_installed_torch_is_windows_rocm_is_non_windows_noop(self): + with patch.object(stack_mod, "IS_WINDOWS", False): + assert stack_mod._installed_torch_is_windows_rocm() is False + + @patch.object(stack_mod, "_repair_bad_anyio") + @patch.object(stack_mod, "_ensure_rocm_torch") + @patch.object(stack_mod, "_ensure_cuda_torch") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True) + @patch.object(stack_mod, "run") + @patch.object(stack_mod, "pip_install") + def test_install_python_stack_skips_torchao_when_windows_rocm_torch_is_installed( + self, mock_pip, mock_run, mock_has_nvidia, mock_cuda, mock_rocm, mock_anyio, tmp_path + ): + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + + subprocess_result = MagicMock() + subprocess_result.returncode = 0 + subprocess_result.stdout = "" + + with ( + patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}), + patch.object(stack_mod, "IS_WINDOWS", True), + patch.object(stack_mod, "IS_MACOS", False), + patch.object(stack_mod, "IS_MAC_ARM", False), + patch.object(stack_mod, "NO_TORCH", False), + patch.object(stack_mod, "_rocm_windows_torch_installed", False), + patch.object(stack_mod, "_bootstrap_uv", return_value = False), + patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = True), + patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin), + patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin), + patch.object(stack_mod.subprocess, "run", return_value = subprocess_result), + ): + assert stack_mod.install_python_stack() == 0 + + installed_specs = [str(arg) for call in mock_pip.call_args_list for arg in call.args] + assert not any("torchao" in arg for arg in installed_specs) + + # TEST: worker.py -- Windows ROCm patches (source-level checks) diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 304de7a9af..209a8a06f1 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -256,7 +256,7 @@ with sync_playwright() as p: composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) - # Detect chat-only mode (/api/health.chat_only): in chat-only mode /studio + /export redirect to /chat. + # Detect chat-only mode (/api/health.chat_only): /studio redirects to /chat while /export stays reachable and self-gated. health_resp = evaluate_fetch( page, f"{BASE}/api/health", @@ -404,15 +404,19 @@ with sync_playwright() as p: # ───────────────────────────────────────────────────── # 3. Export route. # ───────────────────────────────────────────────────── - step(f"Export route ({'chat-only redirect' if chat_only else 'form fields'})") + step(f"Export route ({'chat-only self-gated' if chat_only else 'form fields'})") page.goto(f"{BASE}/export") page.wait_for_timeout(1500) shoot("07-export") if chat_only: - if "/export" in page.url: - soft_fail(f"chat-only mode should redirect /export -> /chat; url={page.url}") + if "/export" not in page.url: + soft_fail(f"chat-only mode should keep /export reachable; url={page.url}") else: - info(f"OK chat-only redirected /export -> {page.url}") + unavailable = page.get_by_text(re.compile(r"Export unavailable", re.I)).first + if unavailable.count() == 0: + soft_fail("chat-only /export did not show the export unavailable gate") + else: + info("OK chat-only /export rendered the unavailable gate") else: # Non-chat-only: verify the export-cta button + HF token field. cta = page.locator('[data-tour="export-cta"]').first diff --git a/tests/studio/run_real_mlx_smoke.py b/tests/studio/run_real_mlx_smoke.py index 0843a2b447..275fe7ac57 100644 --- a/tests/studio/run_real_mlx_smoke.py +++ b/tests/studio/run_real_mlx_smoke.py @@ -403,10 +403,14 @@ def cmd_train(args) -> int: metrics["gguf_dir"] = str(gguf_dir) with Phase("save_gguf", metrics): try: + # q8_0 (the exporter default), not bf16: llama.cpp has optimized q8_0 + # CPU kernels, whereas bf16 CPU decode is unusably slow on the runner + # and made the fresh-process llama-cli reload below time out. q8_0 is + # also what users deploy by default. model.save_pretrained_gguf( str(gguf_dir), tokenizer = tokenizer, - quantization_method = "not_quantized", + quantization_method = "fast_quantized", ) gguf_files = sorted(gguf_dir.glob("*.gguf")) if not gguf_files: @@ -565,30 +569,56 @@ def _reload_gguf(save_dir: Path, metrics: dict) -> int: raise SystemExit(f"no .gguf files in {save_dir}") gguf_path = gguf_files[0] + # Save/reload-integrity smoke (assert below only needs a few chars). The GGUF is + # exported q8_0 (see save_gguf) because llama.cpp bf16 CPU decode is unusably slow + # on the runner. Run CPU-only (-ngl 0), cap the context (-c 256, the model + # advertises 32768), and keep generation short; all env-tunable. + n_predict = os.environ.get("UNSLOTH_GGUF_RELOAD_N", "8") + n_threads = os.environ.get("UNSLOTH_GGUF_RELOAD_THREADS", str(os.cpu_count() or 4)) + n_ctx = os.environ.get("UNSLOTH_GGUF_RELOAD_CTX", "256") + n_gpu_layers = os.environ.get("UNSLOTH_GGUF_RELOAD_NGL", "0") + reload_timeout = int(os.environ.get("UNSLOTH_GGUF_RELOAD_TIMEOUT", "420")) + argv = [ + str(llama_cli), + "-m", + str(gguf_path), + "-p", + PROMPT, + "-n", + n_predict, + "-t", + n_threads, + "-c", + n_ctx, + "-ngl", + n_gpu_layers, + "--temp", + "0", + "--seed", + str(SEED), + "--no-warmup", + ] with Phase("reload_gguf", metrics): - proc = subprocess.run( - [ - str(llama_cli), - "-m", - str(gguf_path), - "-p", - PROMPT, - "-n", - "24", - "--temp", - "0", - "--seed", - str(SEED), - "-no-cnv", - "--no-warmup", - ], - capture_output = True, - text = True, - timeout = 300, - # Hand llama-cli an immediate EOF; without it -no-cnv can still leave the - # process blocked reading stdin, which times out instead of generating. - stdin = subprocess.DEVNULL, - ) + try: + proc = subprocess.run( + argv, + capture_output = True, + text = True, + timeout = reload_timeout, + # Newer llama.cpp keeps llama-cli in chat mode; exit after one reply. + input = "/exit\n", + ) + except subprocess.TimeoutExpired as exc: + + def _decode(stream) -> str: + if isinstance(stream, bytes): + return stream.decode("utf-8", errors = "replace") + return stream or "" + + print(f" [reload:gguf] TIMEOUT running: {' '.join(argv)}", flush = True) + print(f" [reload:gguf] TIMEOUT stdout:\n{_decode(exc.stdout)[:1000]}", flush = True) + print(f" [reload:gguf] TIMEOUT stderr:\n{_decode(exc.stderr)[:1000]}", flush = True) + raise metrics["llama_cli_returncode"] = proc.returncode metrics["generation"] = (proc.stdout or "")[:1500] diff --git a/tests/test_fast_generate_slow_guard.py b/tests/test_fast_generate_slow_guard.py new file mode 100644 index 0000000000..6bfc561e54 --- /dev/null +++ b/tests/test_fast_generate_slow_guard.py @@ -0,0 +1,95 @@ +"""GPU-free test for the fast_generate slow-mode guard in _utils.py. + +When fast_inference=False, model.fast_generate falls back to HuggingFace generate, so vLLM-only +inputs must be rejected with a clear message instead of leaking into transformers.generate. Covers +a string prompt, a vLLM {"prompt":..., "multi_modal_data":...} dict, SamplingParams passed both +positionally and as a kwarg, and a normal tokenized call passing through. +""" + +import ast, functools, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UTILS = os.path.join(HERE, "unsloth", "models", "_utils.py") + + +def _load_factory(): + src = open(UTILS).read() + for node in ast.parse(src).body: + if isinstance(node, ast.FunctionDef) and node.name == "make_fast_generate_wrapper": + ns = {"functools": functools} + exec(ast.get_source_segment(src, node), ns) + return ns["make_fast_generate_wrapper"] + raise AssertionError("make_fast_generate_wrapper not found in _utils.py") + + +make_fast_generate_wrapper = _load_factory() + + +class _SamplingParams: + pass + + +_SamplingParams.__name__ = "SamplingParams" # match by class name, no vllm import needed + + +def _wrapper(): + state = {} + + def original_generate(*a, **k): + state["hit"] = True + return "ok" + + return make_fast_generate_wrapper(original_generate), state + + +def _rejects(fn, needle): + try: + fn() + except ValueError as e: + assert needle in str(e), str(e) + return True + raise AssertionError("expected ValueError") + + +def test_fast_generate_slow_guard(): + w, _ = _wrapper() + # reject every vLLM-only shape + assert _rejects(lambda: w("hello"), "fast_inference=True") + assert _rejects( + lambda: w({"prompt": "hi", "multi_modal_data": {"image": None}}), "fast_inference=True" + ) + assert _rejects(lambda: w(["a", "b"]), "fast_inference=True") + assert _rejects(lambda: w([{"prompt": "hi"}]), "fast_inference=True") # list of prompt dicts + assert _rejects( + lambda: w({"prompt_token_ids": [1, 2, 3]}), "fast_inference=True" + ) # vLLM TokensPrompt + assert _rejects(lambda: w(prompts = "hello"), "fast_inference=True") # vLLM `prompts` kwarg + assert _rejects( + lambda: w(prompts = [{"prompt": "hi"}]), "fast_inference=True" + ) # vLLM `prompts` kwarg list + assert _rejects( + lambda: w(prompt_token_ids = [1, 2, 3]), "fast_inference=True" + ) # vLLM legacy tokenized kwarg + assert _rejects( + lambda: w(prompts = [1, 2, 3]), "fast_inference=True" + ) # token-id list via vLLM-only `prompts` kwarg + assert _rejects( + lambda: w(prompts = None), "fast_inference=True" + ) # vLLM-only kwarg present even if None + assert _rejects(lambda: w({"prompt": "hi"}, _SamplingParams()), "sampling_params") + assert _rejects( + lambda: w({"prompt": "hi"}, [_SamplingParams()]), "sampling_params" + ) # list of SamplingParams + assert _rejects(lambda: w(sampling_params = object()), "sampling_params") + + # pass normal tokenized calls with no false positives + w, state = _wrapper() + assert w(input_ids = "TOKENS", max_new_tokens = 8) == "ok" and state.get("hit") + assert w([1, 2, 3], max_new_tokens = 8) == "ok" # positional token ids + assert w([], max_new_tokens = 8) == "ok" # empty positional + print("13 reject + 3 pass fast_generate slow-mode guard cases passed") + + +if __name__ == "__main__": + test_fast_generate_slow_guard() + print("OK: fast_generate rejects vLLM-style inputs when fast_inference=False") diff --git a/tests/test_generate_kwarg_gate.py b/tests/test_generate_kwarg_gate.py new file mode 100644 index 0000000000..6d1379d3a9 --- /dev/null +++ b/tests/test_generate_kwarg_gate.py @@ -0,0 +1,137 @@ +"""GPU-free test for the generate-kwarg gate in vision.py +(_unsloth_generate_accepts_kwarg), covering both logits_to_keep injection and mm_token_type_ids +stripping, AST-extracted so no unsloth/CUDA import is needed.""" + +import ast, inspect, os + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_helper(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_unsloth_generate_accepts_kwarg": + ns = {"inspect": inspect} + exec(ast.get_source_segment(src, node), ns) + return ns["_unsloth_generate_accepts_kwarg"] + raise AssertionError("_unsloth_generate_accepts_kwarg not found in vision.py") + + +accepts = _load_helper() + + +class PrepHasKwargs_ForwardHasKey: + # **kwargs on prepare unions forward params; key in forward -> ACCEPTED. + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class PrepNoKwargs_ForwardHasKey: + # no **kwargs -> forward not unioned; key only in forward -> REJECTED (gpt-oss shape). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + logits_to_keep = 0, + ): ... + + +class PrepHasKeyDirectly: + # key directly on prepare -> ACCEPTED. + def prepare_inputs_for_generation( + self, + input_ids, + logits_to_keep = 0, + ): ... + def forward(self, input_ids): ... + + +class NoPrepare: + # no prepare -> empty args, no union -> REJECTED. + def forward( + self, + input_ids, + logits_to_keep = 0, + **kwargs, + ): ... + + +class VisionRejectsMM: + # Qwen3-VL shape: neither prepare nor forward names mm_token_type_ids -> REJECTED (stripped). + def prepare_inputs_for_generation( + self, + input_ids, + attention_mask = None, + ): ... + def forward( + self, + input_ids, + pixel_values = None, + ): ... + + +class VisionAcceptsMM: + # forward names mm_token_type_ids and prepare unions it via **kwargs -> ACCEPTED (kept). + def prepare_inputs_for_generation(self, input_ids, **kwargs): ... + def forward( + self, + input_ids, + mm_token_type_ids = None, + **kwargs, + ): ... + + +# (model, key, expected) per gate case. +CASES = [ + ( + "prep(**kwargs)+forward(key) -> accept", + PrepHasKwargs_ForwardHasKey(), + "logits_to_keep", + True, + ), + ( + "prep(no kwargs)+forward(key) -> reject", + PrepNoKwargs_ForwardHasKey(), + "logits_to_keep", + False, + ), + ("prep(key) direct -> accept", PrepHasKeyDirectly(), "logits_to_keep", True), + ("no prepare_inputs_for_gen -> reject", NoPrepare(), "logits_to_keep", False), + ( + "num_logits_to_keep variant -> reject", + PrepNoKwargs_ForwardHasKey(), + "num_logits_to_keep", + False, + ), + ( + "mm_token_type_ids not accepted -> reject (strip)", + VisionRejectsMM(), + "mm_token_type_ids", + False, + ), + ("mm_token_type_ids accepted -> keep", VisionAcceptsMM(), "mm_token_type_ids", True), +] + + +def test_generate_kwarg_gate(): + for name, model, key, expected in CASES: + got = accepts(model, key) + assert got is expected, f"{name}: got {got}, expected {expected}" + + +if __name__ == "__main__": + test_generate_kwarg_gate() + for name, _, _, _ in CASES: + print(f" [PASS] {name}") + print("OK: generate-kwarg gate behaves like transformers _validate_model_kwargs") diff --git a/tests/test_gradient_checkpointing_restore.py b/tests/test_gradient_checkpointing_restore.py new file mode 100644 index 0000000000..4f9f3faccc --- /dev/null +++ b/tests/test_gradient_checkpointing_restore.py @@ -0,0 +1,185 @@ +# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. + +"""Regression for #4735: a plain ``TrainingArguments`` silently disabling the +gradient-checkpointing (GC) mode the model was configured with at setup. + +Setup records the effective GC mode as ``_unsloth_gradient_checkpointing``; the +trainer restores *that* value, falling back to ``args.gradient_checkpointing`` +only when nothing was recorded. The restore lines live inside exec'd template +strings, which ``py_compile`` never sees, so these tests pull the real snippets +out of the source and execute them against fakes. GPU-free. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + +_ROOT = Path(__file__).resolve().parent.parent / "unsloth" / "models" +_RL = (_ROOT / "rl.py").read_text() +_RL_REPLACEMENTS = (_ROOT / "rl_replacements.py").read_text() + +# The single-line ternary form used at the trainer call sites: +# ._unsloth_gradient_checkpointing if hasattr(, '...') else getattr(, 'gradient_checkpointing', True) +_TERNARY = re.compile( + r"(?P[\w.]+)\._unsloth_gradient_checkpointing " + r"if hasattr\((?P=model), '_unsloth_gradient_checkpointing'\) " + r"else getattr\((?P[\w.]+), 'gradient_checkpointing', True\)" +) + +_MISSING = object() + + +class _Obj: + """Bare attribute bag; ``_unsloth_gradient_checkpointing`` present only when recorded.""" + + def __init__( + self, + recorded = _MISSING, + gradient_checkpointing = _MISSING, + ): + if recorded is not _MISSING: + self._unsloth_gradient_checkpointing = recorded + if gradient_checkpointing is not _MISSING: + self.gradient_checkpointing = gradient_checkpointing + + +class _Self: + def __init__( + self, + model = None, + args = None, + ): + if model is not None: + self.model = model + self.args = args + + +# (recorded on model, args.gradient_checkpointing, expected restored value) +# The point of the fix: a recorded mode wins over args, and a recorded ``None`` +# (a valid setup value) is restored verbatim rather than collapsing to the +# args fallback the way a ``None`` sentinel would. +_MATRIX = [ + ("unsloth", False, "unsloth"), # the #4735 case: args=False must NOT win + (True, False, True), + (False, True, False), # user turned GC off; args=True must NOT re-enable it + (None, True, None), # explicit None is restored, not treated as "unrecorded" + (_MISSING, True, True), # nothing recorded -> fall back to args + (_MISSING, False, False), +] + + +def _eval_ternary(expr, recorded, args_gc): + """Eval a restore expression that references either ``model``/``args`` or ``self.model``/``self.args``.""" + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + self = _Self(model = model, args = args) + return eval( + expr, {"hasattr": hasattr, "getattr": getattr}, {"model": model, "args": args, "self": self} + ) + + +def test_ternary_restore_semantics(): + exprs = [m.group(0) for m in _TERNARY.finditer(_RL)] + exprs += [m.group(0) for m in _TERNARY.finditer(_RL_REPLACEMENTS)] + # Also guards against the lines being deleted/renamed (which reinstates the bug). + assert len(exprs) >= 3, f"expected the 3 trainer-call restore sites, found {len(exprs)}" + for expr in exprs: + for recorded, args_gc, expected in _MATRIX: + got = _eval_ternary(expr, recorded, args_gc) + assert got == expected and type(got) is type( + expected + ), f"{expr!r}: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def _extract_prepare_restore_block(): + """Pull the multi-line restore block out of ``prepare_for_training_mode``'s wrapper. + + It lives inside an exec'd template string, so grab it textually: from the + ``_model = getattr(self, 'model', None)`` line through the closing + ``else:``/``use_gc = ...`` pair. + """ + lines = _RL.splitlines() + start = next( + i for i, l in enumerate(lines) if l.strip() == "_model = getattr(self, 'model', None)" + ) + # End at the fallback assignment rather than a fixed line count, so inserting + # lines into the block can't silently truncate what gets exec'd. + end = next( + i + for i, l in enumerate(lines) + if i > start and "use_gc = getattr(self.args, 'gradient_checkpointing', True)" in l + ) + block = lines[start : end + 1] + # dedent to column 0 so it execs as a top-level block + indent = len(block[0]) - len(block[0].lstrip()) + return "\n".join(l[indent:] for l in block) + + +def test_prepare_for_training_mode_block_semantics(): + block = _extract_prepare_restore_block() + # Must be valid Python (it's never seen by py_compile in the outer file). + ast.parse(block) + + for recorded, args_gc, expected in _MATRIX: + model = _Obj(recorded = recorded) + args = _Obj(gradient_checkpointing = args_gc) + ns = {"self": _Self(model = model, args = args), "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + got = ns["use_gc"] + assert ( + got == expected and type(got) is type(expected) + ), f"prepare block: recorded={recorded!r} args={args_gc!r} -> {got!r}, expected {expected!r}" + + +def test_prepare_block_tolerates_missing_model(): + # gemini flagged the unguarded self.model access: the block reads self.model via + # getattr(self, 'model', None), so a trainer without a .model attribute must fall + # back to args rather than raising AttributeError. + block = _extract_prepare_restore_block() + args = _Obj(gradient_checkpointing = True) + self_no_model = _Self(model = None, args = args) # _Self leaves .model unset when model is None + assert not hasattr(self_no_model, "model") + ns = {"self": self_no_model, "hasattr": hasattr, "getattr": getattr} + exec(block, {}, ns) + assert ns["use_gc"] is True + + +def test_recording_sites_are_real_module_code(): + # The recording side (unlike the restore side) is real module code, not a template + # string. Assert it's present at the choke point (patch_peft_model, so loaded adapters + # are covered) and at the pre-wrapped pass-through, both of which bypass the old + # get_peft_model-only recording. + llama = (_ROOT / "llama.py").read_text() + tree = ast.parse(llama) + + def assigns_marker(node): + return any( + isinstance(n, ast.Assign) + and any( + isinstance(t, ast.Attribute) and t.attr == "_unsloth_gradient_checkpointing" + for t in n.targets + ) + for n in ast.walk(node) + ) + + fns = {n.name: n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)} + assert "patch_peft_model" in fns and assigns_marker( + fns["patch_peft_model"] + ), "patch_peft_model must record _unsloth_gradient_checkpointing so loaded adapters are covered" + # The pass-through branch lives in get_peft_model. + assert assigns_marker( + fns["get_peft_model"] + ), "get_peft_model pass-through must record _unsloth_gradient_checkpointing" diff --git a/tests/test_offload_embedding_hooks.py b/tests/test_offload_embedding_hooks.py new file mode 100644 index 0000000000..b8be603b2a --- /dev/null +++ b/tests/test_offload_embedding_hooks.py @@ -0,0 +1,129 @@ +"""Tests _install_offload_embedding_hooks in vision.py: the offloaded lookup must work and +its output must land on the decoder device, read live from the output embeddings (lm_head) +so it tracks model.to() moves. CUDA cases skip without a GPU.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_installer(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_install_offload_embedding_hooks": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_install_offload_embedding_hooks"] + raise AssertionError("_install_offload_embedding_hooks not found in vision.py") + + +install = _load_installer() +CPU = torch.device("cpu") + + +def _emb(): + return nn.Embedding(32, 8) + + +def _lm_head(device): + # Stand-in decoder reference (untied lm_head) whose weight device is the target. + return nn.Linear(8, 32, bias = False).to(device) + + +def test_install_and_idempotent(): + emb = _emb() + lm = _lm_head(CPU) + assert install(emb, lm, CPU) is True + assert emb._unsloth_offload_hooks_installed is True + n_pre = len(emb._forward_pre_hooks) + n_post = len(emb._forward_hooks) + assert install(emb, lm, CPU) is True + assert len(emb._forward_pre_hooks) == n_pre and len(emb._forward_hooks) == n_post + assert install(None, lm, CPU) is False + + +def test_cpu_noop_forward(): + # cpu weight + cpu decoder + cpu input -> output stays cpu. + emb = _emb() + install(emb, _lm_head(CPU), CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.shape == (2, 5, 8) + assert out.device.type == "cpu" + + +def test_cuda_input_roundtrip(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # CPU weight, CUDA decoder + input -> lookup on cpu, output back on cuda. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_cpu_input_still_returns_to_decoder(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P1: offload makes the input arrive on cpu; the output must still reach the cuda decoder. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cpu")) + assert out.device.type == "cuda", out.device + + +def test_live_decoder_over_stale_fallback(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # P2: fallback captured as cpu (model loaded on cpu), but the decoder later lives on cuda. + # The output must follow the live lm_head device, not the stale cpu fallback. + emb = _emb().to("cpu") + install(emb, _lm_head("cuda"), CPU) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +def test_meta_lm_head_falls_back(): + # A disk-offloaded (meta) lm_head must not be used as the return device: moving hidden + # states to meta is unrecoverable, so fall back to the captured device. No GPU needed. + emb = _emb().to("cpu") + lm = _lm_head(CPU) + lm.weight = nn.Parameter(lm.weight.to("meta")) + install(emb, lm, CPU) + out = emb(torch.randint(0, 32, (2, 5))) + assert out.device.type == "cpu", out.device + + +def test_cuda_weight_pulled_back_to_gpu(): + if not torch.cuda.is_available(): + print("[SKIP] CUDA not available") + return + # bf16 weight later pulled back to gpu + cuda input -> no-op, stays on cuda. + emb = _emb().to("cuda") + install(emb, _lm_head("cuda"), torch.device("cuda")) + out = emb(torch.randint(0, 32, (2, 5), device = "cuda")) + assert out.device.type == "cuda", out.device + + +if __name__ == "__main__": + test_install_and_idempotent() + print("[PASS] install + idempotent") + test_cpu_noop_forward() + print("[PASS] cpu no-op forward") + test_cuda_input_roundtrip() + print("[PASS] cuda input roundtrip") + test_cpu_input_still_returns_to_decoder() + print("[PASS] cpu input still returns to cuda decoder (P1)") + test_live_decoder_over_stale_fallback() + print("[PASS] live decoder device beats stale fallback (P2)") + test_meta_lm_head_falls_back() + print("[PASS] meta lm_head falls back to captured device (P2)") + test_cuda_weight_pulled_back_to_gpu() + print("[PASS] cuda weight-on-gpu no-op") + print("OK: offloaded embedding output always lands on the live decoder device") diff --git a/tests/test_offload_tied_guard.py b/tests/test_offload_tied_guard.py new file mode 100644 index 0000000000..096fba116d --- /dev/null +++ b/tests/test_offload_tied_guard.py @@ -0,0 +1,62 @@ +"""Tests _embeddings_are_tied in vision.py: offload_embedding must detect a shared +embed_tokens/lm_head weight so the loader can refuse to offload tied embeddings +(offloading would strand the output projection on CPU). No GPU needed.""" + +import ast, os +import torch +import torch.nn as nn + +HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +VISION = os.path.join(HERE, "unsloth", "models", "vision.py") + + +def _load_fn(): + src = open(VISION).read() + mod = ast.parse(src) + for node in mod.body: + if isinstance(node, ast.FunctionDef) and node.name == "_embeddings_are_tied": + ns = {"torch": torch} + exec(ast.get_source_segment(src, node), ns) + return ns["_embeddings_are_tied"] + raise AssertionError("_embeddings_are_tied not found in vision.py") + + +tied = _load_fn() + + +def test_untied_separate_weights(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + assert tied(emb, lm) is False + + +def test_tied_shared_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = emb.weight # transformers-style weight tying + assert tied(emb, lm) is True + + +def test_tied_by_storage_even_if_distinct_parameter(): + emb = nn.Embedding(32, 8) + lm = nn.Linear(8, 32, bias = False) + lm.weight = nn.Parameter(emb.weight.detach()) # distinct Parameter, shared storage + assert tied(emb, lm) is True + + +def test_none_output_is_untied(): + emb = nn.Embedding(32, 8) + assert tied(emb, None) is False + assert tied(None, nn.Linear(8, 32)) is False + + +if __name__ == "__main__": + test_untied_separate_weights() + print("[PASS] untied separate weights -> False") + test_tied_shared_parameter() + print("[PASS] tied shared parameter -> True") + test_tied_by_storage_even_if_distinct_parameter() + print("[PASS] tied by storage -> True") + test_none_output_is_untied() + print("[PASS] missing lm_head -> untied (safe to offload)") + print("OK: tied embeddings are detected so offload_embedding can refuse them") diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 8202195ca8..04cc600725 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -33,6 +33,27 @@ if platform.system() == "Windows": pass +class _UnslothDeviceStats: + """Portable device metadata used by backend memory-reporting helpers.""" + + def __init__( + self, + name, + total_memory = 0, + ): + """Store a display name and total memory in bytes.""" + self.name = name + self.total_memory = int(total_memory or 0) + self.major = 0 + self.minor = 0 + self.multi_processor_count = 0 + + +def _bytes_to_gb(value): + """Convert byte counts to GiB rounded""" + return round(float(value or 0) / 1024 / 1024 / 1024, 3) + + def _is_mlx_available(): # Transitional import barrier: keep non-Apple-Silicon imports from touching # unsloth_zoo until unsloth_zoo.mlx is import-safe on GPU hosts. Then this @@ -66,7 +87,12 @@ if _IS_MLX: # mlx.trainer / mlx.loader submodules. Surface a friendly install hint # instead of a raw ImportError on the submodule path. try: - from unsloth_zoo.mlx.trainer import MLXTrainer, MLXTrainingConfig + from unsloth_zoo.mlx.trainer import ( + MLXTrainer, + MLXTrainingConfig, + _is_vlm_model, + _normalize_mlx_optimizer_name, + ) from unsloth_zoo.mlx.loader import FastMLXModel except ImportError as _e: raise ImportError( @@ -75,6 +101,53 @@ if _IS_MLX: "`pip install -U unsloth-zoo` or rerun install.sh." ) from _e + import dataclasses as _dataclasses + import importlib.machinery as _machinery + import sys as _sys + import types as _types + import warnings as _warnings + + __version__ = unsloth_zoo.__version__ + DEVICE_TYPE = "mlx" + + def _is_mlx_cuda_device_target(device): + """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" + if device is None: + return False + return str(device).lower().startswith("cuda") + + def _patch_mlx_batch_encoding_to_cuda(): + """Treat tokenizer_output.to("cuda") as a no-op on the MLX backend.""" + try: + from transformers.tokenization_utils_base import BatchEncoding + except Exception: + return + + original_to = getattr(BatchEncoding, "to", None) + if original_to is None or getattr(original_to, "_unsloth_mlx_cuda_noop", False): + return + + def batch_encoding_to( + self, + device = None, + *args, + **kwargs, + ): + target = kwargs.get("device", device) + if _is_mlx_cuda_device_target(target): + return self + # device given by keyword: don't also pass the positional None, or the + # original raises "multiple values for 'device'" (e.g. .to(device="cpu")). + if "device" in kwargs: + return original_to(self, *args, **kwargs) + return original_to(self, device, *args, **kwargs) + + batch_encoding_to._unsloth_mlx_cuda_noop = True + batch_encoding_to._unsloth_original_to = original_to + BatchEncoding.to = batch_encoding_to + + _patch_mlx_batch_encoding_to_cuda() + # Load raw_text helpers without executing dataprep/__init__.py, which # imports synthetic.py -> torch and would defeat the torch-free MLX path. from pathlib import Path as _Path @@ -89,9 +162,6 @@ if _IS_MLX: TextPreprocessor = _raw_text.TextPreprocessor del _raw_text, _raw_text_spec, _raw_text_path, _Path - __version__ = unsloth_zoo.__version__ - DEVICE_TYPE = "mlx" - class FastLanguageModel: @staticmethod def from_pretrained(*args, **kwargs): @@ -141,14 +211,1202 @@ if _IS_MLX: is_bf16_supported = is_bfloat16_supported + def get_gpu_memory_stats(): + """Return MLX device stats, peak memory, and total memory in GiB.""" + import mlx.core as mx + + info = mx.device_info() + total = info.get("memory_size") or info.get("max_recommended_working_set_size") or 0 + get_peak_memory = getattr(mx, "get_peak_memory", None) + if get_peak_memory is None and hasattr(mx, "metal"): + get_peak_memory = getattr(mx.metal, "get_peak_memory", None) + peak = get_peak_memory() if callable(get_peak_memory) else 0 + stats = _UnslothDeviceStats(info.get("device_name", "Apple GPU"), total) + max_memory = _bytes_to_gb(total) or 1.0 + return stats, _bytes_to_gb(peak), max_memory + + def clear_gpu_memory(): + """Clear MLX's cached GPU memory for compatibility cleanup helpers.""" + import mlx.core as mx + + clear_cache = getattr(mx, "clear_cache", None) + if clear_cache is None and hasattr(mx, "metal"): + clear_cache = getattr(mx.metal, "clear_cache", None) + if callable(clear_cache): + clear_cache() + + def _patch_mlx_torch_cuda_compat_api(): + """Expose CUDA-shaped torch helpers for compatibility callers on MLX.""" + try: + import torch + except Exception: + return + + cuda = getattr(torch, "cuda", None) + if cuda is not None and not getattr(cuda, "_unsloth_mlx_cuda_compat_api", False): + + def get_device_properties(device = None): + """Return MLX device stats through torch.cuda's compatibility API.""" + return get_gpu_memory_stats()[0] + + def get_device_name(device = None): + """Return the MLX device name through torch.cuda's compatibility API.""" + return get_device_properties(device).name + + def max_memory_reserved(device = None): + """Return MLX peak memory in bytes for torch.cuda compatibility API.""" + return int(get_gpu_memory_stats()[1] * 1024 * 1024 * 1024) + + def empty_cache(): + """Clear MLX cache through torch.cuda.empty_cache().""" + clear_gpu_memory() + + def _mlx_active_memory_bytes(): + """Current active MLX memory in bytes (not the peak high-water mark).""" + import mlx.core as mx + + get_active = getattr(mx, "get_active_memory", None) + if get_active is None and hasattr(mx, "metal"): + get_active = getattr(mx.metal, "get_active_memory", None) + return int(get_active()) if callable(get_active) else 0 + + def memory_current(device = None): + """Return CURRENT MLX memory in bytes. torch.cuda.memory_reserved / + memory_allocated report live usage, not the peak (that is max_*).""" + return _mlx_active_memory_bytes() + + def mem_get_info(device = None): + """Return (free, total) bytes for torch.cuda compatibility API. + Free uses CURRENT active memory, not the peak high-water mark, so + a capacity check stays accurate after a transient spike.""" + total = int(get_gpu_memory_stats()[2] * 1024 * 1024 * 1024) + return (max(total - _mlx_active_memory_bytes(), 0), total) + + def reset_peak_memory_stats(device = None): + """Reset MLX's peak-memory counter so a later max_memory_reserved / + max_memory_allocated scopes to the run, not earlier model-load peaks.""" + import mlx.core as mx + + reset = getattr(mx, "reset_peak_memory", None) + if reset is None and hasattr(mx, "metal"): + reset = getattr(mx.metal, "reset_peak_memory", None) + if callable(reset): + reset() + + def synchronize(device = None): + """Wait for queued MLX work when torch.cuda.synchronize() is called.""" + import mlx.core as mx + + sync = getattr(mx, "synchronize", None) + if callable(sync): + sync() + + cuda.get_device_properties = get_device_properties + cuda.get_device_name = get_device_name + cuda.max_memory_reserved = max_memory_reserved + cuda.max_memory_allocated = max_memory_reserved + cuda.memory_reserved = memory_current + cuda.memory_allocated = memory_current + cuda.empty_cache = empty_cache + cuda.mem_get_info = mem_get_info + cuda.reset_peak_memory_stats = reset_peak_memory_stats + cuda.synchronize = synchronize + cuda.current_device = lambda: 0 + cuda.device_count = lambda: 1 + cuda.set_device = lambda device = None: None + cuda.get_device_capability = lambda device = None: (0, 0) + cuda.is_bf16_supported = lambda *args, **kwargs: is_bfloat16_supported() + cuda._unsloth_mlx_cuda_compat_api = True + + tensor_to = getattr(torch.Tensor, "to", None) + if tensor_to is not None and not getattr(tensor_to, "_unsloth_mlx_cuda_noop", False): + + def _coerce_mlx_dtype_to_torch(value): + """Map MLX dtype objects to their torch dtype equivalents.""" + try: + import mlx.core as mx + except Exception: + return value + dtype_map = { + mx.bool_: torch.bool, + mx.int8: torch.int8, + mx.int16: torch.int16, + mx.int32: torch.int32, + mx.int64: torch.int64, + mx.uint8: torch.uint8, + mx.float16: torch.float16, + mx.float32: torch.float32, + mx.bfloat16: torch.bfloat16, + } + mapped = dtype_map.get(value, None) + if mapped is not None: + return mapped + dtype_name = str(value).rsplit(".", 1)[-1] + name_map = { + "bool_": torch.bool, + "int8": torch.int8, + "int16": torch.int16, + "int32": torch.int32, + "int64": torch.int64, + "uint8": torch.uint8, + "float16": torch.float16, + "float32": torch.float32, + "bfloat16": torch.bfloat16, + } + return name_map.get(dtype_name, value) + + def mlx_tensor_to(self, *args, **kwargs): + """Ignore CUDA device targets while preserving dtype conversions.""" + args = list(args) + kwargs = dict(kwargs) + removed_cuda_device = False + if args and _is_mlx_cuda_device_target(args[0]): + args.pop(0) + removed_cuda_device = True + if _is_mlx_cuda_device_target(kwargs.get("device", None)): + kwargs.pop("device", None) + removed_cuda_device = True + if removed_cuda_device and not args: + cuda_only_kwargs = ("non_blocking", "copy", "memory_format") + if all(key in cuda_only_kwargs for key in kwargs): + return self + if removed_cuda_device and not args and not kwargs: + return self + if args: + args[0] = _coerce_mlx_dtype_to_torch(args[0]) + if "dtype" in kwargs: + kwargs["dtype"] = _coerce_mlx_dtype_to_torch(kwargs["dtype"]) + return tensor_to(self, *args, **kwargs) + + mlx_tensor_to._unsloth_mlx_cuda_noop = True + mlx_tensor_to._unsloth_original_to = tensor_to + torch.Tensor.to = mlx_tensor_to + + tensor_cuda = getattr(torch.Tensor, "cuda", None) + if tensor_cuda is not None and not getattr(tensor_cuda, "_unsloth_mlx_cuda_noop", False): + + def mlx_tensor_cuda(self, *args, **kwargs): + """Treat tensor.cuda() as a no-op on MLX.""" + return self + + mlx_tensor_cuda._unsloth_mlx_cuda_noop = True + mlx_tensor_cuda._unsloth_original_cuda = tensor_cuda + torch.Tensor.cuda = mlx_tensor_cuda + + _patch_mlx_torch_cuda_compat_api() + + _MLX_TRAINING_CONFIG_FIELDS = {_field.name for _field in _dataclasses.fields(MLXTrainingConfig)} + _MLX_TRAINING_ARGUMENT_ALIASES = { + "max_length": "max_seq_length", + } + _MLX_COMPAT_EXTRA_ARGUMENTS = frozenset( + ( + "bf16", + "dataloader_num_workers", + "dataloader_pin_memory", + "dataset_kwargs", + "ddp_find_unused_parameters", + "disable_tqdm", + "eval_strategy", + "evaluation_strategy", + "fp16", + "full_determinism", + "gradient_checkpointing_kwargs", + "hub_model_id", + "hub_token", + "log_level", + "logging_strategy", + "neftune_noise_alpha", + "optim_args", + "padding_free", + "push_to_hub", + "remove_unused_columns", + "save_on_each_node", + "save_safetensors", + "save_strategy", + "torch_compile", + ) + ) + _MLX_IMPLEMENTED_EXTRA_ARGUMENTS = frozenset( + ( + "image_size", + "preserve_dataset_order", + "warmup_ratio", + ) + ) + _MLX_ALLOWED_EXTRA_ARGUMENTS = _MLX_COMPAT_EXTRA_ARGUMENTS | _MLX_IMPLEMENTED_EXTRA_ARGUMENTS + _MLX_UNSUPPORTED_TASK_ARGUMENTS = frozenset( + ( + "assistant_only_loss", + "completion_only_loss", + ) + ) + + def _is_mlx_no_save_strategy(value): + if hasattr(value, "value"): + value = value.value + strategy = str(value or "").strip().lower() + strategy = strategy.rsplit(".", 1)[-1] + return strategy in ("no", "none", "false") + + _MLX_ADAMW_OPTIMIZER_ALIASES = frozenset( + ( + "adamw_8bit", + "paged_adamw_8bit", + "adamw_bnb_8bit", + "paged_adamw_32bit", + "adamw_torch", + "adamw_torch_fused", + "paged_adamw", + "adamw_32bit", + "adamw_hf", + "adamw_anyprecision", + "adamw_apex_fused", + ) + ) + + def _normalize_mlx_training_value(key, value): + if key == "eval_steps" and value is None: + return 0 + if key == "num_train_epochs" and value is not None and not isinstance(value, bool): + try: + epochs = float(value) + except (TypeError, ValueError): + pass + else: + if epochs.is_integer(): + return int(epochs) + if key == "lr_scheduler_type" and hasattr(value, "value"): + return value.value + if key != "optim": + return value + try: + return _normalize_mlx_optimizer_name(value) + except ValueError: + # Older unsloth-zoo lacks CUDA/TRL optimizer aliases; map common + # adamw_* names so notebook defaults (optim="adamw_8bit") still work. + opt = str(getattr(value, "value", value) or "adamw").strip().lower() + opt = opt.rsplit(".", 1)[-1].replace("-", "_") + if opt in _MLX_ADAMW_OPTIMIZER_ALIASES: + return "adamw" + raise + + def _mlx_training_argument_values(args): + values = {} + for field in _dataclasses.fields(MLXTrainingConfig): + if hasattr(args, field.name): + values[field.name] = _normalize_mlx_training_value( + field.name, + getattr(args, field.name), + ) + for alias, target in _MLX_TRAINING_ARGUMENT_ALIASES.items(): + if target not in values and hasattr(args, alias): + values[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else alias] = getattr( + args, alias + ) + for name in _MLX_ALLOWED_EXTRA_ARGUMENTS: + if hasattr(args, name): + values[name] = getattr(args, name) + for name in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if hasattr(args, name): + value = getattr(args, name) + if ( + name == "completion_only_loss" + and value is not None + and name in _MLX_TRAINING_CONFIG_FIELDS + ): + values[name] = value + elif value is not None and value is not False: + values[name] = value + if _is_mlx_no_save_strategy(values.get("save_strategy", None)): + values["save_steps"] = 0 + return values + + def _split_mlx_trainer_kwargs(kwargs): + trainer_kwargs = {} + config_kwargs = {} + ignored_kwargs = {} + for key, value in kwargs.items(): + if key in _MLX_TRAINER_KWARGS: + trainer_kwargs[key] = value + continue + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if target in _MLX_TRAINING_CONFIG_FIELDS or key in _MLX_ALLOWED_EXTRA_ARGUMENTS: + config_kwargs[key] = value + else: + ignored_kwargs[key] = value + return trainer_kwargs, config_kwargs, ignored_kwargs + + def _is_mlx_training_args_like(value): + if isinstance(value, (MLXTrainingConfig, dict, str, os.PathLike)): + return True + return any( + hasattr(value, name) + for name in ( + "output_dir", + "per_device_train_batch_size", + "gradient_accumulation_steps", + "max_steps", + "learning_rate", + ) + ) + + def _should_use_trl_positional_schema(args): + if len(args) < 2: + return False + if _is_mlx_training_args_like(args[1]): + return True + # TRL callers often pass explicit defaults: + # SFTTrainer(model, None, None, train_dataset, ...) + return len(args) >= 3 and args[1] is None and (args[2] is None or callable(args[2])) + + def _assign_mlx_positional_kwarg(kwargs, name, value): + if name in kwargs: + raise TypeError( + f"UnslothTrainer.__init__() got multiple values for argument " f"{name!r}" + ) + kwargs[name] = value + + def _normalize_mlx_trainer_init_args(args, kwargs): + kwargs = dict(kwargs) + if len(args) == 0: + return kwargs + + use_trl_schema = _should_use_trl_positional_schema(args) + positional_names = ( + _TRL_SFT_TRAINER_POSITIONAL_KWARGS if use_trl_schema else _MLX_TRAINER_POSITIONAL_KWARGS + ) + if len(args) > len(positional_names): + raise TypeError( + f"UnslothTrainer.__init__() takes at most " + f"{len(positional_names)} positional arguments on MLX " + f"({len(args)} given)" + ) + for name, value in zip(positional_names, args): + _assign_mlx_positional_kwarg(kwargs, name, value) + return kwargs + + def _is_meaningful_mlx_extra_value(value): + if value is None or value is False: + return False + if isinstance(value, (str, bytes)) and len(value) == 0: + return False + if isinstance(value, (dict, list, tuple, set, frozenset)) and len(value) == 0: + return False + return True + + def _warn_ignored_mlx_training_args(extra_kwargs): + names = sorted( + key + for key, value in extra_kwargs.items() + if (key in _MLX_COMPAT_EXTRA_ARGUMENTS and _is_meaningful_mlx_extra_value(value)) + ) + if not names: + return + _warnings.warn( + "Unsloth MLX: accepting but not applying unsupported " + "TrainingArguments kwargs: " + f"{', '.join(names)}. These options are not implemented by " + "MLXTrainer yet.", + RuntimeWarning, + stacklevel = 3, + ) + + def _is_meaningful_mlx_trainer_kwarg(key, value): + if key == "optimizers" and value == (None, None): + return False + return _is_meaningful_mlx_extra_value(value) + + def _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs): + names = sorted( + key + for key, value in ignored_kwargs.items() + if _is_meaningful_mlx_trainer_kwarg(key, value) + ) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported SFTTrainer kwargs cannot be ignored safely: " + f"{', '.join(names)}. Remove these kwargs or use a supported MLX " + "trainer configuration." + ) + + def _raise_unknown_mlx_training_args(extra_kwargs): + names = sorted(key for key in extra_kwargs if key not in _MLX_ALLOWED_EXTRA_ARGUMENTS) + if not names: + return + raise NotImplementedError( + "Unsloth MLX: unsupported TrainingArguments/SFTConfig kwargs: " + f"{', '.join(names)}. Remove these kwargs or use fields implemented " + "by MLXTrainingConfig." + ) + + def _positive_mlx_context_length(value): + if value is None or isinstance(value, bool): + return None + try: + length = int(value) + except (TypeError, ValueError, OverflowError): + return None + if length <= 0: + return None + return length + + def _positive_mlx_training_number(value): + if value is None or isinstance(value, bool): + return None + try: + number = float(value) + except (TypeError, ValueError, OverflowError): + return None + if number <= 0: + return None + return number + + def _set_mlx_cuda_style_context_length(args, length): + args.max_seq_length = length + args.max_length = length + args._unsloth_mlx_max_length_value = length + return args + + class UnslothTrainingArguments(MLXTrainingConfig): + """MLX-compatible public training arguments for Unsloth notebooks.""" + + def __init__(self, *args, **kwargs): + if len(args) == 1 and isinstance(args[0], dict): + kwargs = {**args[0], **kwargs} + elif len(args) == 1 and isinstance(args[0], (str, os.PathLike)): + kwargs = {"output_dir": os.fspath(args[0]), **kwargs} + elif args: + raise TypeError( + "UnslothTrainingArguments on MLX accepts keyword arguments, " + "a dict, or a single positional output_dir." + ) + + max_length_value = kwargs.get("max_length", None) + # Only the canonical max_seq_length marks context length explicit; TRL + # max_length stays a compatibility alias and defers to the model's + # context length when one is available. + max_seq_length_explicit = ( + _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ) + if "max_length" in kwargs and "max_seq_length" not in kwargs: + kwargs["max_seq_length"] = kwargs["max_length"] + elif ( + "max_length" in kwargs + and _positive_mlx_context_length(kwargs.get("max_seq_length", None)) is not None + ): + max_length_value = kwargs["max_seq_length"] + if "num_train_epochs" in kwargs and "max_steps" not in kwargs: + kwargs["max_steps"] = -1 + + dataset_order_explicit = "dataset_order" in kwargs or bool( + kwargs.get("preserve_dataset_order", False) + ) + append_eos_explicit = "append_eos" in kwargs + grad_clip_explicit = any( + name in kwargs for name in ("max_grad_norm", "max_grad_value", "max_grad_leaf_norm") + ) + warmup_ratio = kwargs.get("warmup_ratio", None) + warmup_steps_supplied = "warmup_steps" in kwargs + warmup_steps_value = kwargs.get("warmup_steps", None) + warmup_steps_explicit = False + if warmup_steps_supplied: + try: + warmup_steps_explicit = int(warmup_steps_value) > 0 + except (TypeError, ValueError): + warmup_steps_explicit = True + filtered_kwargs = {} + extra_kwargs = {} + for key, value in kwargs.items(): + target = _MLX_TRAINING_ARGUMENT_ALIASES.get(key, key) + if key != target and target in kwargs: + continue + value = _normalize_mlx_training_value(target, value) + if target in _MLX_UNSUPPORTED_TASK_ARGUMENTS: + if ( + target == "completion_only_loss" + and value is not None + and target in _MLX_TRAINING_CONFIG_FIELDS + ): + filtered_kwargs[target] = value + elif _is_meaningful_mlx_extra_value(value): + extra_kwargs[key] = value + continue + if target in _MLX_TRAINING_CONFIG_FIELDS: + filtered_kwargs[target] = value + else: + extra_kwargs[target if target in _MLX_ALLOWED_EXTRA_ARGUMENTS else key] = value + + _raise_unknown_mlx_training_args(extra_kwargs) + + if _is_mlx_no_save_strategy(extra_kwargs.get("save_strategy", None)): + filtered_kwargs["save_steps"] = 0 + + if warmup_ratio is not None and not warmup_steps_explicit: + import math as _math + max_steps = filtered_kwargs.get( + "max_steps", + getattr(MLXTrainingConfig, "max_steps", 60), + ) + try: + if int(max_steps) > 0: + filtered_kwargs["warmup_steps"] = max( + 0, + _math.ceil(int(max_steps) * float(warmup_ratio)), + ) + except (TypeError, ValueError): + pass + + super().__init__(**filtered_kwargs) + self._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + self._unsloth_mlx_append_eos_explicit = append_eos_explicit + self._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + self._unsloth_mlx_max_length_value = max_length_value + if "max_length" in kwargs: + self.max_length = max_length_value + self._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + self._unsloth_mlx_warmup_steps_explicit = warmup_steps_explicit + self._unsloth_mlx_extra_args = extra_kwargs + for key, value in extra_kwargs.items(): + setattr(self, key, value) + _warn_ignored_mlx_training_args(extra_kwargs) + + def _resolve_mlx_cuda_style_max_seq_length(args, model = None): + model_max_seq_length = _positive_mlx_context_length( + getattr(model, "max_seq_length", None), + ) + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + args_max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if args_max_seq_length_explicit is None: + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + args_max_seq_length_explicit = ( + args_max_seq_length is not None and args_max_seq_length != default_max_seq_length + ) + if not args_max_seq_length_explicit: + args_max_seq_length = None + + if args_max_seq_length is None and model_max_seq_length is not None: + args_max_seq_length = model_max_seq_length + elif ( + args_max_seq_length is not None + and model_max_seq_length is not None + and args_max_seq_length > model_max_seq_length + ): + print( + "Unsloth: You set `max_seq_length` as " + f"{args_max_seq_length} but the maximum the model supports is " + f"{model_max_seq_length}. We shall reduce it." + ) + args_max_seq_length = model_max_seq_length + + if args_max_seq_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_seq_length) + return args + + model_max_length = model_max_seq_length + if model_max_length is None: + model_max_length = _positive_mlx_context_length( + getattr(model, "max_length", None), + ) + if model_max_length is not None: + _set_mlx_cuda_style_context_length(args, model_max_length) + return args + + args_max_length = _positive_mlx_context_length( + getattr(args, "max_length", None), + ) + if args_max_length is None: + args_max_length = _positive_mlx_context_length( + getattr(args, "_unsloth_mlx_max_length_value", None), + ) + if args_max_length is not None: + _set_mlx_cuda_style_context_length(args, args_max_length) + if model is not None: + setattr(model, "max_seq_length", args_max_length) + return args + + _set_mlx_cuda_style_context_length(args, 1024) + return args + + def _apply_unsloth_trainer_mlx_defaults( + args, + model = None, + max_seq_length_explicit = False, + ): + if ( + not getattr(args, "streaming", False) + and not getattr(args, "preserve_dataset_order", False) + and not getattr(args, "_unsloth_mlx_dataset_order_explicit", False) + ): + default_order = getattr(MLXTrainingConfig, "dataset_order", "default") + if getattr(args, "dataset_order", default_order) in (None, default_order): + args.dataset_order = "torch_randperm" + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_append_eos_explicit", False + ): + args.append_eos = False + + if isinstance(args, UnslothTrainingArguments) and not getattr( + args, "_unsloth_mlx_grad_clip_explicit", False + ): + max_grad_norm = _positive_mlx_training_number( + getattr(args, "max_grad_norm", None), + ) + max_grad_value = _positive_mlx_training_number( + getattr(args, "max_grad_value", None), + ) + max_grad_leaf_norm = _positive_mlx_training_number( + getattr(args, "max_grad_leaf_norm", None), + ) + if max_grad_norm is None and max_grad_value is None and max_grad_leaf_norm is None: + args.max_grad_norm = 1.0 + + if not max_seq_length_explicit: + _resolve_mlx_cuda_style_max_seq_length(args, model = model) + return args + + def _coerce_mlx_training_args(args, overrides = None): + overrides = overrides or {} + if isinstance(args, MLXTrainingConfig) and not overrides: + return args + dataset_order_explicit = None + append_eos_explicit = None + max_seq_length_explicit = None + max_length_value = None + grad_clip_explicit = None + if args is None: + values = {} + elif isinstance(args, dict): + values = dict(args) + elif isinstance(args, (str, os.PathLike)): + values = {"output_dir": os.fspath(args)} + else: + dataset_order_explicit = getattr( + args, + "_unsloth_mlx_dataset_order_explicit", + False, + ) + append_eos_explicit = getattr( + args, + "_unsloth_mlx_append_eos_explicit", + None, + ) + max_seq_length_explicit = getattr( + args, + "_unsloth_mlx_max_seq_length_explicit", + None, + ) + if max_seq_length_explicit is None: + args_max_seq_length = _positive_mlx_context_length( + getattr(args, "max_seq_length", None), + ) + default_max_seq_length = getattr(MLXTrainingConfig, "max_seq_length", 2048) + max_seq_length_explicit = ( + args_max_seq_length is not None + and args_max_seq_length != default_max_seq_length + ) + max_length_value = getattr( + args, + "_unsloth_mlx_max_length_value", + getattr(args, "max_length", None), + ) + grad_clip_explicit = getattr( + args, + "_unsloth_mlx_grad_clip_explicit", + None, + ) + values = _mlx_training_argument_values(args) + if hasattr(args, "max_length"): + values["max_length"] = getattr(args, "max_length") + values.update(overrides) + coerced = UnslothTrainingArguments(**values) + if ( + dataset_order_explicit is not None + and "dataset_order" not in overrides + and "preserve_dataset_order" not in overrides + ): + coerced._unsloth_mlx_dataset_order_explicit = dataset_order_explicit + if append_eos_explicit is not None and "append_eos" not in overrides: + coerced._unsloth_mlx_append_eos_explicit = append_eos_explicit + if ( + max_seq_length_explicit is not None + and "max_seq_length" not in overrides + and "max_length" not in overrides + ): + coerced._unsloth_mlx_max_seq_length_explicit = max_seq_length_explicit + if max_length_value is not None and "max_length" not in overrides: + coerced._unsloth_mlx_max_length_value = max_length_value + coerced.max_length = max_length_value + if ( + grad_clip_explicit is not None + and "max_grad_norm" not in overrides + and "max_grad_value" not in overrides + and "max_grad_leaf_norm" not in overrides + ): + coerced._unsloth_mlx_grad_clip_explicit = grad_clip_explicit + return coerced + + _MLX_TRAINER_POSITIONAL_KWARGS = ( + "model", + "tokenizer", + "train_dataset", + "eval_dataset", + "dataset_text_field", + "max_seq_length", + "packing", + "data_collator", + "args", + "formatting_func", + "processor", + ) + _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( + "model", + "args", + "data_collator", + "train_dataset", + "eval_dataset", + "processing_class", + "compute_loss_func", + "compute_metrics", + "callbacks", + "optimizers", + "optimizer_cls_and_kwargs", + "preprocess_logits_for_metrics", + "peft_config", + "formatting_func", + ) + _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) + + def _is_mlx_native_text_collator(collator): + """HF pad/copy collators are redundant on MLX; match by class name.""" + for klass in type(collator).__mro__: + name = klass.__name__ + if name in ( + "DataCollatorForSeq2Seq", + "DataCollatorWithPadding", + "DefaultDataCollator", + ): + return True + if name == "DataCollatorForLanguageModeling": + # Plain causal padding is fine; MLM masking changes semantics. + return not bool(getattr(collator, "mlm", False)) + return False + + _MLX_VISION_COLLATOR_FORWARDED_KWARGS = frozenset( + ("completion_only_loss", "formatting_func", "max_seq_length") + ) + _MLX_VISION_COLLATOR_IMAGE_KWARGS = frozenset(("image_size", "resize")) + _MLX_VISION_COLLATOR_POSITIONAL_KWARGS = ( + "max_seq_length", + "formatting_func", + "resize", + "ignore_index", + "train_on_responses_only", + "instruction_part", + "response_part", + "force_match", + "num_proc", + "completion_only_loss", + "pad_to_multiple_of", + "resize_dimension", + "snap_to_patch_size", + "last_response_only", + ) + _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS = { + "ignore_index": -100, + "train_on_responses_only": False, + "instruction_part": None, + "response_part": None, + "force_match": True, + "num_proc": None, + "pad_to_multiple_of": None, + "resize_dimension": 0, + "snap_to_patch_size": False, + "last_response_only": False, + } + + def _is_default_mlx_vision_collator_value(key, value): + """Return whether an unsupported collator value is the CUDA default.""" + if key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS: + return False + default = _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS[key] + if default is None: + return value is None + if isinstance(default, bool): + return value is default + return value == default and type(value) is type(default) + + def _has_mlx_training_arg_value(args, key): + """Return whether training args already carry an explicit config value.""" + if args is None or isinstance(args, (str, os.PathLike)): + return False + if isinstance(args, dict): + return key in args + return getattr(args, key, None) is not None + + def _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs): + """Reject VLM collator kwargs that cannot be ignored safely on MLX.""" + unsupported = sorted( + key + for key, value in collator_kwargs.items() + if ( + key not in _MLX_VISION_COLLATOR_FORWARDED_KWARGS + and key not in _MLX_VISION_COLLATOR_IMAGE_KWARGS + and ( + ( + key in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and not _is_default_mlx_vision_collator_value(key, value) + ) + or ( + key not in _MLX_VISION_COLLATOR_UNSUPPORTED_DEFAULTS + and _is_meaningful_mlx_extra_value(value) + ) + ) + ) + ) + if unsupported: + raise NotImplementedError( + "Unsloth MLX: unsupported UnslothVisionDataCollator kwargs " + f"cannot be ignored safely: {', '.join(unsupported)}." + ) + + class UnslothTrainer(MLXTrainer): + """Backend-aware public trainer that routes supported SFT notebooks to MLX.""" + + def __init__(self, *args, **kwargs): + kwargs = _normalize_mlx_trainer_init_args(args, kwargs) + processing_class = kwargs.pop("processing_class", None) + processor_from_processing_class = False + if processing_class is not None: + if kwargs.get("processor", None) is None: + kwargs["processor"] = processing_class + processor_from_processing_class = True + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + processing_class, + "tokenizer", + processing_class, + ) + kwargs.setdefault("tokenizer", None) + + data_collator = kwargs.pop("data_collator", None) + if data_collator is not None: + if isinstance(data_collator, UnslothVisionDataCollator): + collator_processor = getattr(data_collator, "processor", None) + if collator_processor is not None and ( + kwargs.get("processor", None) is None or processor_from_processing_class + ): + kwargs["processor"] = collator_processor + if kwargs.get("tokenizer", None) is None: + kwargs["tokenizer"] = getattr( + collator_processor, + "tokenizer", + collator_processor, + ) + collator_kwargs = getattr(data_collator, "kwargs", None) or {} + collator_explicit_kwargs = getattr( + data_collator, + "_unsloth_mlx_explicit_kwargs", + set(collator_kwargs), + ) + collator_image_size = collator_kwargs.get( + "image_size", + collator_kwargs.get("resize", None), + ) + if isinstance(collator_image_size, list): + collator_image_size = tuple(collator_image_size) + if ( + isinstance(collator_image_size, str) + and collator_image_size.lower() == "max" + ): + collator_image_size = "max" + if "image_size" not in kwargs and ( + isinstance(collator_image_size, int) + or collator_image_size == "max" + or ( + isinstance(collator_image_size, tuple) + and len(collator_image_size) == 2 + and all(isinstance(x, int) for x in collator_image_size) + ) + ): + kwargs["image_size"] = collator_image_size + for collator_key in _MLX_VISION_COLLATOR_FORWARDED_KWARGS: + collator_defaulted_value = collator_key not in collator_explicit_kwargs + if collator_defaulted_value and _has_mlx_training_arg_value( + kwargs.get("args"), collator_key + ): + continue + if ( + collator_key in collator_kwargs + and collator_key not in kwargs + and collator_kwargs[collator_key] is not None + ): + kwargs[collator_key] = collator_kwargs[collator_key] + _raise_unsupported_mlx_vision_collator_kwargs(collator_kwargs) + elif _is_mlx_native_text_collator(data_collator): + pass # redundant on MLX; MLXTrainer batches/masks/pads natively + else: + raise NotImplementedError( + "Unsloth MLX: custom data_collator is not supported by " + "MLXTrainer. Pass the dataset directly or use the MLX " + "trainer's native batching path." + ) + + trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) + _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) + trainer_kwargs["args"] = _coerce_mlx_training_args( + trainer_kwargs.get("args"), + config_kwargs, + ) + if getattr( + trainer_kwargs["args"], "completion_only_loss", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: completion_only_loss=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + if getattr( + trainer_kwargs["args"], "train_on_completions", None + ) is True and not _is_vlm_model(trainer_kwargs.get("model")): + raise NotImplementedError( + "Unsloth MLX: train_on_completions=True is only supported " + "for VLM training. For text SFT, call train_on_responses_only " + "after constructing the trainer." + ) + trainer_kwargs["args"] = _apply_unsloth_trainer_mlx_defaults( + trainer_kwargs["args"], + model = trainer_kwargs.get("model"), + max_seq_length_explicit = (trainer_kwargs.get("max_seq_length") is not None), + ) + + super().__init__(**trainer_kwargs) + self.processing_class = ( + processing_class + if processing_class is not None + else self.processor or self.tokenizer + ) + if trainer_kwargs.get("max_seq_length") is not None: + _set_mlx_cuda_style_context_length( + self.args, + self.args.max_seq_length, + ) + self._unsloth_mlx_ignored_trainer_kwargs = ignored_kwargs + class UnslothVisionDataCollator: + def __init__( + self, + model = None, + processor = None, + *args, + **kwargs, + ): + explicit_kwargs = set(kwargs) + if len(args) > len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS): + raise TypeError( + "UnslothVisionDataCollator on MLX accepts at most " + f"{len(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS)} positional " + "options after model and processor." + ) + for key, value in zip(_MLX_VISION_COLLATOR_POSITIONAL_KWARGS, args): + if key in kwargs: + raise TypeError( + f"UnslothVisionDataCollator got multiple values for argument {key!r}" + ) + kwargs[key] = value + explicit_kwargs.add(key) + if "completion_only_loss" not in kwargs: + kwargs["completion_only_loss"] = True + self.model = model + self.processor = processor + self.args = () + self.kwargs = kwargs + self._unsloth_mlx_explicit_kwargs = explicit_kwargs + + def __call__(self, features): + raise NotImplementedError( + "Unsloth: UnslothVisionDataCollator is a compatibility placeholder " + "on MLX. Pass the dataset to UnslothTrainer; MLXTrainer performs " + "vision batching internally." + ) + + def get_chat_template(*args, **kwargs): + """Apply an Unsloth chat template through a lazy MLX-safe import.""" + from .chat_templates import get_chat_template as _get_chat_template + return _get_chat_template(*args, **kwargs) + + def apply_chat_template(*args, **kwargs): + """Format a dataset with an Unsloth chat template through a lazy import.""" + from .chat_templates import apply_chat_template as _apply_chat_template + return _apply_chat_template(*args, **kwargs) + + def standardize_data_formats(*args, **kwargs): + """Normalize ShareGPT-style datasets through the shared zoo helper.""" + from unsloth_zoo.dataset_utils import standardize_data_formats as _standardize_data_formats + return _standardize_data_formats(*args, **kwargs) + + def standardize_sharegpt(*args, **kwargs): + """Alias ShareGPT standardization to the shared dataset-format helper.""" + return standardize_data_formats(*args, **kwargs) + + def train_on_responses_only(*args, **kwargs): + """Mask non-response tokens through the shared zoo dataset helper.""" + from unsloth_zoo.dataset_utils import train_on_responses_only as _train_on_responses_only + return _train_on_responses_only(*args, **kwargs) + + def _safe_mlx_trl_star_exports(_trl): + """Return importable TRL star exports plus the MLX SFT shims.""" + exports = list(getattr(_trl, "__all__", ())) + safe_exports = [] + for name in exports: + try: + getattr(_trl, name) + except Exception: + continue + safe_exports.append(name) + for name in ("SFTConfig", "SFTTrainer"): + if name not in safe_exports: + safe_exports.append(name) + return safe_exports + + # trl trainers with no MLX implementation yet. Swap them for stubs that fail + # with a clear message instead of importing the real torch/CUDA trainer and + # crashing deep inside it, so an unmigrated GRPO/DPO/ORPO notebook is legible. + _MLX_UNSUPPORTED_TRL_TRAINERS = ( + "GRPOTrainer", + "DPOTrainer", + "ORPOTrainer", + "KTOTrainer", + "PPOTrainer", + "RewardTrainer", + ) + + def _make_mlx_unsupported_trl_trainer(name): def __init__(self, *args, **kwargs): raise NotImplementedError( - "Unsloth: UnslothVisionDataCollator is not used on MLX. " - "Use the MLX trainer/data path instead." + f"Unsloth: {name} is not yet supported on the MLX (Apple Silicon) " + f"backend. Only SFT training runs on MLX today; use a CUDA/ROCm GPU " + f"for {name}." ) + return type(name, (), {"__init__": __init__, "_unsloth_mlx_unsupported": True}) + + class _MLXSFTConfig(UnslothTrainingArguments): + """`trl.SFTConfig` alias that keeps TRL's default training length. + + TRL/HF SFTConfig defaults to num_train_epochs=3 (max_steps=-1); the + native MLX config defaults to max_steps=60. An unmigrated notebook that + builds SFTConfig without an explicit length would otherwise silently run + 60 MLX steps under this alias, so seed the TRL epoch default when neither + max_steps nor num_train_epochs is given (epoch mode is MLX-supported). + """ + + def __init__(self, *args, **kwargs): + keys = set(kwargs) + if len(args) == 1 and isinstance(args[0], dict): + keys |= set(args[0]) + if not ({"max_steps", "num_train_epochs"} & keys): + kwargs.setdefault("num_train_epochs", 3) + super().__init__(*args, **kwargs) + + def _install_mlx_trl_sft_shim(): + """Install MLX-backed TRL SFT shims without replacing the TRL module.""" + _trl = _sys.modules.get("trl") + if _trl is None: + try: + import trl as _trl + except ImportError: + _trl = _types.ModuleType("trl") + _trl.__version__ = "0.0.0+unsloth-mlx" + _trl.__package__ = "trl" + _trl.__path__ = [] + _trl.__spec__ = _machinery.ModuleSpec("trl", loader = None, is_package = True) + _sys.modules["trl"] = _trl + + _trl.SFTTrainer = UnslothTrainer + _trl.SFTConfig = _MLXSFTConfig + # Only retarget trainers the installed trl actually exposes (don't invent + # attributes); idempotent so re-importing unsloth is a no-op. + # Decide what to stub from trl's declared exports (__all__) and already + # materialized attrs only. A getattr probe here would trigger trl's lazy + # trainer import, pulling torch and breaking `import unsloth` on torch-free + # MLX just to check existence. + _trl_exports = set(getattr(_trl, "__all__", ()) or ()) + # Stub every non-SFT trainer trl exposes, not just a fixed list, so newer + # trainers (RLOOTrainer, ...) also fail with a clear MLX message instead + # of importing the real torch trainer. Names come from __all__ so we never + # resolve them (that would trigger trl's lazy import and pull torch). + _unsupported = set(_MLX_UNSUPPORTED_TRL_TRAINERS) | { + _n for _n in _trl_exports if _n.endswith("Trainer") and _n != "SFTTrainer" + } + for _name in _unsupported: + _current = vars(_trl).get(_name) + if getattr(_current, "_unsloth_mlx_unsupported", False): + continue + if _name in _trl_exports or _current is not None: + setattr(_trl, _name, _make_mlx_unsupported_trl_trainer(_name)) + _trl.__all__ = _safe_mlx_trl_star_exports(_trl) + _trl.__UNSLOTH_MLX_COMPAT__ = True + + def _install_mlx_unsloth_trainer_shim(): + module_name = f"{__name__}.trainer" + _trainer = _types.ModuleType(module_name) + _trainer.__package__ = __name__ + _trainer.__spec__ = _machinery.ModuleSpec(module_name, loader = None) + _trainer.MLXTrainer = MLXTrainer + _trainer.MLXTrainingConfig = MLXTrainingConfig + _trainer.UnslothTrainer = UnslothTrainer + _trainer.UnslothTrainingArguments = UnslothTrainingArguments + _trainer.UnslothVisionDataCollator = UnslothVisionDataCollator + _sys.modules[module_name] = _trainer + globals()["trainer"] = _trainer + + _install_mlx_trl_sft_shim() + _install_mlx_unsloth_trainer_shim() + else: # GPU path: load everything from _gpu_init from ._gpu_init import * from ._gpu_init import __version__ + + def get_gpu_memory_stats(): + """Return CUDA/ROCm/XPU device stats, peak memory, and total memory in GiB.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + props = torch.xpu.get_device_properties(0) + peak = ( + torch.xpu.max_memory_reserved() + if hasattr(torch.xpu, "max_memory_reserved") + else torch.xpu.max_memory_allocated() + ) + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + if hasattr(torch, "cuda") and torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + peak = torch.cuda.max_memory_reserved() + total = getattr(props, "total_memory", 0) + return props, _bytes_to_gb(peak), _bytes_to_gb(total) or 1.0 + except Exception: + pass + stats = _UnslothDeviceStats("Unknown GPU", 0) + return stats, 0.0, 1.0 + + def clear_gpu_memory(): + """Clear cached GPU memory on CUDA, ROCm, or XPU when available.""" + try: + import torch + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() + elif hasattr(torch, "cuda") and torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass diff --git a/unsloth/_compressed_quantize.py b/unsloth/_compressed_quantize.py index 66ebdd75da..f0a843c380 100644 --- a/unsloth/_compressed_quantize.py +++ b/unsloth/_compressed_quantize.py @@ -245,6 +245,10 @@ def main(): # expert even if the sample set does not route tokens to all of them. is_moe = _is_moe(getattr(model, "config", None)) ignore = ["lm_head"] + # Skip the same modules RedHatAI/NVIDIA skip for the Qwen3.5 / Qwen3-Next family (these also have + # shapes not divisible by the grouped-scheme group_size, which would otherwise error). No-ops + # elsewhere. Hybrid linear attention, VLM vision tower, and the MTP/speculative head. + ignore += ["re:.*\\.linear_attn\\..*", "re:.*\\.visual\\..*", "re:.*mtp.*"] if is_moe: # Keep MoE routing layers unquantized: the router gate and (Qwen) shared-expert gate. ignore += ["re:.*\\.gate$", "re:.*\\.shared_expert_gate$"] diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 60eb8de5d7..169b2dbd0e 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -27,18 +27,25 @@ __all__ = [ "test_construct_chat_template", ] -from transformers import StoppingCriteria, StoppingCriteriaList -from torch import LongTensor, FloatTensor -from transformers.models.llama.modeling_llama import logger +from transformers.utils import logging +try: + from torch import LongTensor, FloatTensor +except ImportError: + LongTensor = FloatTensor = None +logger = logging.get_logger(__name__) import os import shutil -from .tokenizer_utils import * import re from .ollama_template_mappers import OLLAMA_TEMPLATES -from unsloth_zoo.dataset_utils import ( - train_on_responses_only, - standardize_data_formats, -) +try: + from unsloth_zoo.dataset_utils import ( + train_on_responses_only, + standardize_data_formats, + ) +except ImportError: + # dataset_utils pulls torch; keep chat_templates importable on torch-free + # (MLX) hosts, which expose these via the backend-specific wrappers instead. + train_on_responses_only = standardize_data_formats = None standardize_sharegpt = standardize_data_formats CHAT_TEMPLATES = {} DEFAULT_SYSTEM_MESSAGE = {} @@ -1838,11 +1845,24 @@ def get_chat_template( map_eos_token = True, system_message = None, patch_saving = True, - use_zoo_tokenizer_patch = False, + use_zoo_tokenizer_patch = None, ): assert(type(map_eos_token) is bool) + import sys + is_mlx_backend = getattr(sys.modules.get("unsloth"), "DEVICE_TYPE", None) == "mlx" + if use_zoo_tokenizer_patch is None: + use_zoo_tokenizer_patch = is_mlx_backend old_tokenizer = tokenizer + # mlx-lm's TokenizerWrapper._tokenizer is the HF tokenizer, not the Rust + # backend the vocab-edit paths below need; unwrap here, re-wrap before return. + _mlx_tokenizer_wrapper = None + if is_mlx_backend and tokenizer.__class__.__name__ == "TokenizerWrapper": + _inner_tokenizer = getattr(tokenizer, "_tokenizer", None) + if _inner_tokenizer is not None and hasattr(_inner_tokenizer, "is_fast"): + _mlx_tokenizer_wrapper = tokenizer + tokenizer = _inner_tokenizer + IS_GEMMA = False if tokenizer.__class__.__name__.startswith("Gemma"): if chat_template == "chatml": chat_template = "gemma_chatml" @@ -1952,6 +1972,7 @@ def get_chat_template( pass # Must fix the sentence piece tokenizer since there's no tokenizer.model file! + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) else: pass @@ -1997,6 +2018,7 @@ def get_chat_template( # Must fix the sentence piece tokenizer since there's no tokenizer.model file! token_mapping = { old_eos_token : stop_word, } + from .tokenizer_utils import fix_sentencepiece_tokenizer tokenizer = fix_sentencepiece_tokenizer(tokenizer, new_tokenizer, token_mapping,) pass @@ -2057,13 +2079,25 @@ def get_chat_template( # stopping_criteria = create_stopping_criteria(tokenizer, stop_word) # Patch saving functions - if patch_saving: + if patch_saving and not is_mlx_backend: from .save import patch_saving_functions tokenizer = patch_saving_functions(tokenizer) # Add Ollama tokenizer._ollama_modelfile = ollama_modelfile tokenizer._system_message = system_message + + # Re-wrap so the trainer gets the same TokenizerWrapper type back. + if _mlx_tokenizer_wrapper is not None: + _mlx_tokenizer_wrapper._tokenizer = tokenizer + eos_token_id = getattr(tokenizer, "eos_token_id", None) + if eos_token_id is not None: + _mlx_tokenizer_wrapper._eos_token_ids = {eos_token_id} + _mlx_tokenizer_wrapper._chat_template = None + _mlx_tokenizer_wrapper.has_chat_template = ( + getattr(tokenizer, "chat_template", None) is not None + ) + tokenizer = _mlx_tokenizer_wrapper return tokenizer#, stopping_criteria @@ -2749,6 +2783,15 @@ extra_eos_tokens = None, def create_stopping_criteria(tokenizer, stop_word = "eos_token"): + try: + import torch + from transformers import StoppingCriteria, StoppingCriteriaList + except ImportError as exc: + raise ImportError( + "Unsloth: create_stopping_criteria requires PyTorch and is only " + "supported on Torch backends." + ) from exc + class StoppingCriteriaSub(StoppingCriteria): __slots__ = "stop_token", "single_match", "length", @@ -2828,10 +2871,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) @@ -2845,10 +2888,10 @@ def test_chat_templates(): for j in range(len(messages)-1): correct_prompt.append_message(correct_prompt.roles[j%2==1], messages[j+1]["content"]) correct_prompt.append_message(correct_prompt.roles[1], "") - correct_prompt = tokenizer.bos_token + correct_prompt.get_prompt() template = vicuna_old_template correct_tokenizer = AutoTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5") + correct_prompt = correct_tokenizer.bos_token + correct_prompt.get_prompt() correct_tokenizer.chat_template = template our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) # We add ourselves diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 599a5c0262..047783c35e 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -3602,8 +3602,27 @@ def make_fast_generate_wrapper(original_generate): @functools.wraps(original_generate) def _fast_generate_wrapper(*args, **kwargs): - # Check for vLLM-specific arguments - if "sampling_params" in kwargs: + def _has_sampling_params(a): + # SamplingParams passed directly or inside a positional list/tuple + return type(a).__name__ == "SamplingParams" or ( + isinstance(a, (list, tuple)) + and any(type(i).__name__ == "SamplingParams" for i in a) + ) + + def _is_vllm_prompt(a): + # str prompt, a vLLM prompt dict (prompt / prompt_token_ids / prompt_embeds / + # multi_modal_data), or a list/tuple of those + head = a[0] if isinstance(a, (list, tuple)) and len(a) > 0 else a + return isinstance(head, str) or ( + isinstance(head, dict) + and any( + k in head + for k in ("prompt", "prompt_token_ids", "prompt_embeds", "multi_modal_data") + ) + ) + + # vLLM-only; also catch SamplingParams passed positionally (fast_generate(prompt, params)) + if "sampling_params" in kwargs or any(_has_sampling_params(a) for a in args): raise ValueError( "Unsloth: `sampling_params` is only supported when `fast_inference=True` (vLLM). " "Since `fast_inference=False`, use HuggingFace generate arguments instead:\n" @@ -3616,33 +3635,26 @@ def make_fast_generate_wrapper(original_generate): "Since `fast_inference=False`, LoRA weights are already merged into the model." ) - # Check if first positional argument is a string or list of strings - if len(args) > 0: - first_arg = args[0] - is_string_input = False - - if isinstance(first_arg, str): - is_string_input = True - elif isinstance(first_arg, (list, tuple)) and len(first_arg) > 0: - if isinstance(first_arg[0], str): - is_string_input = True - - if is_string_input: - raise ValueError( - "Unsloth: Passing text strings to `fast_generate` is only supported " - "when `fast_inference=True` (vLLM). Since `fast_inference=False`, you must " - "tokenize the input first:\n\n" - " messages = tokenizer.apply_chat_template(\n" - ' [{"role": "user", "content": "Your prompt here"}],\n' - " tokenize=True, add_generation_prompt=True,\n" - ' return_tensors="pt", return_dict=True\n' - " )\n" - " output = model.fast_generate(\n" - " **messages.to('cuda'),\n" - " max_new_tokens=64,\n" - " temperature=1.0,\n" - " )" - ) + # A vLLM-style prompt (string, {"prompt":..., "multi_modal_data":...} dict, or a list/tuple + # of either) only works under vLLM; tokenize first when fast_inference=False. A positional + # arg may be HF token ids, so check it conservatively with _is_vllm_prompt. The `prompts` / + # `prompt_token_ids` / `prompt_embeds` keywords are vLLM-only names that HuggingFace generate + # does not accept, so any of them being present is a vLLM-style call (even a bare token list, + # or an explicit None from a defaulted kwargs dict), hence membership rather than a value check. + vllm_prompt_kwarg = any( + k in kwargs for k in ("prompts", "prompt_token_ids", "prompt_embeds") + ) + if (len(args) > 0 and _is_vllm_prompt(args[0])) or vllm_prompt_kwarg: + raise ValueError( + "Unsloth: Passing vLLM-style prompts to `fast_generate` is only supported when " + "`fast_inference=True` (vLLM). Since `fast_inference=False`, tokenize first:\n\n" + " inputs = tokenizer.apply_chat_template(\n" + ' [{"role": "user", "content": "Your prompt here"}],\n' + " tokenize=True, add_generation_prompt=True,\n" + ' return_tensors="pt", return_dict=True,\n' + " )\n" + " output = model.fast_generate(**inputs.to('cuda'), max_new_tokens=64, temperature=1.0)" + ) # Call original generate return original_generate(*args, **kwargs) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 14ee5ee24e..bb7289dfa8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3047,6 +3047,9 @@ class FastLlamaModel: # Pre-wrapped PEFT model passes through here; still arm the detector so an RL # trainer can reset a compile cache poisoned by a pre-train forward. _unsloth_install_pretrain_detector(model) + # This branch returns before patch_peft_model, so record here too; + # apply_unsloth_gradient_checkpointing above already re-patched global state to match (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing model = _exclude_rope_inv_freq_from_ddp(model) return model else: @@ -3406,6 +3409,11 @@ class FastLlamaModel: @staticmethod def patch_peft_model(model, use_gradient_checkpointing = "unsloth"): + # Persist the effective GC mode so the trainer restores it verbatim: for_inference() + # clears the module flags every GRPO step, and a plain TrainingArguments defaults it to + # False, which would otherwise silently disable it at train time (#4735). Recorded here, + # not in get_peft_model, so adapters loaded via loader.py's from_pretrained path are covered. + model._unsloth_gradient_checkpointing = use_gradient_checkpointing if os.environ.get("UNSLOTH_USE_NEW_MODEL", "0") == "1": return FastBaseModel.patch_peft_model( model = model, diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 53668d14d8..62ef9e916a 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -423,8 +423,14 @@ def prepare_for_training_mode(f): pass # Enable training mode _was_training = None - # Get gradient checkpointing setting from training arguments - use_gc = getattr(self.args, 'gradient_checkpointing', True) + # Restore the GC mode the model was configured with at setup; fall back to + # the training args only when it wasn't recorded (issue #4735). Use hasattr, + # not a None sentinel, so a deliberately-recorded None is restored verbatim. + _model = getattr(self, 'model', None) + if hasattr(_model, '_unsloth_gradient_checkpointing'): + use_gc = _model._unsloth_gradient_checkpointing + else: + use_gc = getattr(self.args, 'gradient_checkpointing', True) if hasattr(self, 'model') and hasattr(self.model, "training"): _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): @@ -532,7 +538,8 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): if getattr(args, "_n_gpu", 1) != 1: args._n_gpu = 1 if "model" in locals() and hasattr(model, "for_training"): - model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True)) + _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True) + model.for_training(use_gradient_checkpointing=_use_gc) super().__init__({RLTrainer_call_args}{RLTrainer_kwargs}) if "model" in locals() and hasattr(model, "for_inference"): model.for_inference() @@ -1015,6 +1022,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): "dtype = _get_dtype(dtype)\n" "float16 = dtype == torch.float16\n" "bfloat16 = dtype == torch.bfloat16\n" + "if full_finetuning:\n" + " if bfloat16 and use_fp16: use_fp16 = False\n" + " if float16 and use_bf16: use_bf16 = False\n" "if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n" "if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n" "if force_float32:\n" @@ -1162,7 +1172,8 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): if "model" in call_args: training_check = ( "if model is not None and hasattr(model, 'for_training'):\n" - " model.for_training(use_gradient_checkpointing=getattr(args, 'gradient_checkpointing', True))\n" + " _use_gc = model._unsloth_gradient_checkpointing if hasattr(model, '_unsloth_gradient_checkpointing') else getattr(args, 'gradient_checkpointing', True)\n" + " model.for_training(use_gradient_checkpointing=_use_gc)\n" "if 'tokenizer' in locals() and hasattr(tokenizer, 'padding_side'): tokenizer.padding_side = 'right'\n" "if 'processing_class' in locals():\n" " if hasattr(processing_class, 'padding_side'): processing_class.padding_side = 'right'\n" diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 61a07b686d..3be614cf4a 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -761,7 +761,8 @@ def grpo_trainer__generate_and_score_completions(function_name, function): # Left pad prompt before calculation old and ref hidden states left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id) max_left_pad = torch.max(left_pad_tokens_per_prompt).item() - self.model.for_training(use_gradient_checkpointing=getattr(self.args, 'gradient_checkpointing', True))""" + _use_gc = self.model._unsloth_gradient_checkpointing if hasattr(self.model, '_unsloth_gradient_checkpointing') else getattr(self.args, 'gradient_checkpointing', True) + self.model.for_training(use_gradient_checkpointing=_use_gc)""" function = function.replace(line_to_replace, replacement_lines) @@ -1337,11 +1338,12 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_sizes_chunks.append(slice_sample_axis(image_sizes, start, end)) temperature = self.temperature - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) - logit_scale_multiply = getattr(model.config, "logit_scale", 0) + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) + logit_scale_multiply = getattr(model_config, "logit_scale", 0) if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) + logit_scale_divide = getattr(model_config, "logits_scaling", 0) if logit_scale_divide is None: logit_scale_divide = 0 @@ -1471,7 +1473,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__get_per_token_logps_and_entropies) -def _unsloth_get_final_logit_softcapping(config): +def _unsloth_get_model_config(model): + """Return HuggingFace model config, unwrapping DDP/Accelerate wrappers.""" + config = getattr(model, "config", None) + if config is None and hasattr(model, "module"): + config = getattr(model.module, "config", None) + return config + + +def _unsloth_get_final_logit_softcapping(model): """Return final_logit_softcapping for a model config, falling back to the nested text sub-config for composite models. Handles both: - Gemma-4-style configs where the attribute lives on ``config.text_config`` @@ -1479,6 +1489,9 @@ def _unsloth_get_final_logit_softcapping(config): reachable via ``config.get_text_config()`` Returns 0 if unset, matching the previous behaviour. """ + config = _unsloth_get_model_config(model) + if config is None: + return 0 softcap = getattr(config, "final_logit_softcapping", None) if softcap is None: text_cfg = getattr(config, "text_config", None) @@ -1499,6 +1512,7 @@ grpo_compute_loss_slow = RL_REPLACEMENTS["grpo_compute_loss_slow"] UnslothEfficientGRPO = RL_REPLACEMENTS["UnslothEfficientGRPO"] grpo_accumulated_loss = RL_REPLACEMENTS["grpo_accumulated_loss"] grpo_update_SamplingParams = RL_REPLACEMENTS["grpo_update_SamplingParams"] +RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_model_config)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_final_logit_softcapping)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_get_mm_token_id)) RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_fix_mm_token_type_ids)) @@ -1616,11 +1630,12 @@ def grpo_trainer_compute_loss(function_name, function): input_ids = input_ids[:, -logits_to_keep:] # Get logit softcapping and logit scale - logit_softcapping = _unsloth_get_final_logit_softcapping(model.config) # Gemma - logit_scale_multiply = getattr(model.config, "logit_scale", 0) # Cohere + model_config = _unsloth_get_model_config(model) + logit_softcapping = _unsloth_get_final_logit_softcapping(model) # Gemma + logit_scale_multiply = getattr(model_config, "logit_scale", 0) # Cohere if logit_scale_multiply is None: logit_scale_multiply = 0 - logit_scale_divide = getattr(model.config, "logits_scaling", 0) # Granite + logit_scale_divide = getattr(model_config, "logits_scaling", 0) # Granite if logit_scale_divide is None: logit_scale_divide = 0 diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index bdc2bd9ef6..689e362f95 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -238,6 +238,71 @@ def _attach_bnb_multidevice_hooks( global NUM_LOGITS_TO_KEEP NUM_LOGITS_TO_KEEP = dict() + +def _unsloth_generate_accepts_kwarg(model, key): + # True if the top level accepts this generate kwarg (some models expose it on an inner forward only). + try: + model_args = set(inspect.signature(model.prepare_inputs_for_generation).parameters) + except (TypeError, ValueError, AttributeError): + model_args = set() + if "kwargs" in model_args or "model_kwargs" in model_args: + try: + model_args |= set(inspect.signature(model.forward).parameters) + except (TypeError, ValueError, AttributeError): + pass + return key in model_args + + +def _install_offload_embedding_hooks(embed_tokens, output_embeddings, return_device): + # Lookup runs on the weight's current device (CPU when offloaded); the output returns to the + # decoder device read live from output_embeddings (lm_head, untied here) so it tracks + # model.to() moves. A meta (disk-offloaded) or missing lm_head falls back to return_device. + if embed_tokens is None: + return False + if getattr(embed_tokens, "_unsloth_offload_hooks_installed", False): + return True + + def _decoder_device(): + weight = getattr(output_embeddings, "weight", None) + if weight is not None and weight.device.type != "meta": + return weight.device + return return_device + + def _unsloth_offload_pre_hook(module, args): + if not args: + return args + inp = args[0] + if not hasattr(inp, "device"): + return args + weight = getattr(module, "weight", None) + target = weight.device if weight is not None else _decoder_device() + if target is None or inp.device == target: + return args + return (inp.to(target),) + tuple(args[1:]) + + def _unsloth_offload_post_hook(module, args, output): + target = _decoder_device() + if target is not None and hasattr(output, "device") and output.device != target: + return output.to(target) + return output + + embed_tokens.register_forward_pre_hook(_unsloth_offload_pre_hook, prepend = True) + embed_tokens.register_forward_hook(_unsloth_offload_post_hook, prepend = True) + embed_tokens._unsloth_offload_hooks_installed = True + return True + + +def _embeddings_are_tied(input_embeddings, output_embeddings): + # Tied lm_head reuses this weight; offloading to CPU would strand the output projection. + if input_embeddings is None or output_embeddings is None: + return False + w_in = getattr(input_embeddings, "weight", None) + w_out = getattr(output_embeddings, "weight", None) + if w_in is None or w_out is None: + return False + return w_in is w_out or w_in.data_ptr() == w_out.data_ptr() + + VLLM_SUPPORTED_VLM = [ "qwen2_5_vl", "gemma3", @@ -321,6 +386,13 @@ def unsloth_base_fast_generate(self, *args, **kwargs): kwargs.pop("token_type_ids", None) # kwargs.pop("token_type_ids", None) + # Vision processors emit mm_token_type_ids that generate() rejects (Qwen3-VL); unlike + # logits_to_keep it is an incoming kwarg, so drop it when generate does not accept it. + if "mm_token_type_ids" in kwargs and not _unsloth_generate_accepts_kwarg( + self, "mm_token_type_ids" + ): + kwargs.pop("mm_token_type_ids", None) + # VLMs do not allow logits_to_keep global NUM_LOGITS_TO_KEEP if arch not in NUM_LOGITS_TO_KEEP: @@ -339,7 +411,7 @@ def unsloth_base_fast_generate(self, *args, **kwargs): if arch not in NUM_LOGITS_TO_KEEP: NUM_LOGITS_TO_KEEP[arch] = None key = NUM_LOGITS_TO_KEEP[arch] - if key is not None and key not in kwargs: + if key is not None and key not in kwargs and _unsloth_generate_accepts_kwarg(self, key): kwargs[key] = 1 model_eos_token_id = getattr(self.config, "eos_token_id", None) @@ -1024,6 +1096,12 @@ class FastBaseModel: raise_handler = RaiseUninitialized() try: + if offload_embedding and fast_inference: + # vLLM manages its own weights; embedding offload does not apply. + print( + "Unsloth: Not offloading embeddings; incompatible with fast_inference (vLLM)." + ) + offload_embedding = False if not fast_inference: # Prevent load_in_fp8 from being forwarded into HF internal model loading load_in_fp8 = kwargs.pop("load_in_fp8", None) @@ -1070,21 +1148,26 @@ class FastBaseModel: pass else: embed_tokens = model.get_input_embeddings() + out_embed = ( + model.get_output_embeddings() + if hasattr(model, "get_output_embeddings") + else None + ) + if _embeddings_are_tied(embed_tokens, out_embed): + raise NotImplementedError( + "offload_embedding = True is not supported for models with tied word " + "embeddings (embed_tokens shares its weight with lm_head). Offloading " + "would strand the output projection on CPU and saves no VRAM. Set " + "offload_embedding = False for this model." + ) nbytes = embed_tokens.weight.numel() * embed_tokens.weight.itemsize ngb = round(nbytes / 1024 / 1024 / 1024, 2) print(f"Unsloth: Offloading embeddings to RAM to save {ngb} GB.") + _embed_device = embed_tokens.weight.device # decoder device, before offload embed_tokens.to("cpu") - # Add hooks to move inputs to CPU and back to CUDA - # [TODO] Doesn't seem to work! - # def pre_hook(module, args): - # args[0]._old_device = args[0].device - # return (args[0].to("cpu", non_blocking = True)) - # def post_hook(module, args, output): - # old_device = getattr(args[0], "_old_device", "cuda") - # return output.to(old_device, non_blocking = True) - # embed_tokens.register_forward_pre_hook(pre_hook, prepend = True) - # embed_tokens.register_forward_hook (post_hook, prepend = True) + # Device-safe embedding offload. + _install_offload_embedding_hooks(embed_tokens, out_embed, _embed_device) # Must free GPU memory otherwise will not free! torch.cuda.empty_cache() gc.collect() @@ -1790,6 +1873,11 @@ class FastBaseModel: float32_mixed_precision = float32_mixed_precision, patch_modules_to_save = True, ) + # Persist the configured GC mode so the trainer restores it verbatim. + # for_inference() clears the module flags (GRPO does this every generation + # step), and a plain TrainingArguments defaults gradient_checkpointing=False, + # which would otherwise silently disable this setting at train time (#4735). + model._unsloth_gradient_checkpointing = use_gradient_checkpointing # Gemma3N audio conformer processes variable-length audio tensors # that cause stride mismatches in AOT autograd compiled backward diff --git a/unsloth/save.py b/unsloth/save.py index 226ca5fed8..50ae4119bd 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -205,6 +205,24 @@ COMPRESSED_EXPORT_SCHEMES = { } +# torchao "portable" quant export: device-agnostic FP8 / INT8, no NVIDIA GPU needed. +# alias -> (kind, sibling suffix). FP8 saves to safetensors, INT8 to .bin; both load in vLLM. +TORCHAO_EXPORT_SCHEMES = { + "torchao_fp8": ("fp8", "torchao-fp8"), + "torchao_int8": ("int8", "torchao-int8"), + "portable_fp8": ("fp8", "torchao-fp8"), + "portable_int8": ("int8", "torchao-int8"), +} + + +def _normalize_torchao_method(save_method): + """Return (kind, suffix) if `save_method` is a torchao portable FP8/INT8 export, else None.""" + if not isinstance(save_method, str): + return None + key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + return TORCHAO_EXPORT_SCHEMES.get(key) + + def _normalize_compressed_method(save_method): """Return (scheme, needs_calibration, suffix) if `save_method` is an FP8/FP4 compressed export, else None (so normal lora / merged_16bit / merged_4bit handling proceeds). @@ -215,6 +233,9 @@ def _normalize_compressed_method(save_method): if not isinstance(save_method, str): return None key = save_method.lower().strip().replace("-", "_").replace(" ", "_") + # torchao aliases route to the torchao path, so skip them before the "fp8" near-miss check. + if key in TORCHAO_EXPORT_SCHEMES: + return None if key in COMPRESSED_EXPORT_SCHEMES: return COMPRESSED_EXPORT_SCHEMES[key] if any(tag in key for tag in ("fp8", "fp4", "mxfp", "nvfp", "w4a", "w8a", "int4", "int8")): @@ -1368,6 +1389,53 @@ def install_python_non_blocking(packages = []): # bump deliberately. Floor 0.6.0 keeps torch>=2.4 resolvable (0.7+ need torch>=2.7; torch pinned below). _LLM_COMPRESSOR_SPEC = "llmcompressor>=0.6.0,<=0.12.0" +# Highest transformers release llm-compressor 0.10.x/0.12.x can run against (its metadata pins +# transformers<=4.57.6). Models that require a newer-transformers sidecar (e.g. Qwen3.5 needs +# transformers 5.3.0) cannot be quantized by llm-compressor at all: it imports +# transformers.modeling_utils.TORCH_INIT_FUNCTIONS, which was removed in transformers 5.x, so the +# compressed-export subprocess dies with a cryptic ImportError AFTER the expensive 16bit merge. +# Detect that up front and fail fast with an actionable message. Bump this in lockstep with a +# llm-compressor release that supports newer transformers. +_LLM_COMPRESSOR_MAX_TRANSFORMERS = "4.57.6" + + +def _transformers_exceeds_llm_compressor_ceiling(transformers_version = None): + """Return (exceeds, active_version) comparing the active transformers to the llm-compressor ceiling. + + `exceeds` is True only when we can parse both versions and the active transformers is strictly + newer than `_LLM_COMPRESSOR_MAX_TRANSFORMERS`. Any parse failure returns False (fail open) so a + real quantization attempt still surfaces the underlying error rather than a false positive. + """ + if transformers_version is None: + try: + import transformers as _tf + transformers_version = _tf.__version__ + except Exception: + return False, "unknown" + try: + from packaging.version import parse as _parse + + # Drop any local build suffix ("4.57.6+abc") so it does not skew the comparison. + active = _parse(str(transformers_version).split("+", 1)[0]) + ceiling = _parse(_LLM_COMPRESSOR_MAX_TRANSFORMERS) + return active > ceiling, str(transformers_version) + except Exception: + return False, str(transformers_version) + + +# A caller (e.g. Unsloth Studio) can enable FP8/FP4 export of newer-transformers models (Qwen3.5, +# Gemma-4, ...) by provisioning a dedicated llm-compressor-main "shadow" (transformers>=5.9 layered +# over the existing torch) and pointing us at its sys.path entry via this env var. When set, the +# quantization subprocess uses it instead of the workspace llm-compressor and the ceiling fail-fast +# is bypassed. +_COMPRESSED_QUANTIZE_PYTHONPATH_ENV = "UNSLOTH_COMPRESSED_QUANTIZE_PYTHONPATH" + + +def _compressed_quantize_pythonpath(): + """Return the llm-compressor-main shadow PYTHONPATH, or None if not set.""" + pp = os.environ.get(_COMPRESSED_QUANTIZE_PYTHONPATH_ENV, "").strip() + return pp or None + def install_llm_compressor(): """Import llm-compressor, installing it on first use for FP8/FP4 export. @@ -2023,10 +2091,40 @@ def unsloth_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -2107,6 +2205,37 @@ def unsloth_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -2114,6 +2243,7 @@ def unsloth_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3812,10 +3942,40 @@ def unsloth_generic_save_pretrained_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = save_directory, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = push_to_hub, + token = token, + is_main_process = is_main_process, + # Forward standard save kwargs to the 16bit merge. + state_dict = state_dict, + save_function = save_function, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + variant = variant, + save_peft_format = save_peft_format, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self del arguments["self"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -3896,6 +4056,37 @@ def unsloth_generic_push_to_hub_merged( gc.collect() return + # torchao portable FP8/INT8 export (no NVIDIA GPU) -> separate path. + _torchao = _normalize_torchao_method(save_method) + if _torchao is not None: + kind, suffix = _torchao + _unsloth_save_torchao( + model = self, + save_directory = repo_id, + tokenizer = tokenizer, + kind = kind, + suffix = suffix, + push_to_hub = True, + token = token, + is_main_process = True, + private = private, + commit_message = commit_message, + commit_description = commit_description, + create_pr = create_pr, + revision = revision, + # Forward standard save kwargs to the 16bit merge. + use_temp_dir = use_temp_dir, + max_shard_size = max_shard_size, + safe_serialization = safe_serialization, + tags = tags, + temporary_location = temporary_location, + maximum_memory_usage = maximum_memory_usage, + datasets = datasets, + ) + for _ in range(3): + gc.collect() + return + arguments = dict(locals()) arguments["model"] = self arguments["save_directory"] = repo_id @@ -3903,6 +4094,7 @@ def unsloth_generic_push_to_hub_merged( del arguments["self"] del arguments["repo_id"] del arguments["_compressed"] + del arguments["_torchao"] del arguments["calibration_dataset"] del arguments["num_calibration_samples"] del arguments["max_seq_length"] @@ -4114,22 +4306,36 @@ def _unsloth_save_compressed_tensors( if not is_main_process: return None - # 1) Install llm-compressor and gate on scheme availability BEFORE merging, so an unsupported - # scheme (e.g. mxfp8) fails fast instead of writing a full 16bit checkpoint first. - install_llm_compressor() - if not _scheme_is_available(scheme): - try: - import transformers as _tf - tf_ver = _tf.__version__ - except Exception: - tf_ver = "unknown" - raise RuntimeError( - f"Unsloth: scheme '{scheme}' is not available in your installed " - f"compressed-tensors / llm-compressor.\n" - f"It requires a newer llm-compressor that needs transformers>=5.9 " - f"(you have transformers {tf_ver}).\n" - "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." - ) + # 1) Prepare the quantization runtime BEFORE merging, so an unusable config fails fast instead of + # writing a full 16bit checkpoint first. With the llm-compressor-main shadow the subprocess + # validates everything itself, so skip the workspace install / ceiling / scheme checks; without + # it, install the workspace llm-compressor and fail fast past its transformers ceiling. + _shadow_pythonpath = _compressed_quantize_pythonpath() + if _shadow_pythonpath is None: + install_llm_compressor() + # llm-compressor cannot run under a newer transformers than its ceiling: the quantization + # subprocess would die with a cryptic ImportError (TORCH_INIT_FUNCTIONS) only AFTER the costly + # 16bit merge. Detect and fail fast with an actionable message instead. + _exceeds, _tf_ver = _transformers_exceeds_llm_compressor_ceiling() + if _exceeds: + raise RuntimeError( + f"Unsloth: FP8/FP4 compressed-tensors export is not available for this model. It runs " + f"under transformers {_tf_ver}, but llm-compressor supports transformers " + f"<= {_LLM_COMPRESSOR_MAX_TRANSFORMERS}. Export to GGUF or 16-bit instead." + ) + if not _scheme_is_available(scheme): + try: + import transformers as _tf + tf_ver = _tf.__version__ + except Exception: + tf_ver = "unknown" + raise RuntimeError( + f"Unsloth: scheme '{scheme}' is not available in your installed " + f"compressed-tensors / llm-compressor.\n" + f"It requires a newer llm-compressor that needs transformers>=5.9 " + f"(you have transformers {tf_ver}).\n" + "Use save_method in {fp8, mxfp4, nvfp4}, or upgrade transformers + llm-compressor." + ) # 2) Pick the local working dir. For a hub push, save_directory is a repo id, so merge and # quantize inside an isolated temp dir instead of writing ./ into the cwd. @@ -4307,8 +4513,15 @@ def _unsloth_save_compressed_tensors( env["HF_TOKEN"] = token env["HUGGING_FACE_HUB_TOKEN"] = token + # Clean PYTHONPATH = shadow only. torch still comes from the interpreter's site-packages; + # transformers 5.x + llm-compressor main come from the shadow. Dropping the inherited + # PYTHONPATH removes any parent transformers sidecar so the shadow's is authoritative. + if _shadow_pythonpath is not None: + env["PYTHONPATH"] = _shadow_pythonpath + print( f"Unsloth: Quantizing the merged model to {scheme} with llm-compressor " + f"{'(llm-compressor-main shadow) ' if _shadow_pythonpath is not None else ''}" "(in a separate process)..." ) try: @@ -4377,6 +4590,255 @@ def _unsloth_save_compressed_tensors( torch.cuda.empty_cache() +def _unsloth_save_torchao( + model, + save_directory: Union[str, os.PathLike], + tokenizer, + kind: str, + suffix: str, + push_to_hub: bool = False, + token: Optional[Union[str, bool]] = None, + is_main_process: bool = True, + **merge_kwargs, +): + """Export a device-agnostic torchao FP8 / INT8 "portable" checkpoint (no NVIDIA GPU needed). + + Merges LoRA to 16bit in a staging dir, then applies torchao weight-only quantization via + `TorchAoConfig` into `save_directory + "-" + suffix`. No calibration, subprocess, or CUDA. + `kind` is "fp8" (safetensors) or "int8" (.bin; torchao only whitelists float8 for safetensors). + """ + import tempfile + + if isinstance(tokenizer, (PreTrainedTokenizerBase, ProcessorMixin)): + tokenizer = patch_saving_functions(tokenizer) + if token is None: + token = get_token() + + # Only the main process merges, quantizes, and uploads; other ranks return at once. + if not is_main_process: + return None + + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + AutoProcessor, + TorchAoConfig, + ) + from torchao.quantization import Float8WeightOnlyConfig, Int8WeightOnlyConfig + + if kind == "fp8": + quant_type = Float8WeightOnlyConfig() + safe_serialization = True + elif kind == "int8": + quant_type = Int8WeightOnlyConfig() + safe_serialization = False # torchao only supports safetensors for float8 configs + else: + raise RuntimeError(f"Unsloth: unknown torchao export kind '{kind}' (expected fp8/int8).") + + # Always merge into an isolated temp staging dir (never save_directory itself), so a co-selected + # 16-bit export written to save_directory is not overwritten or deleted; the torchao output is + # the sibling "-" (or the repo id on a hub push). + repo_id, work_tmp, model_dev = None, None, None + work_tmp = tempfile.mkdtemp(prefix = "unsloth-torchao-") + if push_to_hub: + repo_id = os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(repo_id.rstrip("/")) or "model") + out_dir = staging + "-" + suffix + else: + base = os.fspath(save_directory).rstrip("/\\") or os.fspath(save_directory) + staging = os.path.join(work_tmp, os.path.basename(base) or "model") + out_dir = base + "-" + suffix + + api = None + try: + if push_to_hub: + from huggingface_hub import HfApi + api = HfApi(token = token) + api.create_repo( + repo_id = repo_id, + repo_type = "model", + private = merge_kwargs.get("private", None), + exist_ok = True, + ) + + # 1) Merge to 16bit at a staging dir (LoRA and base alike). The reload reads default + # weight filenames, so never write variant-named shards here. + merge_kwargs.pop("variant", None) + print(f"Unsloth: Merging to 16bit before torchao {kind} quantization...") + merge_args = dict(merge_kwargs) + merge_args.update( + dict( + model = model, + tokenizer = tokenizer, + save_directory = staging, + save_method = "merged_16bit", + push_to_hub = False, + token = token, + is_main_process = is_main_process, + ) + ) + unsloth_generic_save(**merge_args) + + # 2) Detect VLM + trust_remote_code so the right auto class reloads the staged checkpoint. + # A bare *ForConditionalGeneration also matches text seq2seq (T5/BART/Whisper), so key off + # vision_config / a vision-named architecture only, like the compressed path. + is_vlm = False + trust_remote_code = False + if hasattr(model, "config"): + archs = getattr(model.config, "architectures", None) or [] + is_vlm = hasattr(model.config, "vision_config") or any( + x.endswith("ForVisionText2Text") for x in archs + ) + trust_remote_code = bool(getattr(model.config, "auto_map", None)) + # Custom code can be declared only in the tokenizer/processor config, so also honor an + # auto_map in any staged config (the original load already had the user's consent). + if not trust_remote_code: + for _cfg in ( + "config.json", + "tokenizer_config.json", + "processor_config.json", + "preprocessor_config.json", + ): + try: + _p = os.path.join(staging, _cfg) + if os.path.exists(_p): + with open(_p, "r", encoding = "utf-8") as _f: + if "auto_map" in json.load(_f): + trust_remote_code = True + break + except Exception: + pass + # Reload with the class that matches the checkpoint: an image-text VLM class (with a + # fallback for older Transformers that lack AutoModelForImageTextToText); the model's own + # architecture class for encoder-decoder seq2seq (T5/BART/Whisper are not causal LMs, and + # AutoModelForCausalLM would fail to load them); otherwise causal-LM. + if is_vlm: + try: + from transformers import AutoModelForImageTextToText as _reload_model + except ImportError: + from transformers import AutoModelForVision2Seq as _reload_model + auto_model = _reload_model + elif getattr(getattr(model, "config", None), "is_encoder_decoder", False): + import transformers as _tf + auto_model = next( + ( + getattr(_tf, _arch) + for _arch in (getattr(model.config, "architectures", None) or []) + if getattr(_tf, _arch, None) is not None + ), + AutoModelForCausalLM, + ) + else: + auto_model = AutoModelForCausalLM + auto_processor = AutoProcessor if is_vlm else AutoTokenizer + + # 3) Free the in-memory model's accelerator memory before reloading a fresh copy from disk. + # Covers CUDA and XPU (torchao runs on Intel GPUs too), so the original doesn't sit + # resident alongside the reloaded copy and OOM a device that fit the model once. + _has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() + try: + if ( + (torch.cuda.is_available() or _has_xpu) + and hasattr(model, "parameters") + and not getattr(model, "is_loaded_in_4bit", False) + and not getattr(model, "is_loaded_in_8bit", False) + and not getattr(model, "is_quantized", False) + ): + _devs = {str(p.device) for p in model.parameters()} + if len(_devs) == 1 and next(iter(_devs)).startswith(("cuda", "xpu")): + _dev = next(model.parameters()).device + model.to("cpu") + model_dev = _dev + except Exception: + model_dev = None + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if _has_xpu: + torch.xpu.empty_cache() + + # 4) Reload the staged 16bit checkpoint with torchao applied. bfloat16 is required; + # device_map="auto" falls back to CPU, so this works on any hardware. + print(f"Unsloth: Quantizing the merged model to torchao {kind}...") + dtype_kw = {"torch_dtype": torch.bfloat16} if HAS_TORCH_DTYPE else {"dtype": torch.bfloat16} + quantized_model = auto_model.from_pretrained( + staging, + device_map = "auto", + quantization_config = TorchAoConfig(quant_type = quant_type), + trust_remote_code = trust_remote_code, + **dtype_kw, + ) + staged_tokenizer = auto_processor.from_pretrained( + staging, trust_remote_code = trust_remote_code + ) + + quantized_model.save_pretrained(out_dir, safe_serialization = safe_serialization) + staged_tokenizer.save_pretrained(out_dir) + del quantized_model + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # 5) Validate the artifact. + cfg_path = os.path.join(out_dir, "config.json") + cfg = {} + if os.path.exists(cfg_path): + with open(cfg_path, "r", encoding = "utf-8") as f: + cfg = json.load(f) + if "quantization_config" not in cfg: + raise RuntimeError( + f"Unsloth: torchao {kind} export failed - no quantization_config written to " + f"{cfg_path}" + ) + + # 6) Optional hub upload of the quantized artifact (the temp staging is cleaned in finally). + if push_to_hub: + print(f"Unsloth: Uploading torchao {kind} checkpoint to '{repo_id}' ...") + api.upload_folder( + folder_path = out_dir, + repo_id = repo_id, + repo_type = "model", + commit_message = merge_kwargs.get("commit_message", None), + commit_description = merge_kwargs.get("commit_description", None), + create_pr = merge_kwargs.get("create_pr", False), + revision = merge_kwargs.get("revision", None), + ) + datasets = merge_kwargs.get("datasets", None) + if datasets: + try: + from huggingface_hub import metadata_update + metadata_update(repo_id, {"datasets": datasets}, overwrite = True, token = token) + except Exception as meta_err: + logger.warning_once( + f"Unsloth: could not update datasets metadata for {repo_id}: {meta_err}" + ) + + result = repo_id if push_to_hub else out_dir + print( + f"Unsloth: Saved torchao {kind} checkpoint to '{result}'.\n" + f"Unsloth: This is portable (produced on any device, no NVIDIA GPU required). Load it " + f"with vLLM or transformers; FP8/INT8 acceleration is available on supported GPUs." + ) + return result + finally: + if model_dev is not None: + try: + model.to(model_dev) + except Exception: + logger.warning_once( + "Unsloth: could not restore the model to its original device after torchao " + "export; it may remain on CPU." + ) + if work_tmp is not None: + shutil.rmtree(work_tmp, ignore_errors = True) + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def unsloth_save_pretrained_torchao( self, save_directory: Union[str, os.PathLike], diff --git a/unsloth_cli/__init__.py b/unsloth_cli/__init__.py index 440b6276cd..b3831f5314 100644 --- a/unsloth_cli/__init__.py +++ b/unsloth_cli/__init__.py @@ -11,7 +11,7 @@ from importlib.metadata import version as package_version, PackageNotFoundError from unsloth_cli.commands.train import train from unsloth_cli.commands.inference import inference from unsloth_cli.commands.chat import chat -from unsloth_cli.commands.connect import connect_app +from unsloth_cli.commands.start import start_app from unsloth_cli.commands.export import export, list_checkpoints from unsloth_cli.commands.studio import ( run as studio_run, @@ -79,9 +79,16 @@ app.command()(export) app.command("list-checkpoints")(list_checkpoints) app.add_typer(studio_app, name = "studio", help = "Unsloth Studio commands.") app.add_typer( - connect_app, + start_app, + name = "start", + help = "Start a coding agent (Claude, Codex, OpenClaw, OpenCode, Hermes, Pi) against Studio.", +) +# Backwards-compatible hidden alias: `unsloth connect` routes to `unsloth start`. +app.add_typer( + start_app, name = "connect", - help = "Connect a coding agent (Claude Code, Codex) to Studio.", + hidden = True, + help = "Deprecated alias for `unsloth start`.", ) # Top-level `unsloth run` aliases `unsloth studio run`; same context diff --git a/unsloth_cli/commands/connect.py b/unsloth_cli/commands/connect.py deleted file mode 100644 index 096a7925e5..0000000000 --- a/unsloth_cli/commands/connect.py +++ /dev/null @@ -1,777 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""`unsloth connect` — launch a coding agent against a running Studio server.""" - -import json -import os -import re -import shlex -import shutil -import signal -import subprocess -import urllib.error -import urllib.request -from pathlib import Path -from typing import NoReturn, Optional - -import typer - -from unsloth_cli._inference import ( - _USER_AGENT, - _studio_token, - ensure_studio_backend_path, - find_studio_server, - is_loopback_url, - urlopen_no_redirect, - verify_studio_identity, -) - -connect_app = typer.Typer( - help = "Connect a coding agent to a running Studio server.", - no_args_is_help = True, - context_settings = {"help_option_names": ["-h", "--help"]}, -) - -_CODEX_PROFILE = "unsloth_api" -_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" -_HERMES_ENV_KEY = "UNSLOTH_API_KEY" -_HERMES_PROVIDER = "unsloth" -_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" -_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} -_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") - -# Shared by every agent command; only the config/env/command differ. -_MODEL_OPTION = typer.Option( - None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." -) -_KEY_OPTION = typer.Option( - None, - "--api-key", - envvar = "UNSLOTH_API_KEY", - help = ( - "Studio API key. For a local Studio it is minted automatically and " - "remembered per server. For a remote server, pass one with --api-key " - "(or UNSLOTH_API_KEY); it is remembered for next time." - ), -) -_LAUNCH_OPTION = typer.Option( - True, - "--launch/--no-launch", - help = "--no-launch prints the env and command instead (remote shells, WSL).", -) - - -def _fail(message: str) -> NoReturn: - typer.echo(message, err = True) - raise typer.Exit(code = 1) - - -def _http_error_detail(exc: urllib.error.HTTPError) -> str: - try: - body = json.loads(exc.read().decode()) - return body.get("detail") or body["error"]["message"] - except Exception: - return str(exc) - - -def _http_json( - method: str, - url: str, - token: str, - payload = None, - timeout = 30, - error = None, -): - """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" - request = urllib.request.Request( - url, - data = None if payload is None else json.dumps(payload).encode(), - headers = { - "Authorization": f"Bearer {token}", - "Content-Type": "application/json", - "User-Agent": _USER_AGENT, - }, - method = method, - ) - try: - # No redirects: a 3xx would leak this bearer token to an unvetted base. - with urlopen_no_redirect(request, timeout = timeout) as response: - return json.loads(response.read().decode() or "{}") - except urllib.error.HTTPError as exc: - if error is None: - raise - _fail(f"{error}: {_http_error_detail(exc)}") - except (urllib.error.URLError, TimeoutError) as exc: - if error is None: - raise - _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") - - -def _require_studio() -> str: - base = find_studio_server() - if base is None: - expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888") - _fail( - f"No running Studio server found at {expected}. Start one with " - "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." - ) - return base - - -def _key_cache_path() -> Path: - ensure_studio_backend_path() - from utils.paths import auth_root - return auth_root() / "agent_api_key.json" - - -def _read_cache(cache: Path) -> dict: - try: - data = json.loads(cache.read_text(encoding = "utf-8")) - except Exception: - return {} - return data if isinstance(data, dict) else {} - - -def _server_buckets(servers: dict, base: str) -> dict: - # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a - # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). - entry = servers.get(base) if isinstance(servers, dict) else None - if isinstance(entry, list): - return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} - if not isinstance(entry, dict): - return {"saved": [], "minted": []} - - def _strs(name: str) -> list: - value = entry.get(name) - return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] - - return {"saved": _strs("saved"), "minted": _strs("minted")} - - -def _cached_keys(cache: Path, base: str, source: str) -> list: - # Keys are scoped per server. `source` splits user-supplied --api-key keys - # ("saved", trusted for that base) from auto-minted ones ("minted", replayed - # only after the identity check). Legacy unscoped caches are ignored. - return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] - - -def _write_private_json(path: Path, data: dict) -> None: - # O_CREAT with 0o600 so a file holding an API key is never world-readable, - # even briefly (existing files keep whatever perms the user set). - path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) - fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) - with os.fdopen(fd, "w") as handle: - handle.write(json.dumps(data, indent = 2) + "\n") - - -def _read_json_object(path: Path) -> Optional[dict]: - # {} when missing, None when it can't be parsed as an object (so the caller - # leaves a user-managed file untouched rather than clobbering it). - if not path.exists(): - return {} - try: - data = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - return None - return data if isinstance(data, dict) else None - - -def _subdict(parent: dict, key: str) -> dict: - child = parent.get(key) - if not isinstance(child, dict): - child = parent[key] = {} - return child - - -def _remember_key(cache: Path, base: str, key: str, source: str) -> None: - data = _read_cache(cache) - servers = data.get("servers") - if not isinstance(servers, dict): - servers = data["servers"] = {} - buckets = _server_buckets(servers, base) - other = "minted" if source == "saved" else "saved" - buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] - buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance - new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} - if servers.get(base) == new_entry: - return - servers[base] = new_entry - # Collapse legacy unscoped fields. - data.pop("keys", None) - data.pop("key", None) - try: - _write_private_json(cache, data) - except OSError: - pass # worst case the next launch mints another key - - -def _key_accepted(base: str, key: str) -> bool: - try: - _http_json("GET", f"{base}/v1/models", key) - return True - except Exception: - return False - - -def _agent_api_key(base: str, explicit: Optional[str]) -> str: - cache = _key_cache_path() - if explicit: - _remember_key(cache, base, explicit, "saved") - return explicit - - # Replay a key the user saved for *this exact* server first (scoped per base, - # so it only goes back there -- including a remote/SSH-tunnelled Studio whose - # secret the local handshake can't match). Skip ones the server rejects. - for key in _cached_keys(cache, base, "saved"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "saved") - return key - - # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() - # trusts a base after only a health check, so both are limited to a loopback - # server we can cryptographically confirm is ours. - if not is_loopback_url(base): - _fail( - f"No saved API key for {base} and automatic minting only runs against " - "a local Studio. Create an API key in Studio → Settings → API and " - "pass it with --api-key (it is remembered per server), or set " - "UNSLOTH_API_KEY." - ) - if not verify_studio_identity(base): - _fail( - f"Couldn't verify that {base} is your Studio (it may be running as a " - "different OS user, or another process took the port). Create an API " - "key in Studio → Settings → API and pass it with --api-key, or set " - "UNSLOTH_API_KEY." - ) - - # Identity verified: replay a previously auto-minted key, else mint a new one. - for key in _cached_keys(cache, base, "minted"): - if _key_accepted(base, key): - _remember_key(cache, base, key, "minted") - return key - - # Self-issue a JWT (signed with the local secret) and mint a key. - token = _studio_token() - if token is None: - _fail( - "Couldn't authenticate with the Studio server automatically. Create " - "an API key in Studio → Settings → API and pass it with --api-key, " - "or set UNSLOTH_API_KEY." - ) - key = _http_json( - "POST", - f"{base}/api/auth/api-keys", - token, - {"name": "Coding agents (unsloth connect)"}, - error = "Couldn't create an API key", - )["key"] - _remember_key(cache, base, key, "minted") - return key - - -def _loaded_models(base: str, key: str) -> list: - return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) - - -def _resolve_model(base: str, key: str, requested: Optional[str]) -> dict: - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] == requested), None) - if requested and match is None: - typer.echo(f"Loading {requested} on the Studio server (this can take a while)…") - loaded = _http_json( - "POST", - f"{base}/api/inference/load", - key, - {"model_path": requested}, - timeout = 3600, - error = "Model load failed", - ) - # Studio registers the model under a canonical id (resolved identifier, - # casing) that /v1/models echoes but which may differ from the path we - # passed; match on the id the load reports so we don't silently fall - # through to models[0] and connect to a different loaded model. - wanted = {requested} - if isinstance(loaded, dict): - wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} - models = _loaded_models(base, key) - match = next((m for m in models if m["id"] in wanted), None) - if match is not None: - return match - if requested: - # We asked Studio to load it and it didn't surface in /v1/models; don't - # silently hand back an unrelated loaded model. - _fail( - f"Studio didn't report '{requested}' as loaded. Double-check the model " - "id, or load it from the model dropdown in the UI." - ) - if not models: - _fail( - "No model is loaded in Studio. Load one from the model dropdown in " - "the UI, or pass --model to load it from here." - ) - return models[0] - - -def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: - # Codex always streams, and Studio only streams /v1/responses from llama-server. - try: - status = _http_json("GET", f"{base}/api/inference/status", key) - except urllib.error.HTTPError as exc: - if exc.code == 404: - return # older server without the endpoint; don't block the launch - raise - if status.get("is_gguf"): - return - hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" - _fail( - f"Codex needs a GGUF model served by llama-server, but {model_id} is on " - f"the transformers backend. Try: unsloth connect codex --model {hint}" - ) - - -def claude_settings_path() -> Path: - return Path.home() / ".claude" / "settings.json" - - -def ensure_claude_attribution_header() -> None: - # The header invalidates the llama.cpp KV cache (~90% slower) and Claude - # Code only honors this setting from settings.json, not the env var. - path = claude_settings_path() - settings = {} - if path.exists(): - try: - settings = json.loads(path.read_text(encoding = "utf-8")) - except (ValueError, OSError): - settings = None - if not isinstance(settings, dict): - typer.echo( - f"Warning: couldn't parse {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - env = settings.get("env") - if not isinstance(env, dict): - env = settings["env"] = {} - if str(env.get("CLAUDE_CODE_ATTRIBUTION_HEADER")) == "0": - return - env["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "0" - try: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(json.dumps(settings, indent = 2) + "\n", encoding = "utf-8") - except OSError: - typer.echo( - f"Warning: couldn't write {path} — set CLAUDE_CODE_ATTRIBUTION_HEADER " - 'to "0" in its "env" section yourself, or local inference will be much slower.', - err = True, - ) - return - typer.echo(f"Disabled Claude Code's attribution header in {path} (it breaks KV-cache reuse).") - - -_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" - - -def _claude_cache_flags() -> list: - # The flag moves per-machine context (cwd, env info, git status) out of - # the system prompt, where it changes every session and defeats llama.cpp - # prefix caching. As of 2.1.175 it only takes effect in print mode (`-p` - # passed through ctx.args); interactive sessions accept and ignore it. - # Claude Code < 2.1.98 aborts on the unknown flag, so check the version - # first; no local binary means a --no-launch printout for another machine. - executable = shutil.which("claude") - if executable is None: - return [_DYNAMIC_SECTIONS_FLAG] - try: - result = subprocess.run( - [executable, "--version"], capture_output = True, text = True, timeout = 10 - ) - version = tuple(int(part) for part in result.stdout.split()[0].split(".")) - except Exception: - return [] - return [_DYNAMIC_SECTIONS_FLAG] if version >= (2, 1, 98) else [] - - -def codex_home() -> Path: - return Path(os.environ.get("CODEX_HOME") or Path.home() / ".codex") - - -def _merge_codex_config(existing: str, base: str) -> str: - chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table - if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): - if chunks[0] and not chunks[0].endswith("\n"): - chunks[0] += "\n" - chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' - # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. - stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") - text = "".join(c for c in chunks if not c.startswith(stale)) - if not text.endswith("\n"): - text += "\n" - if not text.endswith("\n\n"): - text += "\n" - return text + ( - f"{_PROVIDER_HEADER}\n" - 'name = "Unsloth Studio"\n' - f"base_url = {json.dumps(base + '/v1')}\n" - f'env_key = "{_CODEX_ENV_KEY}"\n' - 'wire_api = "responses"\n' - "requires_openai_auth = false\n" - ) - - -def write_codex_config(base: str, model: dict) -> None: - home = codex_home() - home.mkdir(parents = True, exist_ok = True) - - config = home / "config.toml" - existing = config.read_text(encoding = "utf-8") if config.exists() else "" - merged = _merge_codex_config(existing, base) - if merged != existing: - config.write_text(merged, encoding = "utf-8") - typer.echo(f"Updated {config}") - - # oss_provider here too: codex --oss picks the provider from it, and the - # profile layer must beat a user-set value (e.g. "ollama") in config.toml. - profile_text = ( - f'oss_provider = "{_CODEX_PROFILE}"\n' - f'model_provider = "{_CODEX_PROFILE}"\n' - f"model = {json.dumps(model['id'])}\n" - ) - window = model.get("context_length") or model.get("max_context_length") - if window: - profile_text += f"model_context_window = {int(window)}\n" - profile = home / f"{_CODEX_PROFILE}.config.toml" - if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: - profile.write_text(profile_text, encoding = "utf-8") - typer.echo(f"Updated {profile}") - - -def _wsl_windows_executable(command: list) -> Optional[str]: - if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): - return None - executable = shutil.which(command[0]) - if executable and executable.startswith("/mnt/"): - return executable - return None - - -def _merge_wslenv(current: str, names: tuple) -> str: - entries = [entry for entry in current.split(":") if entry] - existing = {entry.split("/", 1)[0] for entry in entries} - entries.extend(name for name in names if name not in existing) - return ":".join(entries) - - -def _print_env( - env: dict, - command: list, - unset_env: tuple = (), - wsl_env_bridge: tuple = (), -) -> None: - if os.name == "nt": - for name in unset_env: - typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") - for name, value in env.items(): - # PowerShell: ` is the escape char, and $ triggers expansion inside "". - escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") - typer.echo(f'$env:{name} = "{escaped}"') - typer.echo(subprocess.list2cmdline(command)) - return - for name in unset_env: - typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") - for name, value in env.items(): - typer.echo(f"export {name}={shlex.quote(value)}") - if wsl_env_bridge: - typer.echo( - f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" - ) - typer.echo(shlex.join(command)) - - -def _launch( - command: list, - env: dict, - install_hint: str, - unset_env: tuple = (), -) -> NoReturn: - executable = shutil.which(command[0]) - if executable is None: - _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - child_env = dict(os.environ) - if wsl_env_bridge: - child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) - for name in unset_env: - child_env[name] = "" - else: - for name in unset_env: - child_env.pop(name, None) - child_env.update(env) - # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. - previous = signal.signal(signal.SIGINT, signal.SIG_IGN) - try: - code = subprocess.run([executable, *command[1:]], env = child_env).returncode - finally: - signal.signal(signal.SIGINT, previous) - # Negative returncode means killed by signal N; shells expect 128+N. - raise typer.Exit(code = code if code >= 0 else 128 - code) - - -def _connect(api_key: Optional[str], model: Optional[str]) -> tuple: - base = _require_studio() - key = _agent_api_key(base, api_key) - return base, key, _resolve_model(base, key, model) - - -def _run( - base: str, - entry: dict, - env: dict, - command: list, - *, - launch: bool, - install_hint: str, - unset_env: tuple = (), -) -> None: - typer.echo(f"Studio {base} · model {entry['id']}") - wsl_env_bridge = ( - tuple(dict.fromkeys((*env.keys(), *unset_env))) if _wsl_windows_executable(command) else () - ) - if not launch: - _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) - return - _launch(command, env, install_hint = install_hint, unset_env = unset_env) - - -def openclaw_config_path() -> Path: - return Path.home() / ".openclaw" / "openclaw.json" - - -def write_openclaw_config(base: str, key: str, model: dict) -> None: - path = openclaw_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). - provider_model = {"id": model["id"], "name": model["id"]} - window = model.get("context_length") or model.get("max_context_length") - if window: - provider_model["contextWindow"] = int(window) - models = _subdict(config, "models") - models.setdefault("mode", "merge") - _subdict(models, "providers")["unsloth"] = { - "baseUrl": f"{base}/v1", - "apiKey": key, - "api": "openai-completions", - "models": [provider_model], - } - # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). - defaults = _subdict(_subdict(config, "agents"), "defaults") - _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" - # Unauthenticated loopback gateway: without auth.mode=none the client won't open - # the websocket. The daemon must still be started separately (`openclaw gateway`). - gateway = _subdict(config, "gateway") - gateway.setdefault("mode", "local") - _subdict(gateway, "auth").setdefault("mode", "none") - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def opencode_config_path() -> Path: - config_home = os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config" - return Path(config_home) / "opencode" / "opencode.json" - - -def write_opencode_config(base: str, key: str, model: dict) -> None: - path = opencode_config_path() - config = _read_json_object(path) - if config is None: - typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", - err = True, - ) - return - before = json.dumps(config, sort_keys = True) - config.setdefault("$schema", "https://opencode.ai/config.json") - _subdict(config, "provider")["unsloth"] = { - "npm": "@ai-sdk/openai-compatible", - "name": "Unsloth Studio", - "options": {"baseURL": f"{base}/v1", "apiKey": key}, - "models": {model["id"]: {"name": model["id"]}}, - } - # OpenCode selects a model by "/". - config["model"] = f"unsloth/{model['id']}" - if json.dumps(config, sort_keys = True) != before: - _write_private_json(path, config) - typer.echo(f"Updated {path}") - - -def hermes_config_path() -> Path: - return Path.home() / ".hermes" / "config.yaml" - - -def write_hermes_config(base: str, model: dict) -> None: - import yaml - - path = hermes_config_path() - config: dict = {} - if path.exists(): - try: - loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) - except (yaml.YAMLError, OSError): - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - if isinstance(loaded, dict): - config = loaded - elif loaded is not None: - # Non-empty, non-mapping YAML is a user-managed file; leave it. - typer.echo( - f"Warning: couldn't parse {path} — configure the custom endpoint " - "there yourself, or move the file aside and re-run.", - err = True, - ) - return - # Hermes only reads the key for a *named* custom provider (a bare - # `provider: custom` ignores it), so register it under providers.*. - _subdict(config, "model").update( - provider = f"custom:{_HERMES_PROVIDER}", - default = model["id"], - api_mode = "openai", - ) - _subdict(config, "providers")[_HERMES_PROVIDER] = { - "base_url": f"{base}/v1", - "api_mode": "openai", - "key_env": _HERMES_ENV_KEY, - } - text = yaml.safe_dump(config, sort_keys = False) - if not path.exists() or path.read_text(encoding = "utf-8") != text: - path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(text, encoding = "utf-8") - typer.echo(f"Updated {path}") - - -@connect_app.command("claude", context_settings = _PASSTHROUGH) -def claude( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Claude Code at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - model_id = entry["id"] - ensure_claude_attribution_header() - - env = { - "ANTHROPIC_BASE_URL": base, - "ANTHROPIC_AUTH_TOKEN": key, - "ANTHROPIC_MODEL": model_id, - # Update checks, beta features, and other background requests either - # stall against a local server or evict the conversation from - # llama-server's KV-cache slots, so turn off everything nonessential. - "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", - "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", - } - command = ["claude", "--model", model_id, *_claude_cache_flags(), *ctx.args] - install_hint = ( - "irm https://claude.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://claude.ai/install.sh | bash" - ) - _run( - base, - entry, - env, - command, - launch = launch, - install_hint = install_hint, - unset_env = _CLAUDE_ENV_UNSET, - ) - - -@connect_app.command("codex", context_settings = _PASSTHROUGH) -def codex( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenAI Codex at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - _require_gguf_for_codex(base, key, entry["id"]) - write_codex_config(base, entry) - - env = {_CODEX_ENV_KEY: key} - command = ["codex", "--oss", "--profile", _CODEX_PROFILE, *ctx.args] - _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") - - -@connect_app.command("openclaw", context_settings = _PASSTHROUGH) -def openclaw( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenClaw at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_openclaw_config(base, key, entry) # key lives in the config, not the env - - command = ["openclaw", *ctx.args] - install_hint = ( - "iwr -useb https://openclaw.ai/install.ps1 | iex" - if os.name == "nt" - else "curl -fsSL https://openclaw.ai/install.sh | bash" - ) - _run(base, entry, {}, command, launch = launch, install_hint = install_hint) - - -@connect_app.command("opencode", context_settings = _PASSTHROUGH) -def opencode( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point OpenCode at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_opencode_config(base, key, entry) # key lives in the config, not the env - - command = ["opencode", *ctx.args] - _run(base, entry, {}, command, launch = launch, install_hint = "npm install -g opencode-ai") - - -@connect_app.command("hermes", context_settings = _PASSTHROUGH) -def hermes( - ctx: typer.Context, - model: Optional[str] = _MODEL_OPTION, - api_key: Optional[str] = _KEY_OPTION, - launch: bool = _LAUNCH_OPTION, -): - """Point Hermes (Nous Research) at the running Studio server and start it.""" - base, key, entry = _connect(api_key, model) - write_hermes_config(base, entry) - - env = {_HERMES_ENV_KEY: key} - command = ["hermes", *ctx.args] - install_hint = ( - "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" - "/main/scripts/install.sh | bash" - ) - _run(base, entry, env, command, launch = launch, install_hint = install_hint) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py new file mode 100644 index 0000000000..b188180188 --- /dev/null +++ b/unsloth_cli/commands/start.py @@ -0,0 +1,1497 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""`unsloth start` — launch a coding agent against a running Studio server.""" + +import atexit +import contextlib +import json +import os +import re +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import NamedTuple, NoReturn, Optional +from urllib.parse import urlparse + +import click +import typer + +from unsloth_cli._inference import ( + _USER_AGENT, + _studio_token, + ensure_studio_backend_path, + find_studio_server, + is_loopback_url, + urlopen_no_redirect, + verify_studio_identity, +) + +start_app = typer.Typer( + help = "Start a coding agent against a running Studio server.", + no_args_is_help = True, + context_settings = {"help_option_names": ["-h", "--help"]}, +) + +_CODEX_PROFILE = "unsloth_api" +_CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" +_HERMES_ENV_KEY = "UNSLOTH_API_KEY" +_HERMES_PROVIDER = "unsloth" +# Hermes refuses to initialize when the model window is under 64,000 tokens; its +# error message points at the model.context_length / auxiliary.compression +# overrides in config.yaml. write_hermes_config claims this value for smaller +# windows and scales the compaction threshold back down to the real window. +_HERMES_MIN_CONTEXT = 65536 +_PI_PROVIDER = "unsloth" +_PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" +_PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} +_CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") + +# Shared by every agent command; only the config/env/command differ. +_MODEL_OPTION = typer.Option( + None, "--model", "-m", help = "Model for the agent; defaults to the one loaded in Studio." +) +_KEY_OPTION = typer.Option( + None, + "--api-key", + envvar = "UNSLOTH_API_KEY", + help = ( + "Studio API key. For a local Studio it is minted automatically and " + "remembered per server. For a remote server, pass one with --api-key " + "(or UNSLOTH_API_KEY); it is remembered for next time." + ), +) +_LAUNCH_OPTION = typer.Option( + True, + "--launch/--no-launch", + help = "--no-launch prints the env and command instead (remote shells, WSL).", +) +_SERVE_OPTION = typer.Option( + True, + "--serve/--no-serve", + help = ( + "If no Studio server is running, auto-start one for --model and stop it when the " + "agent exits. --no-serve keeps the old behavior of erroring out." + ), +) +# Model-load knobs mirrored from `unsloth run`; only used when --model triggers a +# load on the server. Server-startup flags (--host/--port/--cloudflare/...) do not +# apply here because `unsloth start` attaches to an already-running server. +_GGUF_VARIANT_OPTION = typer.Option( + None, "--gguf-variant", help = "GGUF quant variant to load (e.g. UD-Q4_K_XL)." +) +_CONTEXT_OPTION = typer.Option( + 0, + "--max-seq-length", + "--context-length", + help = "Context length in tokens for the load (0 = model default).", +) +_LOAD_4BIT_OPTION = typer.Option( + True, "--load-in-4bit/--no-load-in-4bit", help = "Load hub models in 4-bit (ignored for GGUF)." +) +_TENSOR_PARALLEL_OPTION = typer.Option( + False, + "--tensor-parallel/--no-tensor-parallel", + help = "Split a GGUF across GPUs by tensor instead of by layer (multi-GPU only).", +) +# One normalized "run tools without prompting" switch. Each agent spells this +# differently and it's easy to forget which is which, so accept every spelling and +# route to the agent's own mechanism in _yolo_command_flags / the config writers. +_YOLO_OPTION = typer.Option( + False, + "--yolo", + "--dangerously-skip-permissions", + "--dangerously-bypass-approvals-and-sandbox", + help = ( + "Auto-approve all tool actions for this session; routed to the agent's own " + "flag/config. Any of the three spellings works for any agent." + ), +) + +# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no +# such flag (config only) and are handled in their config writers, so they are absent. +_YOLO_COMMAND_FLAGS = { + "claude": ["--dangerously-skip-permissions"], + "codex": ["--dangerously-bypass-approvals-and-sandbox"], + "hermes": ["--yolo"], + # Pi never prompts per tool call; its only approval gate is project trust, so -a + # (trust project resources) is the closest "don't ask me" equivalent. + "pi": ["--approve"], +} + + +def _yolo_command_flags(agent: str, yolo: bool) -> list: + # .get so a config-based agent (or a typo) yields no flag instead of a KeyError. + return _YOLO_COMMAND_FLAGS.get(agent, []) if yolo else [] + + +class LoadOptions(NamedTuple): + """Model-load knobs forwarded to /api/inference/load when --model triggers a load.""" + + gguf_variant: Optional[str] = None + max_seq_length: int = 0 + load_in_4bit: bool = True + tensor_parallel: bool = False + + +def _split_repo_variant(model: str) -> tuple: + """Split ``org/name:QUANT`` into ``(repo, variant)`` -> ``("org/name", "QUANT")``. + + ``unsloth run`` and llama.cpp accept ``--model org/name:QUANT`` as shorthand for + ``--model org/name --gguf-variant QUANT``. Mirror that here so a ``:variant`` suffix + resolves against the already-loaded ``org/name`` (which /v1/models lists without the + suffix) instead of trying to load a repo id containing ``:`` -- which Hugging Face + rejects, and which would evict a model another session is using. Local paths, Windows + drive letters, and ids without a ``:`` pass through unchanged. + """ + s = (model or "").strip() + if not s or s.startswith(("/", "./", "../", "~")) or s == ".": + return s, None + if len(s) >= 2 and s[1] == ":" and s[0].isalpha(): # Windows drive, e.g. C:\models\x + return s, None + if ":" not in s: + return s, None + repo, _, variant = s.rpartition(":") + if not repo or not variant or "/" in variant: + return s, None + return repo, variant + + +def _fail(message: str) -> NoReturn: + typer.echo(message, err = True) + raise typer.Exit(code = 1) + + +def _http_error_detail(exc: urllib.error.HTTPError) -> str: + try: + body = json.loads(exc.read().decode()) + return body.get("detail") or body["error"]["message"] + except Exception: + return str(exc) + + +def _http_json( + method: str, + url: str, + token: str, + payload = None, + timeout = 30, + error = None, +): + """On HTTPError: raise if `error` is None, else fail with `error` plus the server's detail.""" + request = urllib.request.Request( + url, + data = None if payload is None else json.dumps(payload).encode(), + headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "User-Agent": _USER_AGENT, + }, + method = method, + ) + try: + # No redirects: a 3xx would leak this bearer token to an unvetted base. + with urlopen_no_redirect(request, timeout = timeout) as response: + return json.loads(response.read().decode() or "{}") + except urllib.error.HTTPError as exc: + if error is None: + raise + _fail(f"{error}: {_http_error_detail(exc)}") + except (urllib.error.URLError, TimeoutError) as exc: + if error is None: + raise + _fail(f"{error}: {getattr(exc, 'reason', None) or exc}") + + +# A server that WE auto-started (never one we merely found). Kept at module scope so +# _run's finally and the atexit backstop can tear it down without threading a handle +# through all six agent commands. Only one agent runs per process, so one slot is enough. +_auto_served_server: Optional[subprocess.Popen] = None +# Model download + load can be slow; give the auto-started server room before giving up. +_SERVER_START_TIMEOUT_S = 900 + + +def _studio_healthy(base: str, timeout: float = 3.0) -> bool: + request = urllib.request.Request(f"{base}/api/health", headers = {"User-Agent": _USER_AGENT}) + try: + with urllib.request.urlopen(request, timeout = timeout) as response: + return json.loads(response.read(65536).decode() or "{}").get("status") == "healthy" + except Exception: + return False + + +def _log_tail(path: Path, lines: int = 20) -> str: + try: + return "\n".join(path.read_text(encoding = "utf-8", errors = "replace").splitlines()[-lines:]) + except OSError: + return "(no server log)" + + +def _shutdown_server(server: Optional[subprocess.Popen]) -> None: + # Idempotent teardown of a server WE started, plus its own children (llama-server, + # cloudflared). A no-op once the process is already gone. + if server is None or server.poll() is not None: + return + if os.name == "nt": + # terminate()/kill() reach only the parent `unsloth run`; taskkill /T walks the + # whole tree so the llama-server child doesn't keep the port and GPU (matches the + # taskkill /T /F pattern already used in unsloth/dataprep/synthetic.py). + try: + subprocess.run( + ["taskkill", "/PID", str(server.pid), "/T", "/F"], + capture_output = True, + timeout = 15, + check = False, + ) + server.wait(timeout = 5) + except Exception: + with contextlib.suppress(Exception): + server.kill() + return + try: + os.killpg(os.getpgid(server.pid), signal.SIGTERM) + except OSError: + server.terminate() + try: + server.wait(timeout = 15) + except Exception: + try: + os.killpg(os.getpgid(server.pid), signal.SIGKILL) + except OSError: + server.kill() + + +def _shutdown_auto_served() -> None: + global _auto_served_server + server, _auto_served_server = _auto_served_server, None + if server is not None and server.poll() is None: + typer.echo("Stopping the auto-started Studio server…") + _shutdown_server(server) + + +def _start_studio_server(base: str, model: str, load: LoadOptions) -> subprocess.Popen: + """Spawn `unsloth run` for `model`, wait until it is fully ready, and return it.""" + global _auto_served_server + unsloth = shutil.which("unsloth") or "unsloth" + parsed = urlparse(base) + # --disable-tools = passthrough mode (relay the agent's own tools); --no-cloudflare = + # loopback only, no tunnel. Mirrors .github/scripts/serve-unsloth-run.sh. + command = [ + unsloth, + "run", + "-H", + parsed.hostname or "127.0.0.1", + "-p", + str(parsed.port or 8888), + "--disable-tools", + "--no-cloudflare", + "--model", + model, + ] + if load.gguf_variant: + command += ["--gguf-variant", load.gguf_variant] + if load.max_seq_length: + command += ["--context-length", str(load.max_seq_length)] + if not load.load_in_4bit: + command += ["--no-load-in-4bit"] + if load.tensor_parallel: + command += ["--tensor-parallel"] + + log_path = Path(tempfile.gettempdir()) / f"unsloth-start-server-{os.getpid()}.log" + typer.echo( + f"No Studio server at {base}. Starting one for {model} (loading the model can take a while)…" + ) + typer.echo(f"Server log: {log_path}") + # 0600: the `unsloth run` banner in this log carries the minted sk-unsloth- key, and + # the tempdir is world-traversable. Unlink first so a stale looser-mode file (pid + # reuse) can't survive with its old permissions. + log_path.unlink(missing_ok = True) + log = os.fdopen(os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb") + # Own session/process group so a mid-session Ctrl+C (cancel a turn) doesn't reach the + # server; we tear it down explicitly when the agent exits. + kwargs: dict = {"stdout": log, "stderr": subprocess.STDOUT, "stdin": subprocess.DEVNULL} + if os.name == "nt": + kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + kwargs["start_new_session"] = True + try: + server = subprocess.Popen(command, **kwargs) + finally: + log.close() # Popen dup'd the fd; drop the parent's copy + _auto_served_server = server + atexit.register(_shutdown_auto_served) + + deadline = time.monotonic() + _SERVER_START_TIMEOUT_S + while time.monotonic() < deadline: + if server.poll() is not None: + tail = _log_tail(log_path) + _shutdown_auto_served() + _fail(f"The Studio server stopped before it was ready. Last log lines:\n{tail}") + # `unsloth run` prints the minted key only after the server is up AND the model is + # loaded, so it is the fully-ready signal (same contract serve-unsloth-run.sh uses). + if _studio_healthy(base) and "sk-unsloth-" in _log_tail(log_path, lines = 400): + typer.echo(f"Studio server ready at {base}.") + return server + time.sleep(2.0) + _shutdown_auto_served() + _fail( + f"The Studio server didn't become ready within {_SERVER_START_TIMEOUT_S}s. See {log_path}." + ) + + +def _effective_base(base: str) -> str: + # `unsloth run` binds to `parsed.port or 8888` and serves at the root, so normalize + # UNSLOTH_STUDIO_URL to plain scheme://host:port. A portless http://127.0.0.1 would + # otherwise launch on 8888 but poll port 80, and a path like /studio would poll + # /studio/api/health (404) -- either way hitting the startup timeout. IPv6 literals + # stay bracketed. + parsed = urlparse(base) + host = parsed.hostname or "127.0.0.1" + if ":" in host: # bare IPv6 literal (urlparse strips the brackets) + host = f"[{host}]" + return f"{parsed.scheme or 'http'}://{host}:{parsed.port or 8888}" + + +def _require_studio( + model: Optional[str] = None, + load: Optional[LoadOptions] = None, + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + """Return (base, server). server is a Popen only when WE auto-started it.""" + base = find_studio_server() + if base is not None: + return base, None + expected = os.environ.get("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888").rstrip("/") + # Auto-start a local server only for an interactive launch with a model to serve, and + # only for a plain-HTTP loopback target: never stand in for an explicit remote + # UNSLOTH_STUDIO_URL, and never for an https:// one -- `unsloth run` serves plain + # HTTP, so the health poll against https would spin until the startup timeout. + if ( + serve + and launch + and model + and is_loopback_url(expected) + and urlparse(expected).scheme == "http" + ): + # Normalize to the port unsloth run actually binds, so the health poll and the + # returned base hit the same server we launch (not a portless :80). + expected = _effective_base(expected) + return expected, _start_studio_server(expected, model, load or LoadOptions()) + model_hint = "" if model else " Pass --model to have it start one for you, or" + _fail( + f"No running Studio server found at {expected}.{model_hint} start one with " + "`unsloth studio`, or point UNSLOTH_STUDIO_URL at a remote server." + ) + + +def _key_cache_path() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agent_api_key.json" + + +def _read_cache(cache: Path) -> dict: + try: + data = json.loads(cache.read_text(encoding = "utf-8")) + except Exception: + return {} + return data if isinstance(data, dict) else {} + + +def _server_buckets(servers: dict, base: str) -> dict: + # Normalise a server's entry to {"saved": [...], "minted": [...]}, tolerating a + # corrupt/legacy value (bare string/list -> treated as minted, behind the handshake). + entry = servers.get(base) if isinstance(servers, dict) else None + if isinstance(entry, list): + return {"saved": [], "minted": [k for k in entry if isinstance(k, str)]} + if not isinstance(entry, dict): + return {"saved": [], "minted": []} + + def _strs(name: str) -> list: + value = entry.get(name) + return [k for k in value if isinstance(k, str)] if isinstance(value, list) else [] + + return {"saved": _strs("saved"), "minted": _strs("minted")} + + +def _cached_keys(cache: Path, base: str, source: str) -> list: + # Keys are scoped per server. `source` splits user-supplied --api-key keys + # ("saved", trusted for that base) from auto-minted ones ("minted", replayed + # only after the identity check). Legacy unscoped caches are ignored. + return _server_buckets(_read_cache(cache).get("servers", {}), base)[source] + + +def _write_private_json(path: Path, data: dict) -> None: + # O_CREAT with 0o600 so a file holding an API key is never world-readable, + # even briefly (existing files keep whatever perms the user set). + path.parent.mkdir(parents = True, exist_ok = True, mode = 0o700) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as handle: + handle.write(json.dumps(data, indent = 2) + "\n") + + +def _read_json_object(path: Path) -> Optional[dict]: + # {} when missing, None when it can't be parsed as an object (so the caller + # leaves a user-managed file untouched rather than clobbering it). + if not path.exists(): + return {} + try: + data = json.loads(path.read_text(encoding = "utf-8")) + except (ValueError, OSError): + return None + return data if isinstance(data, dict) else None + + +def _subdict(parent: dict, key: str) -> dict: + child = parent.get(key) + if not isinstance(child, dict): + child = parent[key] = {} + return child + + +def _remember_key(cache: Path, base: str, key: str, source: str) -> None: + data = _read_cache(cache) + servers = data.get("servers") + if not isinstance(servers, dict): + servers = data["servers"] = {} + buckets = _server_buckets(servers, base) + other = "minted" if source == "saved" else "saved" + buckets[source] = ([key] + [k for k in buckets[source] if k != key])[:8] + buckets[other] = [k for k in buckets[other] if k != key] # a key has one provenance + new_entry = {"saved": buckets["saved"], "minted": buckets["minted"]} + if servers.get(base) == new_entry: + return + servers[base] = new_entry + # Collapse legacy unscoped fields. + data.pop("keys", None) + data.pop("key", None) + try: + _write_private_json(cache, data) + except OSError: + pass # worst case the next launch mints another key + + +def _key_accepted(base: str, key: str) -> bool: + # Only a genuine auth rejection (401/403) means "this key is bad -- skip it and try + # the next cached key or mint a fresh one". A 5xx or a network blip is a server-side + # outage, not a bad key: fail with a clean message (never a traceback) instead of + # silently discarding a working key and minting extras against a struggling server. + try: + _http_json("GET", f"{base}/v1/models", key) + return True + except urllib.error.HTTPError as exc: + if exc.code in (401, 403): + return False + _fail( + f"Studio server error while checking an API key ({exc.code}). " + "The server may be starting up or unhealthy; try again shortly." + ) + except (urllib.error.URLError, TimeoutError) as exc: + _fail( + "Couldn't reach the Studio server while checking an API key: " + f"{getattr(exc, 'reason', None) or exc}" + ) + + +def _agent_api_key( + base: str, + explicit: Optional[str], + *, + auto_started: bool = False, +) -> str: + cache = _key_cache_path() + if explicit: + if not auto_started or _key_accepted(base, explicit): + _remember_key(cache, base, explicit, "saved") + return explicit + # The server was auto-started for this run, so an exported + # UNSLOTH_API_KEY meant for some other server must not fail the + # launch: the loopback mint path below is guaranteed to work. + # (An explicit key that the fresh server accepts, e.g. one persisted + # in this Studio home's auth db, is still honored above.) + + # Replay a key the user saved for *this exact* server first (scoped per base, + # so it only goes back there -- including a remote/SSH-tunnelled Studio whose + # secret the local handshake can't match). Skip ones the server rejects. + for key in _cached_keys(cache, base, "saved"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "saved") + return key + + # Beyond here we auto-mint or replay an auto-minted key. find_studio_server() + # trusts a base after only a health check, so both are limited to a loopback + # server we can cryptographically confirm is ours. + if not is_loopback_url(base): + _fail( + f"No saved API key for {base} and automatic minting only runs against " + "a local Studio. Create an API key in Studio → Settings → API and " + "pass it with --api-key (it is remembered per server), or set " + "UNSLOTH_API_KEY." + ) + if not verify_studio_identity(base): + _fail( + f"Couldn't verify that {base} is your Studio (it may be running as a " + "different OS user, or another process took the port). Create an API " + "key in Studio → Settings → API and pass it with --api-key, or set " + "UNSLOTH_API_KEY." + ) + + # Identity verified: replay a previously auto-minted key, else mint a new one. + for key in _cached_keys(cache, base, "minted"): + if _key_accepted(base, key): + _remember_key(cache, base, key, "minted") + return key + + # Self-issue a JWT (signed with the local secret) and mint a key. + token = _studio_token() + if token is None: + _fail( + "Couldn't authenticate with the Studio server automatically. Create " + "an API key in Studio → Settings → API and pass it with --api-key, " + "or set UNSLOTH_API_KEY." + ) + key = _http_json( + "POST", + f"{base}/api/auth/api-keys", + token, + {"name": "Coding agents (unsloth start)"}, + error = "Couldn't create an API key", + )["key"] + _remember_key(cache, base, key, "minted") + return key + + +def _loaded_models(base: str, key: str) -> list: + return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) + + +def _resolve_model( + base: str, + key: str, + requested: Optional[str], + load: LoadOptions = LoadOptions(), +) -> dict: + models = _loaded_models(base, key) + # /v1/models reports the model id but not the active GGUF variant or runtime load + # settings, so an id match alone can hide the wrong quant (Q8_0 serving while the + # user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to + # /api/inference/load: the server's already-loaded dedup answers "already_loaded" + # without reloading when the variant AND settings match, so a second session running + # the same command still attaches without evicting the first. + load_has_overrides = bool( + load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel + ) + match = ( + None + if requested and load_has_overrides + else next((m for m in models if m["id"] == requested), None) + ) + if requested and match is None: + typer.echo( + f"Ensuring {requested} is loaded with the requested settings…" + if load_has_overrides + else f"Loading {requested} on the Studio server (this can take a while)…" + ) + # Mirror `unsloth run`'s load knobs; keep the default payload as just + # model_path so a bare `--model` load is unchanged. + payload = {"model_path": requested} + if load.gguf_variant: + payload["gguf_variant"] = load.gguf_variant + if load.max_seq_length: + payload["max_seq_length"] = load.max_seq_length + if not load.load_in_4bit: + payload["load_in_4bit"] = False + if load.tensor_parallel: + payload["tensor_parallel"] = True + loaded = _http_json( + "POST", + f"{base}/api/inference/load", + key, + payload, + timeout = 3600, + error = "Model load failed", + ) + # Studio registers the model under a canonical id (resolved identifier, + # casing) that /v1/models echoes but which may differ from the path we + # passed; match on the id the load reports so we don't silently fall + # through to models[0] and connect to a different loaded model. + wanted = {requested} + if isinstance(loaded, dict): + wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} + models = _loaded_models(base, key) + match = next((m for m in models if m["id"] in wanted), None) + if match is not None: + return match + if requested: + # We asked Studio to load it and it didn't surface in /v1/models; don't + # silently hand back an unrelated loaded model. + _fail( + f"Studio didn't report '{requested}' as loaded. Double-check the model " + "id, or load it from the model dropdown in the UI." + ) + if not models: + _fail( + "No model is loaded in Studio. Load one from the model dropdown in " + "the UI, or pass --model to load it from here." + ) + return models[0] + + +def _require_gguf_for_codex(base: str, key: str, model_id: str) -> None: + # Codex always streams, and Studio only streams /v1/responses from llama-server. + try: + status = _http_json("GET", f"{base}/api/inference/status", key) + except urllib.error.HTTPError as exc: + if exc.code == 404: + return # older server without the endpoint; don't block the launch + raise + if status.get("is_gguf"): + return + hint = model_id if "gguf" in model_id.lower() else f"{model_id}-GGUF" + _fail( + f"Codex needs a GGUF model served by llama-server, but {model_id} is on " + f"the transformers backend. Try: unsloth start codex --model {hint}" + ) + + +_DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" +# Session overlay applied via `claude --settings`; suppresses the attribution header +# for THIS run only (no ~/.claude write) so llama.cpp KV-cache reuse is preserved. It +# reinforces the CLAUDE_CODE_ATTRIBUTION_HEADER env var on builds that read the setting +# only from settings.json. +_CLAUDE_SETTINGS_OVERLAY = '{"env":{"CLAUDE_CODE_ATTRIBUTION_HEADER":"0"}}' + + +def _claude_version() -> Optional[tuple]: + # None = no local `claude` (a --no-launch printout for another machine; assume a + # current build). An unparseable version is treated as too old for the new flags. + executable = shutil.which("claude") + if executable is None: + return None + try: + result = subprocess.run( + [executable, "--version"], capture_output = True, text = True, timeout = 10 + ) + # Pull the X.Y.Z out of the output rather than assuming it is the first token. + # claude prints it first today ("2.1.98 (Claude Code)"), but a format change + # (e.g. "claude version 2.1.98") shouldn't silently drop the optimization flags; + # no match falls through to "too old", same as an unparseable version. + match = re.search(r"(\d+)\.(\d+)\.(\d+)", result.stdout) + return tuple(int(part) for part in match.groups()) if match else (0,) + except Exception: + return (0,) + + +def _claude_flags() -> list: + # Both knobs preserve llama.cpp KV-cache reuse: --exclude-dynamic-system-prompt-sections + # moves per-session context out of the system prompt, and --settings suppresses the + # attribution header for this session only (no persistent ~/.claude write; the env var + # sets it too). Claude Code < 2.1.98 aborts on unknown flags, so gate on the version; + # no local binary means a printout for another machine, so assume a current build. + version = _claude_version() + if version is not None and version < (2, 1, 98): + return [] + return [_DYNAMIC_SECTIONS_FLAG, "--settings", _CLAUDE_SETTINGS_OVERLAY] + + +def _merge_codex_config(existing: str, base: str) -> str: + chunks = re.split(r"(?m)^(?=\[)", existing) # preamble, then one chunk per table + if not re.search(r"(?m)^\s*oss_provider\s*=", chunks[0]): + if chunks[0] and not chunks[0].endswith("\n"): + chunks[0] += "\n" + chunks[0] += f'oss_provider = "{_CODEX_PROFILE}"\n' + # Drop the provider table and any stale [model_providers.unsloth_api.*] subtables. + stale = (_PROVIDER_HEADER, _PROVIDER_HEADER[:-1] + ".") + text = "".join(c for c in chunks if not c.startswith(stale)) + if not text.endswith("\n"): + text += "\n" + if not text.endswith("\n\n"): + text += "\n" + return text + ( + f"{_PROVIDER_HEADER}\n" + 'name = "Unsloth Studio"\n' + f"base_url = {json.dumps(base + '/v1')}\n" + f'env_key = "{_CODEX_ENV_KEY}"\n' + 'wire_api = "responses"\n' + "requires_openai_auth = false\n" + ) + + +def write_codex_config(base: str, model: dict, home: Path) -> None: + home.mkdir(parents = True, exist_ok = True) + + config = home / "config.toml" + existing = config.read_text(encoding = "utf-8") if config.exists() else "" + merged = _merge_codex_config(existing, base) + if merged != existing: + config.write_text(merged, encoding = "utf-8") + typer.echo(f"Updated {config}") + + # oss_provider here too: codex --oss picks the provider from it, and the + # profile layer must beat a user-set value (e.g. "ollama") in config.toml. + profile_text = ( + f'oss_provider = "{_CODEX_PROFILE}"\n' + f'model_provider = "{_CODEX_PROFILE}"\n' + f"model = {json.dumps(model['id'])}\n" + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + profile_text += f"model_context_window = {int(window)}\n" + profile = home / f"{_CODEX_PROFILE}.config.toml" + if not profile.exists() or profile.read_text(encoding = "utf-8") != profile_text: + profile.write_text(profile_text, encoding = "utf-8") + typer.echo(f"Updated {profile}") + + +def _wsl_windows_executable(command: list) -> Optional[str]: + if os.name == "nt" or not os.environ.get("WSL_DISTRO_NAME"): + return None + executable = shutil.which(command[0]) + if executable and executable.startswith("/mnt/"): + return executable + return None + + +def _looks_like_path(value: str) -> bool: + # A var only wants the WSLENV /p flag if its value is a filesystem path: an + # absolute POSIX path (/...), a UNC path (\\...), or a drive-qualified Windows + # path (C:...). Scalar knobs (e.g. a numeric context window) must pass through + # untranslated, so they get no flag. + return bool(value) and (value.startswith(("/", "\\")) or (len(value) >= 2 and value[1] == ":")) + + +def _wsl_bridge_names(env: dict, unset_env: tuple) -> tuple: + # Build the WSLENV share list for a Windows shim reached from WSL. Path-valued + # vars get /p so WSLENV translates them to the Windows path the /mnt shim can + # actually open; a cleared var carries no value to translate. + names = [name + ("/p" if _looks_like_path(value) else "") for name, value in env.items()] + names.extend(unset_env) + return tuple(dict.fromkeys(names)) + + +def _merge_wslenv(current: str, names: tuple) -> str: + # Index WSLENV entries by bare var name, preserving first-seen order. The vars we + # bridge are applied last so our entry wins: a user's pre-existing unflagged "HOME" + # is upgraded to "HOME/p" (rather than left as-is), since WSLENV ignores a duplicate + # name and a bare entry would leave the path untranslated for a Windows shim. + ordered = [] + by_name = {} + for entry in (*current.split(":"), *names): + if not entry: + continue + base = entry.split("/", 1)[0] + if base not in by_name: + ordered.append(base) + by_name[base] = entry + return ":".join(by_name[base] for base in ordered) + + +def _powershell_quote(arg: str) -> str: + # PowerShell reads single-quoted strings literally (an embedded ' is doubled), so + # JSON args such as `--settings {"env":...}` survive intact. list2cmdline's + # backslash-escaped double quotes are cmd.exe syntax and PowerShell mis-parses them. + if arg and re.fullmatch(r"[A-Za-z0-9_./:=+-]+", arg): + return arg + return "'" + arg.replace("'", "''") + "'" + + +def _print_env( + env: dict, + command: list, + unset_env: tuple = (), + wsl_env_bridge: tuple = (), +) -> None: + if os.name == "nt": + for name in unset_env: + typer.echo(f"Remove-Item Env:{name} -ErrorAction SilentlyContinue") + for name, value in env.items(): + # PowerShell: ` is the escape char, and $ triggers expansion inside "". + escaped = value.replace("`", "``").replace('"', '`"').replace("$", "`$") + typer.echo(f'$env:{name} = "{escaped}"') + typer.echo(" ".join(_powershell_quote(arg) for arg in command)) + return + for name in unset_env: + typer.echo(f"export {name}=" if wsl_env_bridge else f"unset {name}") + for name, value in env.items(): + typer.echo(f"export {name}={shlex.quote(value)}") + if wsl_env_bridge: + typer.echo( + f"export WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + # The final line is a SELF-CONTAINED one-liner (inline env, VAR=... cmd) rather than a + # bare command. People copy just the last line, and a bare `codex`/`claude` would then + # run against their real ~/.codex or Anthropic credentials with zero isolation -- e.g. + # inheriting a pre-existing damaged ~/.codex state DB and blaming the recipe. Inline + # assignments scope every var (and empty-string the conflicting ones) to this single + # invocation, so a partial copy behaves the same as pasting the whole block. + inline = [f"{name}=" for name in unset_env] + inline += [f"{name}={shlex.quote(value)}" for name, value in env.items()] + if wsl_env_bridge: + inline.append( + f"WSLENV={shlex.quote(_merge_wslenv(os.environ.get('WSLENV', ''), wsl_env_bridge))}" + ) + typer.echo(" ".join((*inline, shlex.join(command)))) + + +def _install_agent(name: str, install_hint: str) -> Optional[str]: + # Missing agent under --launch: offer to run its documented install command, then + # re-resolve it on PATH. Consent-based (we never auto-run a remote install script + # silently), and a non-interactive stdin cannot answer the prompt, so both the + # no-TTY and declined cases return None and let the caller print the hint and exit. + if not sys.stdin.isatty(): + return None + typer.echo(f"`{name}` is not installed.") + if not typer.confirm(f"Install it now with `{install_hint}`?", default = False): + return None + # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm) + # on Windows, /bin/sh (curl | bash, or npm) everywhere else. + if os.name == "nt": + install_command = ["powershell", "-NoProfile", "-Command", install_hint] + else: + install_command = ["/bin/sh", "-c", install_hint] + if subprocess.run(install_command).returncode != 0: + _fail(f"Install command failed. Run it yourself, then re-run: {install_hint}") + executable = shutil.which(name) + if executable is None: + _fail( + f"`{name}` installed but isn't on PATH yet. Open a new shell (or add it to " + f"PATH), then re-run. Install command: {install_hint}" + ) + return executable + + +def _launch( + command: list, + env: dict, + install_hint: str, + unset_env: tuple = (), +) -> NoReturn: + executable = shutil.which(command[0]) or _install_agent(command[0], install_hint) + if executable is None: + _fail(f"`{command[0]}` not found on PATH. Install it with: {install_hint}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + child_env = dict(os.environ) + if wsl_env_bridge: + child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_env_bridge) + for name in unset_env: + child_env[name] = "" + else: + for name in unset_env: + child_env.pop(name, None) + child_env.update(env) + # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. + previous = signal.signal(signal.SIGINT, signal.SIG_IGN) + try: + code = subprocess.run([executable, *command[1:]], env = child_env).returncode + finally: + signal.signal(signal.SIGINT, previous) + # Negative returncode means killed by signal N; shells expect 128+N. + raise typer.Exit(code = code if code >= 0 else 128 - code) + + +def _connect( + api_key: Optional[str], + model: Optional[str], + load: LoadOptions = LoadOptions(), + *, + serve: bool = False, + launch: bool = True, +) -> tuple: + # `--model org/name:QUANT` is shorthand for `--model org/name --gguf-variant QUANT`. + # Split it before we match/serve so the attach path resolves against the already-loaded + # `org/name` (listed without the suffix) instead of reloading a `:`-suffixed repo id -- + # which Studio rejects and which would evict a model another session is using. + if model: + repo, variant = _split_repo_variant(model) + if variant: + model = repo + if not load.gguf_variant: + load = load._replace(gguf_variant = variant) + base, server = _require_studio(model, load, serve = serve, launch = launch) + try: + key = _agent_api_key(base, api_key, auto_started = server is not None) + # A server we just started has exactly the requested model loaded, so resolve to + # whatever it is serving instead of re-matching the raw --model string. + entry = _resolve_model(base, key, None if server is not None else model, load) + except BaseException: + _shutdown_auto_served() + raise + return base, key, entry + + +def _run( + base: str, + entry: dict, + env: dict, + command: list, + *, + launch: bool, + install_hint: str, + unset_env: tuple = (), + clear_screen: bool = False, +) -> None: + # Some agents (Pi) render inline from wherever the cursor sits: their first + # paint assumes a clean screen rather than clearing or entering the + # alternate screen themselves. Hand them one so the session doesn't start + # mid-scroll under our connection output. click.clear() is cross-platform + # and a no-op when stdout is not a terminal (piped/CI), so transcripts and + # --no-launch recipes stay intact. + if launch and clear_screen: + click.clear() + typer.echo(f"Studio {base} · model {entry['id']}") + wsl_env_bridge = _wsl_bridge_names(env, unset_env) if _wsl_windows_executable(command) else () + if not launch: + _print_env(env, command, unset_env = unset_env, wsl_env_bridge = wsl_env_bridge) + return + try: + _launch(command, env, install_hint = install_hint, unset_env = unset_env) + finally: + # Tear down a server we auto-started once the agent session ends (no-op otherwise). + _shutdown_auto_served() + + +def _agents_config_root() -> Path: + ensure_studio_backend_path() + from utils.paths import auth_root + return auth_root() / "agents" + + +@contextlib.contextmanager +def _session_config(agent: str, launch: bool): + """Yield a private directory for an agent's session config (never the user's own). + + launch: an ephemeral temp dir removed after the agent process exits, so nothing + persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later + on this machine), reused across runs. Either way the user's real ~/. + config is left untouched. + """ + if launch: + path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-")) + try: + yield path + finally: + shutil.rmtree(path, ignore_errors = True) + else: + # Never wipe this dir: a previously printed recipe may still be running + # an agent whose sessions/state live here, and every config writer + # merges idempotently into an existing home anyway. + path = _agents_config_root() / agent + path.mkdir(parents = True, exist_ok = True, mode = 0o700) + yield path + + +def write_openclaw_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Studio is a generic OpenAI-compatible /v1 endpoint (the vLLM/LM Studio path). + provider_model = {"id": model["id"], "name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + provider_model["contextWindow"] = int(window) + models = _subdict(config, "models") + models.setdefault("mode", "merge") + _subdict(models, "providers")["unsloth"] = { + "baseUrl": f"{base}/v1", + "apiKey": key, + "api": "openai-completions", + "models": [provider_model], + } + # Pin a default model, else OpenClaw drops into its setup agent ("no models available"). + defaults = _subdict(_subdict(config, "agents"), "defaults") + _subdict(defaults, "model")["primary"] = f"unsloth/{model['id']}" + # Unauthenticated loopback gateway: without auth.mode=none the client won't open + # the websocket. The daemon must still be started separately (`openclaw gateway`). + gateway = _subdict(config, "gateway") + gateway.setdefault("mode", "local") + _subdict(gateway, "auth").setdefault("mode", "none") + if yolo: + # OpenClaw has no --yolo flag, and it gates tool execution on BOTH the + # tools.exec config AND a host-local approvals file (the stricter wins), so + # setting only the config still lets the agent prompt/deny. Set both, mirroring + # `openclaw exec-policy preset yolo`. + exec_policy = _subdict(_subdict(config, "tools"), "exec") + exec_policy["host"] = "gateway" + exec_policy["security"] = "full" + exec_policy["ask"] = "off" + # Approvals file in OPENCLAW_STATE_DIR (== this config's dir). ask=off means + # nothing is ever prompted, so the runtime socket block is unnecessary here. + approvals = path.parent / "exec-approvals.json" + _write_private_json( + approvals, + {"version": 1, "defaults": {"security": "full", "ask": "off", "askFallback": "full"}}, + ) + typer.echo(f"Updated {approvals}") + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_opencode_config( + base: str, + key: str, + model: dict, + path: Path, + yolo: bool = False, +) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + config.setdefault("$schema", "https://opencode.ai/config.json") + model_entry = {"name": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # A custom-provider model with no limit defaults to context 0, which silently + # disables OpenCode's auto-compaction; declare the real window (and a sane + # output cap) so it compacts instead of overflowing the server. + model_entry["limit"] = {"context": window, "output": min(window // 4, 8192)} + _subdict(config, "provider")["unsloth"] = { + "npm": "@ai-sdk/openai-compatible", + "name": "Unsloth Studio", + "options": {"baseURL": f"{base}/v1", "apiKey": key}, + "models": {model["id"]: model_entry}, + } + # OpenCode selects a model by "/". + config["model"] = f"unsloth/{model['id']}" + if window: + # Compact with ~10% headroom (near 90% full). The fixed 20k-token default + # buffer over-compacts, or never settles, on a small local context. + compaction = _subdict(config, "compaction") + compaction["auto"] = True + compaction["reserved"] = max(1, window // 10) + if yolo: + # OpenCode has no --yolo flag; auto-approve is the config `permission` block + # (singular). Allow the prompting tools so tool calls don't block on the TUI. + config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +def write_hermes_config(base: str, model: dict, path: Path) -> None: + import yaml + + config: dict = {} + if path.exists(): + try: + loaded = yaml.safe_load(path.read_text(encoding = "utf-8")) + except (yaml.YAMLError, OSError): + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + if isinstance(loaded, dict): + config = loaded + elif loaded is not None: + # Non-empty, non-mapping YAML is a user-managed file; leave it. + typer.echo( + f"Warning: couldn't parse {path} — configure the custom endpoint " + "there yourself, or move the file aside and re-run.", + err = True, + ) + return + # Hermes only reads the key for a *named* custom provider (a bare + # `provider: custom` ignores it), so register it under providers.*. + _subdict(config, "model").update( + provider = f"custom:{_HERMES_PROVIDER}", + default = model["id"], + api_mode = "openai", + ) + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # Hermes auto-detects context from GET /v1/models, but OpenAI's schema has no + # context field, so it can fall back to a 256k default that overflows a small + # local model. Pin the real window (top-level model.context_length is the + # highest-priority override) and compact at 90% of it (Hermes defaults to 50%). + if window >= _HERMES_MIN_CONTEXT: + _subdict(config, "model")["context_length"] = window + _subdict(config, "compression").update(enabled = True, threshold = 0.9) + else: + # Below Hermes' 64,000-token floor it refuses to initialize, so claim + # the floor and shrink the threshold so compaction still fires at 90% + # of the REAL window (the threshold is a fraction of the claimed + # context_length). The auxiliary override keeps the same floor check + # from rejecting the compression model mid-session. + _subdict(config, "model")["context_length"] = _HERMES_MIN_CONTEXT + threshold = round(0.9 * window / _HERMES_MIN_CONTEXT, 4) + _subdict(config, "compression").update(enabled = True, threshold = threshold) + auxiliary = _subdict(_subdict(config, "auxiliary"), "compression") + auxiliary["context_length"] = _HERMES_MIN_CONTEXT + _subdict(config, "providers")[_HERMES_PROVIDER] = { + "base_url": f"{base}/v1", + "api_mode": "openai", + "key_env": _HERMES_ENV_KEY, + } + text = yaml.safe_dump(config, sort_keys = False) + if not path.exists() or path.read_text(encoding = "utf-8") != text: + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(text, encoding = "utf-8") + typer.echo(f"Updated {path}") + + +def write_pi_config(base: str, key: str, model: dict, path: Path) -> None: + config = _read_json_object(path) + if config is None: + typer.echo( + f"Warning: couldn't parse {path} — add an 'unsloth' provider there " + "yourself, or move the file aside and re-run.", + err = True, + ) + return + before = json.dumps(config, sort_keys = True) + # Pi reads custom providers from ~/.pi/agent/models.json (HOME-relocated for the + # session). Studio is a generic OpenAI-compatible /v1 endpoint, and the key lives + # in the config rather than the env (matching openclaw/opencode). + provider_model = {"id": model["id"]} + window = model.get("context_length") or model.get("max_context_length") + if window: + window = int(window) + # An unspecified model defaults to contextWindow 128000 / maxTokens 16384, + # far larger than a small Studio context, so Pi compacts too late and overflows + # the server. Pin the real window and a sane output cap (mirrors OpenCode). + provider_model["contextWindow"] = window + provider_model["maxTokens"] = min(window // 4, 8192) + _subdict(config, "providers")[_PI_PROVIDER] = { + "api": "openai-completions", + "baseUrl": f"{base}/v1", + "apiKey": key, + "models": [provider_model], + } + if json.dumps(config, sort_keys = True) != before: + _write_private_json(path, config) + typer.echo(f"Updated {path}") + + +@start_app.command("claude", context_settings = _PASSTHROUGH) +def claude( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Claude Code at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + model_id = entry["id"] + + env = { + "ANTHROPIC_BASE_URL": base, + "ANTHROPIC_AUTH_TOKEN": key, + "ANTHROPIC_MODEL": model_id, + # Session-only (no ~/.claude write): suppress the attribution header so + # llama.cpp KV-cache reuse is preserved; --settings below reinforces it. + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + # Update checks, beta features, and other background requests either + # stall against a local server or evict the conversation from + # llama-server's KV-cache slots, so turn off everything nonessential. + "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", + "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", + # A local server streams in bursts; disable the full-screen TUI redraw so the + # terminal doesn't flicker between tokens. + "CLAUDE_CODE_NO_FLICKER": "1", + } + # Claude Code auto-compacts against its native (~600k token) window; a local + # model's context is usually far smaller, so size the window to the loaded + # model's real context length. Otherwise the conversation overflows the + # server's window (silent truncation) long before Claude decides to compact. + # codex/openclaw get the same value through their config (model_context_window + # / contextWindow); Claude has no config file, so it rides on the env var. + window = entry.get("context_length") or entry.get("max_context_length") + if window: + env["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = str(int(window)) + # Compact at 90% of that window; the override only takes effect once the + # window is set, and it can only lower the threshold, so it just guarantees + # headroom before the server's context limit instead of relying on Claude's + # default (which is tuned for its native 200K/1M window). + env["CLAUDE_AUTOCOMPACT_PCT_OVERRIDE"] = "90" + # --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions. + # IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a + # sandbox is detected, and we don't want to falsely claim one on the user's host. + command = [ + "claude", + "--model", + model_id, + *_claude_flags(), + *_yolo_command_flags("claude", yolo), + *ctx.args, + ] + install_hint = ( + "irm https://claude.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://claude.ai/install.sh | bash" + ) + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + unset_env = _CLAUDE_ENV_UNSET, + ) + + +@start_app.command("codex", context_settings = _PASSTHROUGH) +def codex( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenAI Codex at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # This preflight runs after _connect may have auto-started a server but before _run + # installs its teardown finally, so tear the server down here if it rejects the model + # (e.g. a transformers-backend model) rather than leaving it on the atexit backstop. + try: + _require_gguf_for_codex(base, key, entry["id"]) + except BaseException: + _shutdown_auto_served() + raise + command = [ + "codex", + "--oss", + "--profile", + _CODEX_PROFILE, + *_yolo_command_flags("codex", yolo), + *ctx.args, + ] + with _session_config("codex", launch) as home: + write_codex_config(base, entry, home) + env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") + + +@start_app.command("openclaw", context_settings = _PASSTHROUGH) +def openclaw( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenClaw at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["openclaw", *ctx.args] + install_hint = ( + "iwr -useb https://openclaw.ai/install.ps1 | iex" + if os.name == "nt" + else "curl -fsSL https://openclaw.ai/install.sh | bash" + ) + with _session_config("openclaw", launch) as cfg: + config_path = cfg / "openclaw.json" + # key lives in the config, not the env; --yolo writes the exec policy here too. + write_openclaw_config(base, key, entry, config_path, yolo = yolo) + # Scope both config and state so OpenClaw never touches the user's ~/.openclaw. + env = {"OPENCLAW_CONFIG_PATH": str(config_path), "OPENCLAW_STATE_DIR": str(cfg)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("opencode", context_settings = _PASSTHROUGH) +def opencode( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point OpenCode at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["opencode", *ctx.args] + with _session_config("opencode", launch) as cfg: + config_path = cfg / "opencode.json" + # OPENCODE_CONFIG is an overlay (loaded between the user's global and project + # configs), so this adds the Unsloth provider/model for the session without + # changing the user's default model. Key lives in the config, not the env. + write_opencode_config(base, key, entry, config_path, yolo = yolo) + # A project's own opencode.json outranks OPENCODE_CONFIG, so the session model + # pin (and --yolo permissions) would silently lose to a repo config. Carry the + # settings that must win in OPENCODE_CONFIG_CONTENT, which outranks project + # config; the API key stays in the private file, never in the printed env. + inline_config: dict = {"model": f"unsloth/{entry['id']}"} + if yolo: + inline_config["permission"] = {"edit": "allow", "bash": "allow", "webfetch": "allow"} + env = { + "OPENCODE_CONFIG": str(config_path), + "OPENCODE_CONFIG_CONTENT": json.dumps(inline_config), + } + _run(base, entry, env, command, launch = launch, install_hint = "npm install -g opencode-ai") + + +@start_app.command("hermes", context_settings = _PASSTHROUGH) +def hermes( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Hermes (Nous Research) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] + install_hint = ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash" + ) + with _session_config("hermes", launch) as home: + # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) + # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. + write_hermes_config(base, entry, home / "config.yaml") + env = {_HERMES_ENV_KEY: key, "HERMES_HOME": str(home)} + _run(base, entry, env, command, launch = launch, install_hint = install_hint) + + +@start_app.command("pi", context_settings = _PASSTHROUGH) +def pi( + ctx: typer.Context, + model: Optional[str] = _MODEL_OPTION, + api_key: Optional[str] = _KEY_OPTION, + launch: bool = _LAUNCH_OPTION, + gguf_variant: Optional[str] = _GGUF_VARIANT_OPTION, + max_seq_length: int = _CONTEXT_OPTION, + load_in_4bit: bool = _LOAD_4BIT_OPTION, + tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, + serve: bool = _SERVE_OPTION, + yolo: bool = _YOLO_OPTION, +): + """Point Pi (coding agent) at the running Studio server and start it.""" + base, key, entry = _connect( + api_key, + model, + LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), + serve = serve, + launch = launch, + ) + # Pi defaults to the google provider, so pin our provider/model on the command + # line; the custom OpenAI-compatible endpoint itself is only configurable via + # ~/.pi/agent/models.json. + command = [ + "pi", + "--provider", + _PI_PROVIDER, + "--model", + entry["id"], + *_yolo_command_flags("pi", yolo), + *ctx.args, + ] + # --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs + # no install scripts), so accepting the prompt skips dependency lifecycle scripts. + install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent" + with _session_config("pi", launch) as home: + # Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers + # it over $HOME/.pi/agent), so pin it at the session dir: an inherited + # PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real + # config and skip our provider/key. HOME is relocated too so any other ~/.pi paths + # stay in the session. The key rides in the config rather than the env. + pi_agent_dir = home / ".pi" / "agent" + write_pi_config(base, key, entry, pi_agent_dir / "models.json") + env = {"HOME": str(home), "PI_CODING_AGENT_DIR": str(pi_agent_dir)} + if os.name == "nt" or os.environ.get("WSL_DISTRO_NAME"): + # Node resolves ~/.pi via USERPROFILE (then HOMEDRIVE + HOMEPATH) on Windows, + # not HOME. Set them whenever Pi may run as a Windows process: native Windows, + # or a /mnt Windows shim launched from WSL (the WSLENV bridge then translates + # the path). Otherwise the Windows process falls back to the user's real + # %USERPROFILE%\.pi. splitdrive yields no drive off a POSIX path, so + # HOMEDRIVE/HOMEPATH stay unset there. + env["USERPROFILE"] = str(home) + drive, tail = os.path.splitdrive(str(home)) + if drive: + env["HOMEDRIVE"], env["HOMEPATH"] = drive, tail + # Pi paints inline from the current cursor position (no alternate screen, + # no clear on first render), so give it the clean screen it assumes. + _run( + base, + entry, + env, + command, + launch = launch, + install_hint = install_hint, + clear_screen = True, + ) diff --git a/unsloth_cli/tests/test_connect.py b/unsloth_cli/tests/test_connect.py deleted file mode 100644 index e76a892647..0000000000 --- a/unsloth_cli/tests/test_connect.py +++ /dev/null @@ -1,954 +0,0 @@ -# SPDX-License-Identifier: AGPL-3.0-only -# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - -"""Tests for `unsloth connect` — config merging and launch env, no network.""" - -from __future__ import annotations - -import json -import os -import sys -import urllib.error -from pathlib import Path -from types import SimpleNamespace - -_REPO_ROOT = Path(__file__).resolve().parents[2] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - - -import pytest -from typer.testing import CliRunner - -import unsloth_cli.commands.connect as connect - -BASE = "http://127.0.0.1:8888" -MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} - - -# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and -# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. -def _assert_env_set(output: str, name: str, value: str) -> None: - needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -def _assert_env_unset(output: str, name: str) -> None: - needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" - assert needle in output, f"{needle!r} not found in:\n{output}" - - -@pytest.fixture() -def claude_settings(tmp_path, monkeypatch): - path = tmp_path / "claude" / "settings.json" - monkeypatch.setattr(connect, "claude_settings_path", lambda: path) - return path - - -def test_claude_settings_created_when_missing(claude_settings): - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_merge_preserves_existing(claude_settings): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text( - json.dumps({"effortLevel": "high", "env": {"CLAUDE_CODE_ENABLE_TELEMETRY": "0"}}) - ) - connect.ensure_claude_attribution_header() - settings = json.loads(claude_settings.read_text()) - assert settings["effortLevel"] == "high" - assert settings["env"]["CLAUDE_CODE_ENABLE_TELEMETRY"] == "0" - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_claude_settings_already_set_untouched(claude_settings): - claude_settings.parent.mkdir(parents = True) - original = json.dumps({"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}}) - claude_settings.write_text(original) - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == original - - -def test_claude_settings_bad_json_left_alone(claude_settings, capsys): - claude_settings.parent.mkdir(parents = True) - claude_settings.write_text("{not json") - connect.ensure_claude_attribution_header() - assert claude_settings.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def _fake_claude(monkeypatch, version_output: str) -> None: - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr( - connect.subprocess, - "run", - lambda *args, **kwargs: SimpleNamespace(stdout = version_output), - ) - - -def test_cache_flags_passed_to_supported_claude(monkeypatch): - _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") - assert connect._claude_cache_flags() == ["--exclude-dynamic-system-prompt-sections"] - - -def test_cache_flags_skipped_on_old_claude(monkeypatch): - _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") - assert connect._claude_cache_flags() == [] - - -def test_cache_flags_skipped_on_unparseable_version(monkeypatch): - _fake_claude(monkeypatch, "weird build string\n") - assert connect._claude_cache_flags() == [] - - -def _parse_toml(text: str) -> dict: - tomllib = pytest.importorskip("tomllib") - return tomllib.loads(text) - - -def test_merge_codex_config_fresh(): - merged = connect._merge_codex_config("", BASE) - parsed = _parse_toml(merged) - assert parsed["oss_provider"] == "unsloth_api" - provider = parsed["model_providers"]["unsloth_api"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["wire_api"] == "responses" - assert provider["requires_openai_auth"] is False - - -def test_merge_codex_config_replaces_stale_block(): - existing = ( - 'model = "gpt-5"\n' - "\n" - "[model_providers.unsloth_api]\n" - 'base_url = "http://old-host:9999/v1"\n' - 'wire_api = "chat"\n' - "\n" - "[model_providers.unsloth_api.http_headers]\n" - 'x-old = "1"\n' - "\n" - "[model_providers.ollama]\n" - 'base_url = "http://localhost:11434/v1"\n' - ) - merged = connect._merge_codex_config(existing, BASE) - parsed = _parse_toml(merged) - assert parsed["model"] == "gpt-5" - assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" - assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" - assert "http_headers" not in parsed["model_providers"]["unsloth_api"] - assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" - assert connect._merge_codex_config(merged, BASE) == merged - - -def test_merge_codex_config_keeps_user_oss_provider(): - merged = connect._merge_codex_config('oss_provider = "ollama"\n', BASE) - assert _parse_toml(merged)["oss_provider"] == "ollama" - - -def test_write_codex_config_profile(tmp_path, monkeypatch): - monkeypatch.setenv("CODEX_HOME", str(tmp_path)) - connect.write_codex_config(BASE, MODEL) - profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) - assert profile["oss_provider"] == "unsloth_api" - assert profile["model_provider"] == "unsloth_api" - assert profile["model"] == MODEL["id"] - assert profile["model_context_window"] == 131072 - config = _parse_toml((tmp_path / "config.toml").read_text()) - assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" - - -@pytest.fixture() -def fake_studio(tmp_path, monkeypatch, claude_settings): - calls = [] - state = {"models": [MODEL]} - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - calls.append((method, url, payload)) - if url.endswith("/v1/models"): - return {"object": "list", "data": state["models"]} - if url.endswith("/api/inference/status"): - return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} - if url.endswith("/api/auth/api-keys"): - return {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/inference/load"): - state["models"] = [{"id": payload["model_path"], "context_length": 4096}] - return {} - raise AssertionError(f"unexpected request: {method} {url}") - - monkeypatch.setattr(connect, "find_studio_server", lambda: BASE) - # Identity handshake has its own tests; trust the loopback server here. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: True) - # _studio_token / api-keys are faked so the mint flow stays offline. - monkeypatch.setattr(connect, "_studio_token", lambda: "jwt-token") - monkeypatch.setattr(connect, "_http_json", http_json) - monkeypatch.setattr(connect, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") - # No `claude` on PATH, so _claude_cache_flags never probes the real binary. - monkeypatch.setattr(connect.shutil, "which", lambda _: None) - monkeypatch.setenv("CODEX_HOME", str(tmp_path / "codex")) - monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) - return calls - - -def test_connect_claude_no_launch(fake_studio, claude_settings): - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_unset(result.output, "ANTHROPIC_API_KEY") - _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") - _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") - _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") - assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output - settings = json.loads(claude_settings.read_text()) - assert settings["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" - - -def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr(connect.shutil, "which", lambda _: "/usr/local/bin/claude") - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] - assert "ANTHROPIC_API_KEY" not in captured["env"] - assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): - captured = {} - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") - monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - monkeypatch.setattr(connect, "_claude_cache_flags", lambda: []) - - def run(command, env): - captured["command"] = command - captured["env"] = env - return SimpleNamespace(returncode = 0) - - monkeypatch.setattr(connect.subprocess, "run", run) - result = CliRunner().invoke(connect.connect_app, ["claude"]) - - assert result.exit_code == 0, result.output - assert captured["command"] == [ - "/mnt/c/Users/samle/AppData/Roaming/npm/claude", - "--model", - MODEL["id"], - ] - assert captured["env"]["ANTHROPIC_API_KEY"] == "" - assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" - assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" - assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE - assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] - for name in ( - "ANTHROPIC_AUTH_TOKEN", - "ANTHROPIC_BASE_URL", - "ANTHROPIC_MODEL", - "ANTHROPIC_API_KEY", - "CLAUDE_CODE_OAUTH_TOKEN", - ): - assert name in captured["env"]["WSLENV"].split(":") - - -@pytest.mark.skipif( - os.name == "nt", - reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " - "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", -) -def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): - monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setattr( - connect.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" - ) - - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - - assert result.exit_code == 0, result.output - assert "export ANTHROPIC_API_KEY=" in result.output - assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output - assert "export WSLENV=" in result.output - assert "ANTHROPIC_AUTH_TOKEN" in result.output - assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output - - -def test_connect_codex_no_launch(fake_studio, tmp_path): - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - assert "codex --oss --profile unsloth_api" in result.output - assert (tmp_path / "codex" / "config.toml").exists() - assert (tmp_path / "codex" / "unsloth_api.config.toml").exists() - - -def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - # First run mints; second reuses the minted key cached for this server. - mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] - assert len(mints) == 1 - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - - -def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): - CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - # Reused, not re-minted (a mint would return the feedface stand-in). - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - # An explicit key is remembered as "saved" so it replays without the handshake. - assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] - - -def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): - cache = tmp_path / "agent_api_key.json" - cache.write_text( - json.dumps( - {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} - ) - ) - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/v1/models") and token == "sk-unsloth-stale": - raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - # The working key moves to the front so the next run tries it first. - cached = json.loads(cache.read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] - - -def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): - # Legacy unscoped caches have no server binding (could leak across servers), - # so they're ignored: a fresh key is minted and stored scoped to this server. - (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") - cached = json.loads((tmp_path / "agent_api_key.json").read_text()) - assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] - assert "key" not in cached # legacy field collapsed away - - -def test_connect_model_flag_loads_on_server(fake_studio): - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] - ) - assert result.exit_code == 0, result.output - loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] - assert loads == [ - ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) - ] - _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") - - -def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): - # Studio registers a loaded model under a canonical id (resolved identifier - # / casing) that can differ from the path we passed. The agent must connect - # to that model, not silently fall through to the first loaded one. - requested = "Unsloth/Qwen3.5-35B-A3B" - canonical = "unsloth/Qwen3.5-35B-A3B" - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {"model": canonical, "display_name": canonical} - if url.endswith("/v1/models"): - # Decoy sorts first, so models[0] is the wrong pick on the old code. - return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", requested] - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) - - -def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): - monkeypatch.setattr( - connect, - "_http_json", - lambda method, url, token, payload = None, timeout = 30, error = None: ( - {"key": "sk-unsloth-feedfacefeedface"} - if url.endswith("/api/auth/api-keys") - else {"object": "list", "data": []} - ), - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No model is loaded" in result.output - - -def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): - # Studio never surfaces the requested model; fail loudly rather than - # silently connecting to whatever else happens to be loaded. - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/load"): - return {} - if url.endswith("/v1/models"): - return {"object": "list", "data": [MODEL]} # decoy; request never appears - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke( - connect.connect_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] - ) - assert result.exit_code == 1 - assert "unsloth/Missing-7B" in result.output - - -def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): - inner = connect._http_json - - def http_json( - method, - url, - token, - payload = None, - timeout = 30, - error = None, - ): - if url.endswith("/api/inference/status"): - return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} - return inner(method, url, token, payload, timeout, error) - - monkeypatch.setattr(connect, "_http_json", http_json) - result = CliRunner().invoke(connect.connect_app, ["codex", "--no-launch"]) - assert result.exit_code == 1 - assert "GGUF" in result.output - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): - # A server known only by URL + health check is unverified: keyless connect - # must refuse and make no request at all. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.evil.example:8888") - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 1 - assert "Settings → API" in result.output - assert "--api-key" in result.output - assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) - - -def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): - # User named both server and key, so it's their choice; only auto-send is blocked. - monkeypatch.setattr(connect, "find_studio_server", lambda: "http://studio.example:8888") - result = CliRunner().invoke( - connect.connect_app, - ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - - -def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): - # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; - # auto-minting stays blocked for non-loopback. - remote = "http://studio.example:8888" - monkeypatch.setattr(connect, "find_studio_server", lambda: remote) - (tmp_path / "agent_api_key.json").write_text( - json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) - ) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): - # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an - # error, not a silent local model load (which they did not ask for). - import typer - - import unsloth_cli._inference as inference - - monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") - monkeypatch.setattr( - inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" - ) - with pytest.raises(typer.Exit): - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - - -def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): - # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback - # server can't be verified, fall back to a local load rather than erroring. - import unsloth_cli._inference as inference - - monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) - monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") - monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) - assert ( - inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) - is None - ) - - -def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( - fake_studio, tmp_path, monkeypatch -): - # With no saved key, the next step would auto-mint; an unverified loopback - # server (port squatter) must be refused, with nothing sent. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted - - -def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): - # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) - # replays on keyless runs without the handshake, scoped to its own base. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted - - -def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): - # A "minted" key is NOT replayed to an unverified loopback server: minting and - # minted-key replay both sit behind the handshake, so a squatter can't grab it. - cache = tmp_path / "agent_api_key.json" - cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "--api-key" in result.output - assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent - - -def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): - # An explicit key is the user's deliberate choice, so it does not require - # the automatic identity handshake. - monkeypatch.setattr(connect, "verify_studio_identity", lambda base: False) - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - - -def _serve_identity(proof_for): - """Start a localhost HTTP server answering /api/auth/identity with - proof_for(nonce_bytes). Returns (base_url, shutdown).""" - import base64 - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - from urllib.parse import parse_qs, urlparse - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - parsed = urlparse(self.path) - if parsed.path != "/api/auth/identity": - self.send_response(404) - self.end_headers() - return - nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) - host, port = self.server.server_address[0], self.server.server_address[1] - body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(body) - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): - # Real crypto end to end: verify_studio_identity reads the install secret from - # an isolated DB; a "good" server proves the same secret, a spoofing one can't. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: # backend not importable here (e.g. missing deps) - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - good = lambda nonce, host, port: storage.compute_identity_proof( - nonce, host, port - ) # real secret - bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret - base_ok, stop_ok = _serve_identity(good) - base_bad, stop_bad = _serve_identity(bad) - try: - assert inference.verify_studio_identity(base_ok) is True - assert inference.verify_studio_identity(base_bad) is False - finally: - stop_ok() - stop_bad() - - -def _serve_redirect(target): - """Start a localhost server that 302-redirects every GET to target+path.""" - import threading - from http.server import BaseHTTPRequestHandler, HTTPServer - - class Handler(BaseHTTPRequestHandler): - def do_GET(self): - self.send_response(302) - self.send_header("Location", target + self.path) - self.end_headers() - - def log_message(self, *a): - pass - - server = HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target = server.serve_forever, daemon = True).start() - base = f"http://127.0.0.1:{server.server_address[1]}" - return base, server.shutdown - - -def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): - # A squatter could 302 /api/auth/identity to the real Studio and relay its - # proof; redirects must be refused so the squatter's base isn't accepted. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - squatter_base, stop_squatter = _serve_redirect(real_base) - try: - assert inference.verify_studio_identity(real_base) is True # direct: ok - assert inference.verify_studio_identity(squatter_base) is False # relayed: refused - finally: - stop_real() - stop_squatter() - - -def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): - # A squatter that proxies the nonce to the real Studio on another port gets a - # proof bound to *that* port; the client expects one bound to the port it - # connected to, so the relayed proof is rejected. - import unsloth_cli._inference as inference - - inference.ensure_studio_backend_path() - try: - from studio.backend.auth import storage - except Exception as exc: - pytest.skip(f"studio backend not importable: {exc}") - - monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") - monkeypatch.setattr(storage, "_identity_secret_cache", None) - - real_base, stop_real = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) - ) - real_port = int(real_base.rsplit(":", 1)[1]) - # The squatter answers on its own port but returns the proof for the real port. - squatter_base, stop_squatter = _serve_identity( - lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) - ) - try: - assert inference.verify_studio_identity(real_base) is True - assert inference.verify_studio_identity(squatter_base) is False - finally: - stop_real() - stop_squatter() - - -@pytest.mark.parametrize( - "url, loopback", - [ - ("http://127.0.0.1:8888", True), - ("http://localhost:8888", True), - ("http://[::1]:8888", True), - ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 - ("http://0.0.0.0:8888", False), - ("http://10.0.0.5:8888", False), - ("http://studio.evil.example:8888", False), - ("https://studio.example.com", False), - ], -) -def test_is_loopback_url(url, loopback): - assert connect.is_loopback_url(url) is loopback - - -def test_connect_no_studio_errors(fake_studio, monkeypatch): - monkeypatch.setattr(connect, "find_studio_server", lambda: None) - result = CliRunner().invoke(connect.connect_app, ["claude", "--no-launch"]) - assert result.exit_code == 1 - assert "No running Studio server" in result.output - - -def test_connect_explicit_api_key_skips_mint(fake_studio): - result = CliRunner().invoke( - connect.connect_app, - ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], - ) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") - assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) - - -# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── - - -def test_write_openclaw_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["models"]["providers"]["unsloth"] - assert provider["baseUrl"] == f"{BASE}/v1" - assert provider["apiKey"] == "sk-unsloth-abc" - assert provider["api"] == "openai-completions" - assert provider["models"] == [ - {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} - ] - # The default model must be pinned or OpenClaw has nothing active. - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["gateway"]["mode"] == "local" - assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway - if os.name != "nt": # the file holds an API key - assert path.stat().st_mode & 0o777 == 0o600 - - -def test_write_openclaw_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text( - json.dumps( - { - "theme": "dark", - "agents": {"defaults": {"temperature": 0.5}}, - "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, - } - ) - ) - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "dark" - assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - assert config["models"]["mode"] == "replace" # user's mode is left as-is - assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" - assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" - before = path.read_text() - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_write_openclaw_config_corrupt_left_alone(tmp_path, monkeypatch, capsys): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - path.write_text("{not json") - connect.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == "{not json" - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_openclaw_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "openclaw.json" - monkeypatch.setattr(connect, "openclaw_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["openclaw", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "openclaw" in result.output - assert "export" not in result.output # key lives in the config, not the env - config = json.loads(path.read_text()) - assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" - # OpenAI /v1/chat/completions works on either backend — no GGUF gate. - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── - - -def test_write_opencode_config_fresh(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - provider = config["provider"]["unsloth"] - assert provider["npm"] == "@ai-sdk/openai-compatible" - assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} - assert provider["models"] == {MODEL["id"]: {"name": MODEL["id"]}} - assert config["model"] == f"unsloth/{MODEL['id']}" - - -def test_write_opencode_config_preserves_and_idempotent(tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - path.write_text( - json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) - ) - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - config = json.loads(path.read_text()) - assert config["theme"] == "tokyonight" - assert config["provider"]["anthropic"]["name"] == "Anthropic" - assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" - before = path.read_text() - connect.write_opencode_config(BASE, "sk-unsloth-abc", MODEL) - assert path.read_text() == before - - -def test_connect_opencode_no_launch(fake_studio, tmp_path, monkeypatch): - path = tmp_path / "opencode.json" - monkeypatch.setattr(connect, "opencode_config_path", lambda: path) - result = CliRunner().invoke(connect.connect_app, ["opencode", "--no-launch"]) - assert result.exit_code == 0, result.output - assert "opencode" in result.output - config = json.loads(path.read_text()) - assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["model"] == f"unsloth/{MODEL['id']}" - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) - - -# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── - - -@pytest.fixture() -def hermes_config(tmp_path, monkeypatch): - path = tmp_path / "config.yaml" - monkeypatch.setattr(connect, "hermes_config_path", lambda: path) - return path - - -def test_write_hermes_config_fresh(hermes_config): - yaml = pytest.importorskip("yaml") - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - # Hermes only honors the key for a *named* custom provider, so the endpoint - # is registered under providers.* and model.provider points at it. - assert config["model"]["provider"] == "custom:unsloth" - assert config["model"]["default"] == MODEL["id"] - assert config["model"]["api_mode"] == "openai" - provider = config["providers"]["unsloth"] - assert provider["base_url"] == f"{BASE}/v1" - assert provider["api_mode"] == "openai" - assert provider["key_env"] == "UNSLOTH_API_KEY" - # The key is resolved from the launch env, never written to disk. - assert "sk-unsloth" not in hermes_config.read_text() - - -def test_write_hermes_config_preserves_and_idempotent(hermes_config): - yaml = pytest.importorskip("yaml") - hermes_config.write_text( - yaml.safe_dump( - { - "terminal": {"backend": "local"}, - "model": {"temperature": 0.7}, - "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, - } - ) - ) - connect.write_hermes_config(BASE, MODEL) - config = yaml.safe_load(hermes_config.read_text()) - assert config["terminal"] == {"backend": "local"} # unrelated sections kept - assert config["model"]["temperature"] == 0.7 # unrelated model keys kept - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - before = hermes_config.read_text() - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == before - - -def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): - pytest.importorskip("yaml") - original = "- just\n- a\n- list\n" # valid YAML, but not a mapping - hermes_config.write_text(original) - connect.write_hermes_config(BASE, MODEL) - assert hermes_config.read_text() == original # user-managed file left untouched - assert "couldn't parse" in capsys.readouterr().err - - -def test_connect_hermes_no_launch(fake_studio, hermes_config): - yaml = pytest.importorskip("yaml") - result = CliRunner().invoke(connect.connect_app, ["hermes", "--no-launch"]) - assert result.exit_code == 0, result.output - _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") - assert "hermes" in result.output - config = yaml.safe_load(hermes_config.read_text()) - assert config["model"]["provider"] == "custom:unsloth" - assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" - assert config["model"]["default"] == MODEL["id"] - assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py new file mode 100644 index 0000000000..a6a092a17c --- /dev/null +++ b/unsloth_cli/tests/test_start.py @@ -0,0 +1,1848 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for `unsloth start` — config merging and launch env, no network.""" + +from __future__ import annotations + +import json +import os +import shlex +import sys +import urllib.error +from pathlib import Path +from types import SimpleNamespace + +_REPO_ROOT = Path(__file__).resolve().parents[2] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +import pytest +from typer.testing import CliRunner + +import unsloth_cli.commands.start as start + +BASE = "http://127.0.0.1:8888" +MODEL = {"id": "unsloth/gemma-4-26B-A4B-it-GGUF", "context_length": 131072} + + +# --no-launch prints shell setup as POSIX (export/unset) on Unix/WSL and +# PowerShell ($env:/Remove-Item) on native Windows; assert the host's form. +def _assert_env_set(output: str, name: str, value: str) -> None: + needle = f'$env:{name} = "{value}"' if os.name == "nt" else f"export {name}={value}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _assert_env_unset(output: str, name: str) -> None: + needle = f"Remove-Item Env:{name}" if os.name == "nt" else f"unset {name}" + assert needle in output, f"{needle!r} not found in:\n{output}" + + +def _launch_command(output: str) -> list: + # The --no-launch recipe ends with a self-contained one-liner: inline NAME=value + # assignments, then the command. Return just the command argv. + last = [ln for ln in output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + for i, part in enumerate(parts): + name = part.partition("=")[0] + if "=" not in part or not name.replace("_", "").isalnum(): + return parts[i:] + return [] + + +def _fake_claude(monkeypatch, version_output: str) -> None: + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr( + start.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace(stdout = version_output), + ) + + +def test_claude_flags_passed_to_supported_claude(monkeypatch): + _fake_claude(monkeypatch, "2.1.98 (Claude Code)\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_claude_flags_skipped_on_old_claude(monkeypatch): + _fake_claude(monkeypatch, "2.0.14 (Claude Code)\n") + assert start._claude_flags() == [] + + +def test_claude_flags_skipped_on_unparseable_version(monkeypatch): + _fake_claude(monkeypatch, "weird build string\n") + assert start._claude_flags() == [] + + +def test_claude_flags_detected_when_version_not_first_token(monkeypatch): + # The X.Y.Z is pulled from anywhere in the output, so a format change (version not + # the first token) doesn't silently drop the optimization flags. + _fake_claude(monkeypatch, "claude version 2.1.98\n") + assert start._claude_flags() == [ + "--exclude-dynamic-system-prompt-sections", + "--settings", + start._CLAUDE_SETTINGS_OVERLAY, + ] + + +def test_install_agent_prompts_then_installs(monkeypatch): + # TTY + yes: run the documented install command, then re-resolve the now-present binary. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) + ran = [] + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, *a, **k: ran.append(command) or SimpleNamespace(returncode = 0), + ) + # _install_agent only re-resolves after installing (the pre-install check is the + # caller's job), so `which` reports the now-present binary. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + executable = start._install_agent("codex", "npm install -g @openai/codex") + assert executable == "/usr/local/bin/codex" + assert ran == [["/bin/sh", "-c", "npm install -g @openai/codex"]] + + +def test_install_agent_declined_returns_none(monkeypatch): + # TTY + no: never runs anything; caller falls back to the print-hint failure. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install when declined") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def test_install_agent_non_interactive_returns_none(monkeypatch): + # No TTY (piped stdin): cannot prompt, so don't install; return None silently. + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: False)) + monkeypatch.setattr( + start.subprocess, "run", lambda *a, **k: pytest.fail("should not install without a TTY") + ) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + + +def _parse_toml(text: str) -> dict: + tomllib = pytest.importorskip("tomllib") + return tomllib.loads(text) + + +def test_merge_codex_config_fresh(): + merged = start._merge_codex_config("", BASE) + parsed = _parse_toml(merged) + assert parsed["oss_provider"] == "unsloth_api" + provider = parsed["model_providers"]["unsloth_api"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["wire_api"] == "responses" + assert provider["requires_openai_auth"] is False + + +def test_merge_codex_config_replaces_stale_block(): + existing = ( + 'model = "gpt-5"\n' + "\n" + "[model_providers.unsloth_api]\n" + 'base_url = "http://old-host:9999/v1"\n' + 'wire_api = "chat"\n' + "\n" + "[model_providers.unsloth_api.http_headers]\n" + 'x-old = "1"\n' + "\n" + "[model_providers.ollama]\n" + 'base_url = "http://localhost:11434/v1"\n' + ) + merged = start._merge_codex_config(existing, BASE) + parsed = _parse_toml(merged) + assert parsed["model"] == "gpt-5" + assert parsed["model_providers"]["unsloth_api"]["base_url"] == f"{BASE}/v1" + assert parsed["model_providers"]["unsloth_api"]["wire_api"] == "responses" + assert "http_headers" not in parsed["model_providers"]["unsloth_api"] + assert parsed["model_providers"]["ollama"]["base_url"] == "http://localhost:11434/v1" + assert start._merge_codex_config(merged, BASE) == merged + + +def test_merge_codex_config_keeps_user_oss_provider(): + merged = start._merge_codex_config('oss_provider = "ollama"\n', BASE) + assert _parse_toml(merged)["oss_provider"] == "ollama" + + +def test_write_codex_config_profile(tmp_path): + start.write_codex_config(BASE, MODEL, tmp_path) + profile = _parse_toml((tmp_path / "unsloth_api.config.toml").read_text()) + assert profile["oss_provider"] == "unsloth_api" + assert profile["model_provider"] == "unsloth_api" + assert profile["model"] == MODEL["id"] + assert profile["model_context_window"] == 131072 + config = _parse_toml((tmp_path / "config.toml").read_text()) + assert config["model_providers"]["unsloth_api"]["env_key"] == "UNSLOTH_STUDIO_AUTH_TOKEN" + + +@pytest.fixture() +def fake_studio(tmp_path, monkeypatch): + calls = [] + state = {"models": [MODEL]} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return {"object": "list", "data": state["models"]} + if url.endswith("/api/inference/status"): + return {"is_gguf": True, "model_identifier": state["models"][0]["id"]} + if url.endswith("/api/auth/api-keys"): + return {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/inference/load"): + state["models"] = [{"id": payload["model_path"], "context_length": 4096}] + return {} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "find_studio_server", lambda: BASE) + # Identity handshake has its own tests; trust the loopback server here. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: True) + # _studio_token / api-keys are faked so the mint flow stays offline. + monkeypatch.setattr(start, "_studio_token", lambda: "jwt-token") + monkeypatch.setattr(start, "_http_json", http_json) + monkeypatch.setattr(start, "_key_cache_path", lambda: tmp_path / "agent_api_key.json") + # --no-launch session configs land under tmp instead of the real Unsloth dir. + monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents") + # No `claude` on PATH, so _claude_flags never probes the real binary. + monkeypatch.setattr(start.shutil, "which", lambda _: None) + monkeypatch.delenv("UNSLOTH_API_KEY", raising = False) + return calls + + +def test_connect_claude_no_launch(fake_studio): + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_unset(result.output, "ANTHROPIC_API_KEY") + _assert_env_unset(result.output, "CLAUDE_CODE_OAUTH_TOKEN") + _assert_env_set(result.output, "ANTHROPIC_BASE_URL", BASE) + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC", "1") + _assert_env_set(result.output, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS", "1") + # Suppress the full-screen TUI redraw so a bursty local server doesn't flicker. + _assert_env_set(result.output, "CLAUDE_CODE_NO_FLICKER", "1") + # Attribution header is suppressed for the session via env + --settings, never + # by writing the user's ~/.claude/settings.json. + _assert_env_set(result.output, "CLAUDE_CODE_ATTRIBUTION_HEADER", "0") + # Auto-compact window is sized to the loaded model's real context length so the + # session compacts before it overflows the local server's (much smaller) window, + # and compaction is forced at 90% of it for headroom. + _assert_env_set(result.output, "CLAUDE_CODE_AUTO_COMPACT_WINDOW", str(MODEL["context_length"])) + _assert_env_set(result.output, "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE", "90") + assert f"claude --model {MODEL['id']} --exclude-dynamic-system-prompt-sections" in result.output + # Overlay is passed inline (session-only), not a path into the user's ~/.claude. + assert "--settings" in result.output + assert ".claude/settings.json" not in result.output + + +def test_connect_claude_compact_window_omitted_without_context(fake_studio, monkeypatch): + # A model that doesn't report a context length -> leave Claude's default window + # rather than guessing one. + monkeypatch.setattr(start, "_resolve_model", lambda *a, **k: {"id": "local-model"}) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" not in result.output + assert "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE" not in result.output + + +def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == ["/usr/local/bin/claude", "--model", MODEL["id"]] + assert "ANTHROPIC_API_KEY" not in captured["env"] + assert "CLAUDE_CODE_OAUTH_TOKEN" not in captured["env"] + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + assert captured["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_windows_shim_from_wsl_bridges_env(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") + monkeypatch.setenv("CLAUDE_CODE_OAUTH_TOKEN", "oauth-stale") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + monkeypatch.setattr(start, "_claude_flags", lambda: []) + + def run(command, env): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["claude"]) + + assert result.exit_code == 0, result.output + assert captured["command"] == [ + "/mnt/c/Users/samle/AppData/Roaming/npm/claude", + "--model", + MODEL["id"], + ] + assert captured["env"]["ANTHROPIC_API_KEY"] == "" + assert captured["env"]["CLAUDE_CODE_OAUTH_TOKEN"] == "" + assert captured["env"]["ANTHROPIC_AUTH_TOKEN"] == "sk-unsloth-feedfacefeedface" + assert captured["env"]["ANTHROPIC_BASE_URL"] == BASE + assert captured["env"]["ANTHROPIC_MODEL"] == MODEL["id"] + for name in ( + "ANTHROPIC_AUTH_TOKEN", + "ANTHROPIC_BASE_URL", + "ANTHROPIC_MODEL", + "ANTHROPIC_API_KEY", + "CLAUDE_CODE_OAUTH_TOKEN", + ): + assert name in captured["env"]["WSLENV"].split(":") + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario (calling a Windows agent .exe from inside WSL); " + "os.name is 'posix' under WSL, so this path can't run on a native Windows runner.", +) +def test_connect_claude_no_launch_windows_shim_from_wsl_prints_wslenv(fake_studio, monkeypatch): + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr( + start.shutil, "which", lambda _: "/mnt/c/Users/samle/AppData/Roaming/npm/claude" + ) + + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + + assert result.exit_code == 0, result.output + assert "export ANTHROPIC_API_KEY=" in result.output + assert "export CLAUDE_CODE_OAUTH_TOKEN=" in result.output + assert "export WSLENV=" in result.output + assert "ANTHROPIC_AUTH_TOKEN" in result.output + assert "CLAUDE_CODE_OAUTH_TOKEN" in result.output + + +def test_connect_codex_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_STUDIO_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + assert "codex --oss --profile unsloth_api" in result.output + # Config lands in the session-scoped CODEX_HOME, not the user's ~/.codex. + home = tmp_path / "agents" / "codex" + _assert_env_set(result.output, "CODEX_HOME", str(home)) + assert (home / "config.toml").exists() + assert (home / "unsloth_api.config.toml").exists() + + +def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): + # Launch mode writes config to a throwaway temp CODEX_HOME and removes it after + # the agent exits; the user's real ~/.codex is never the target. + captured = {} + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + + def run(command, env): + captured["home"] = env["CODEX_HOME"] + captured["config_present"] = (Path(env["CODEX_HOME"]) / "config.toml").exists() + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["codex"]) + assert result.exit_code == 0, result.output + home = Path(captured["home"]) + assert captured["config_present"] # config existed while codex ran + assert "unsloth-codex-" in home.name # an ephemeral temp dir, not ~/.codex + assert not home.exists() # cleaned up after the agent exits + + +@pytest.mark.skipif( + os.name == "nt", + reason = "the #6547 CI parser is bash-only; on Windows --no-launch prints PowerShell", +) +def test_no_launch_output_is_parseable(fake_studio): + # Mirror the #6547 CI parser: status lines, then `export`/`unset`, then exactly + # one launch command on the last line (now an inline-env one-liner, so the parser + # matches by substring rather than prefix). + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + lines = [ln for ln in result.output.splitlines() if ln.strip()] + skip = ("export ", "unset ", "Studio ", "Updated ", "Disabled ", "Warning", "Loading") + body = [ln for ln in lines if not ln.startswith(skip)] + assert "codex --oss --profile unsloth_api" in body[-1] + assert any(ln.startswith("export CODEX_HOME=") for ln in lines) + + +def test_no_launch_last_line_is_self_contained(fake_studio, tmp_path): + # People copy just the last line. A bare `codex` there would run against the user's + # real ~/.codex (e.g. a pre-existing damaged state DB) with zero isolation, so the + # last line must inline every session env var ahead of the command. + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + parts = shlex.split(last) + assignments = {} + command = [] + for i, part in enumerate(parts): + if "=" not in part: + command = parts[i:] + break + name, _, value = part.partition("=") + assignments[name] = value + assert command and command[0] == "codex" + assert assignments["CODEX_HOME"] == str(tmp_path / "agents" / "codex") + assert assignments["UNSLOTH_STUDIO_AUTH_TOKEN"].startswith("sk-unsloth-") + + +def test_no_launch_claude_last_line_blanks_conflicting_auth(fake_studio): + # The unset vars must be neutralized inline too, or a partial copy would send the + # user's own ANTHROPIC_API_KEY to the Studio base. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + last = [ln for ln in result.output.splitlines() if ln.strip()][-1] + assert "ANTHROPIC_API_KEY= " in last + assert "CLAUDE_CODE_OAUTH_TOKEN= " in last + assert "ANTHROPIC_AUTH_TOKEN=" in last # the real key still applied after the blanks + + +def test_opencode_inline_config_beats_project_config(fake_studio): + # A project's opencode.json outranks OPENCODE_CONFIG, so the model pin (and --yolo + # permissions) ride in OPENCODE_CONFIG_CONTENT, which outranks project config. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--yolo"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline["model"] == f"unsloth/{MODEL['id']}" + assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + assert "sk-unsloth" not in content_line # key stays in the private file + + +def test_opencode_inline_config_omits_permissions_without_yolo(fake_studio): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + content_line = next( + ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") + ) + inline = json.loads( + shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] + ) + assert inline == {"model": f"unsloth/{MODEL['id']}"} + + +def test_https_loopback_never_auto_serves(fake_studio, monkeypatch): + # `unsloth run` serves plain HTTP; auto-serving behind an https:// target would poll + # the wrong scheme until the startup timeout. Keep the plain "no server" error. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "https://127.0.0.1:8443") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_connect_alias_still_works(fake_studio): + # `unsloth connect` remains a compat alias for `unsloth start`. + from unsloth_cli import app + + result = CliRunner().invoke(app, ["connect", "claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_key_minted_once_then_cached(fake_studio, tmp_path): + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + # First run mints; second reuses the minted key cached for this server. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert len(mints) == 1 + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + + +def test_connect_explicit_key_remembered_for_keyless_runs(fake_studio, tmp_path): + CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + # Reused, not re-minted (a mint would return the feedface stand-in). + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + # An explicit key is remembered as "saved" so it replays without the handshake. + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_connect_skips_cached_keys_the_server_rejects(fake_studio, tmp_path, monkeypatch): + cache = tmp_path / "agent_api_key.json" + cache.write_text( + json.dumps( + {"servers": {BASE: {"minted": ["sk-unsloth-stale", "sk-unsloth-feedfacefeedface"]}}} + ) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-stale": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + # The working key moves to the front so the next run tries it first. + cached = json.loads(cache.read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface", "sk-unsloth-stale"] + + +def test_connect_saved_key_server_outage_surfaces_not_reminted(fake_studio, tmp_path, monkeypatch): + # A 5xx/timeout while checking a saved key is a server outage, not a rejected key: + # surface it instead of discarding the key and minting a new one against a sick server. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-saved"]}}})) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-saved": + raise urllib.error.HTTPError(url, 503, "Service Unavailable", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code != 0, result.output + # The outage did not cause a fresh key to be minted. + mints = [c for c in fake_studio if c[1].endswith("/api/auth/api-keys")] + assert mints == [] + + +def test_connect_legacy_unscoped_cache_not_replayed(fake_studio, tmp_path): + # Legacy unscoped caches have no server binding (could leak across servers), + # so they're ignored: a fresh key is minted and stored scoped to this server. + (tmp_path / "agent_api_key.json").write_text(json.dumps({"key": "sk-unsloth-oldformat"})) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-feedfacefeedface") + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["minted"] == ["sk-unsloth-feedfacefeedface"] + assert "key" not in cached # legacy field collapsed away + + +def test_connect_model_flag_loads_on_server(fake_studio): + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Qwen3.5-35B-A3B"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": "unsloth/Qwen3.5-35B-A3B"}) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", "unsloth/Qwen3.5-35B-A3B") + + +def test_connect_model_flag_forwards_load_options(fake_studio): + # The model-load knobs mirrored from `unsloth run` reach /api/inference/load. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF", + "--gguf-variant", + "UD-Q4_K_XL", + "--context-length", + "8192", + "--no-load-in-4bit", + "--tensor-parallel", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + { + "model_path": "unsloth/Qwen3-4B-GGUF", + "gguf_variant": "UD-Q4_K_XL", + "max_seq_length": 8192, + "load_in_4bit": False, + "tensor_parallel": True, + }, + ) + ] + + +def test_connect_model_flag_matches_canonical_id(fake_studio, monkeypatch): + # Studio registers a loaded model under a canonical id (resolved identifier + # / casing) that can differ from the path we passed. The agent must connect + # to that model, not silently fall through to the first loaded one. + requested = "Unsloth/Qwen3.5-35B-A3B" + canonical = "unsloth/Qwen3.5-35B-A3B" + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {"model": canonical, "display_name": canonical} + if url.endswith("/v1/models"): + # Decoy sorts first, so models[0] is the wrong pick on the old code. + return {"object": "list", "data": [MODEL, {"id": canonical, "context_length": 4096}]} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", requested]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_MODEL", canonical) + + +@pytest.mark.parametrize( + "model, expected", + [ + ("unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", ("unsloth/Qwen3-1.7B-GGUF", "UD-Q4_K_XL")), + ("unsloth/gemma-4-E2B-it-GGUF:Q8_0", ("unsloth/gemma-4-E2B-it-GGUF", "Q8_0")), + ("unsloth/Qwen3-1.7B-GGUF", ("unsloth/Qwen3-1.7B-GGUF", None)), # no suffix + ("/models/local.gguf", ("/models/local.gguf", None)), # absolute path + ("./rel.gguf", ("./rel.gguf", None)), # relative path + ("C:\\models\\x.gguf", ("C:\\models\\x.gguf", None)), # Windows drive + ("repo:with/slash", ("repo:with/slash", None)), # slash in variant -> not a variant + ("", ("", None)), + ], +) +def test_split_repo_variant(model, expected): + assert start._split_repo_variant(model) == expected + + +def test_connect_model_bare_id_matches_loaded_without_reload(fake_studio): + # A bare `--model ` (no load knobs) attaches to the already-loaded model + # without touching /api/inference/load, so it can never evict another session. + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch", "--model", MODEL["id"]]) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_model_variant_suffix_defers_to_server_dedup(fake_studio): + # `--model repo:QUANT` splits into a VALID load payload (bare repo + gguf_variant), + # never the `:`-suffixed repo id Studio rejects. The variant knob defers to + # /api/inference/load, whose already-loaded dedup answers without reloading when the + # active variant+settings match -- so a second session running the same command + # attaches without evicting the first, while a genuinely different quant reloads. + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", MODEL["id"] + ":UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": MODEL["id"], "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + _assert_env_set(result.output, "ANTHROPIC_MODEL", MODEL["id"]) + + +def test_connect_load_knobs_reach_server_even_when_id_loaded(fake_studio): + # /v1/models can't reveal the active quant, so an id match alone would silently keep + # the wrong variant loaded. Explicit knobs must always consult the load endpoint. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", MODEL["id"], "--gguf-variant", "Q8_0"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ("POST", f"{BASE}/api/inference/load", {"model_path": MODEL["id"], "gguf_variant": "Q8_0"}) + ] + + +def test_connect_model_variant_suffix_loads_split_repo(fake_studio): + # When the model is not already loaded, the `:QUANT` suffix becomes the gguf_variant + # and the load uses the bare (valid) repo id, mirroring `unsloth run repo --gguf-variant`. + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--model", "unsloth/Qwen3-4B-GGUF:UD-Q4_K_XL"], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_explicit_gguf_variant_wins_over_suffix(fake_studio): + # An explicit --gguf-variant takes precedence; the suffix is still stripped so the + # repo id stays valid. + result = CliRunner().invoke( + start.start_app, + [ + "claude", + "--no-launch", + "--model", + "unsloth/Qwen3-4B-GGUF:Q8_0", + "--gguf-variant", + "UD-Q4_K_XL", + ], + ) + assert result.exit_code == 0, result.output + loads = [c for c in fake_studio if c[1].endswith("/api/inference/load")] + assert loads == [ + ( + "POST", + f"{BASE}/api/inference/load", + {"model_path": "unsloth/Qwen3-4B-GGUF", "gguf_variant": "UD-Q4_K_XL"}, + ) + ] + + +def test_connect_no_model_loaded_errors(fake_studio, monkeypatch): + monkeypatch.setattr( + start, + "_http_json", + lambda method, url, token, payload = None, timeout = 30, error = None: ( + {"key": "sk-unsloth-feedfacefeedface"} + if url.endswith("/api/auth/api-keys") + else {"object": "list", "data": []} + ), + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No model is loaded" in result.output + + +def test_connect_requested_model_not_loaded_fails(fake_studio, monkeypatch): + # Studio never surfaces the requested model; fail loudly rather than + # silently connecting to whatever else happens to be loaded. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/load"): + return {} + if url.endswith("/v1/models"): + return {"object": "list", "data": [MODEL]} # decoy; request never appears + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["claude", "--no-launch", "--model", "unsloth/Missing-7B"] + ) + assert result.exit_code == 1 + assert "unsloth/Missing-7B" in result.output + + +def test_connect_codex_rejects_non_gguf_model(fake_studio, monkeypatch): + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "unsloth/Qwen3-0.6B"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke(start.start_app, ["codex", "--no-launch"]) + assert result.exit_code == 1 + assert "GGUF" in result.output + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_keyless_refuses_to_send_credential(fake_studio, monkeypatch): + # A server known only by URL + health check is unverified: keyless connect + # must refuse and make no request at all. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.evil.example:8888") + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 1 + assert "Settings → API" in result.output + assert "--api-key" in result.output + assert fake_studio == [] # no HTTP request of any kind (no mint, no /v1/models) + + +def test_connect_nonloopback_explicit_key_is_allowed(fake_studio, monkeypatch): + # User named both server and key, so it's their choice; only auto-send is blocked. + monkeypatch.setattr(start, "find_studio_server", lambda: "http://studio.example:8888") + result = CliRunner().invoke( + start.start_app, + ["opencode", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + + +def test_connect_nonloopback_replays_saved_key(fake_studio, tmp_path, monkeypatch): + # A key saved for a remote (non-loopback) Studio is replayed on keyless runs; + # auto-minting stays blocked for non-loopback. + remote = "http://studio.example:8888" + monkeypatch.setattr(start, "find_studio_server", lambda: remote) + (tmp_path / "agent_api_key.json").write_text( + json.dumps({"servers": {remote: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}}) + ) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_studio_server_errors_on_explicit_remote(monkeypatch): + # A user who pointed UNSLOTH_STUDIO_URL at a remote Studio should get an + # error, not a silent local model load (which they did not ask for). + import typer + + import unsloth_cli._inference as inference + + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://studio.example:8888") + monkeypatch.setattr( + inference, "find_studio_server", lambda *a, **k: "http://studio.example:8888" + ) + with pytest.raises(typer.Exit): + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + + +def test_connect_studio_server_falls_back_locally_on_default_discovery(monkeypatch): + # Opportunistic local discovery (no UNSLOTH_STUDIO_URL): if the loopback + # server can't be verified, fall back to a local load rather than erroring. + import unsloth_cli._inference as inference + + monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False) + monkeypatch.setattr(inference, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888") + monkeypatch.setattr(inference, "verify_studio_identity", lambda *a, **k: False) + assert ( + inference.connect_studio_server("m", hf_token = None, max_seq_length = 4096, load_in_4bit = False) + is None + ) + + +def test_connect_unverified_loopback_without_cached_key_refuses_to_mint( + fake_studio, tmp_path, monkeypatch +): + # With no saved key, the next step would auto-mint; an unverified loopback + # server (port squatter) must be refused, with nothing sent. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # never minted + + +def test_connect_replays_saved_key_without_identity_check(fake_studio, tmp_path, monkeypatch): + # A "saved" key (e.g. for an SSH-tunnelled Studio the handshake can't match) + # replays on keyless runs without the handshake, scoped to its own base. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"saved": ["sk-unsloth-deadbeefdeadbeef"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) # reused, not minted + + +def test_connect_minted_cache_requires_identity_check(fake_studio, tmp_path, monkeypatch): + # A "minted" key is NOT replayed to an unverified loopback server: minting and + # minted-key replay both sit behind the handshake, so a squatter can't grab it. + cache = tmp_path / "agent_api_key.json" + cache.write_text(json.dumps({"servers": {BASE: {"minted": ["sk-unsloth-feedfacefeedface"]}}})) + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "--api-key" in result.output + assert not any(c[1].endswith("/v1/models") for c in fake_studio) # minted key never sent + + +def test_connect_explicit_key_skips_identity_check(fake_studio, monkeypatch): + # An explicit key is the user's deliberate choice, so it does not require + # the automatic identity handshake. + monkeypatch.setattr(start, "verify_studio_identity", lambda base: False) + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + + +def _serve_identity(proof_for): + """Start a localhost HTTP server answering /api/auth/identity with + proof_for(nonce_bytes). Returns (base_url, shutdown).""" + import base64 + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + from urllib.parse import parse_qs, urlparse + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + parsed = urlparse(self.path) + if parsed.path != "/api/auth/identity": + self.send_response(404) + self.end_headers() + return + nonce = base64.urlsafe_b64decode(parse_qs(parsed.query)["nonce"][0]) + host, port = self.server.server_address[0], self.server.server_address[1] + body = json.dumps({"proof": proof_for(nonce, host, port)}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(body) + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_end_to_end(tmp_path, monkeypatch): + # Real crypto end to end: verify_studio_identity reads the install secret from + # an isolated DB; a "good" server proves the same secret, a spoofing one can't. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: # backend not importable here (e.g. missing deps) + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + good = lambda nonce, host, port: storage.compute_identity_proof( + nonce, host, port + ) # real secret + bad = lambda nonce, host, port: "00" * 32 # spoofer without the secret + base_ok, stop_ok = _serve_identity(good) + base_bad, stop_bad = _serve_identity(bad) + try: + assert inference.verify_studio_identity(base_ok) is True + assert inference.verify_studio_identity(base_bad) is False + finally: + stop_ok() + stop_bad() + + +def _serve_redirect(target): + """Start a localhost server that 302-redirects every GET to target+path.""" + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + class Handler(BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(302) + self.send_header("Location", target + self.path) + self.end_headers() + + def log_message(self, *a): + pass + + server = HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target = server.serve_forever, daemon = True).start() + base = f"http://127.0.0.1:{server.server_address[1]}" + return base, server.shutdown + + +def test_verify_studio_identity_rejects_redirect(tmp_path, monkeypatch): + # A squatter could 302 /api/auth/identity to the real Studio and relay its + # proof; redirects must be refused so the squatter's base isn't accepted. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + squatter_base, stop_squatter = _serve_redirect(real_base) + try: + assert inference.verify_studio_identity(real_base) is True # direct: ok + assert inference.verify_studio_identity(squatter_base) is False # relayed: refused + finally: + stop_real() + stop_squatter() + + +def test_verify_studio_identity_rejects_relayed_proof(tmp_path, monkeypatch): + # A squatter that proxies the nonce to the real Studio on another port gets a + # proof bound to *that* port; the client expects one bound to the port it + # connected to, so the relayed proof is rejected. + import unsloth_cli._inference as inference + + inference.ensure_studio_backend_path() + try: + from studio.backend.auth import storage + except Exception as exc: + pytest.skip(f"studio backend not importable: {exc}") + + monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db") + monkeypatch.setattr(storage, "_identity_secret_cache", None) + + real_base, stop_real = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, port) + ) + real_port = int(real_base.rsplit(":", 1)[1]) + # The squatter answers on its own port but returns the proof for the real port. + squatter_base, stop_squatter = _serve_identity( + lambda nonce, host, port: storage.compute_identity_proof(nonce, host, real_port) + ) + try: + assert inference.verify_studio_identity(real_base) is True + assert inference.verify_studio_identity(squatter_base) is False + finally: + stop_real() + stop_squatter() + + +@pytest.mark.parametrize( + "url, loopback", + [ + ("http://127.0.0.1:8888", True), + ("http://localhost:8888", True), + ("http://[::1]:8888", True), + ("http://127.0.0.5:9001", True), # SSH tunnels can land anywhere in 127/8 + ("http://0.0.0.0:8888", False), + ("http://10.0.0.5:8888", False), + ("http://studio.evil.example:8888", False), + ("https://studio.example.com", False), + ], +) +def test_is_loopback_url(url, loopback): + assert start.is_loopback_url(url) is loopback + + +def test_connect_no_studio_errors(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude", "--no-launch"]) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + + +@pytest.fixture(autouse = True) +def _reset_auto_served(): + # Never let a test leave a fake server in the module slot (an atexit backstop would + # otherwise try to signal it at interpreter shutdown). + yield + start._auto_served_server = None + + +def test_start_studio_server_builds_command_and_waits(monkeypatch): + captured = {} + + class FakePopen: + def __init__(self, command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + self.pid = 4321 + + def poll(self): + return None + + monkeypatch.setattr(start.subprocess, "Popen", FakePopen) + monkeypatch.setattr(start, "_studio_healthy", lambda base, timeout = 3.0: True) + monkeypatch.setattr(start, "_log_tail", lambda path, lines = 20: "API Key: sk-unsloth-abc123") + monkeypatch.setattr(start.time, "sleep", lambda _s: None) + + server = start._start_studio_server( + "http://127.0.0.1:8888", + "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL", + start.LoadOptions( + gguf_variant = "UD-Q4_K_XL", max_seq_length = 8192, load_in_4bit = True, tensor_parallel = True + ), + ) + cmd = captured["command"] + assert cmd[1] == "run" + assert "--disable-tools" in cmd and "--no-cloudflare" in cmd + assert cmd[cmd.index("--model") + 1] == "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL" + assert cmd[cmd.index("--gguf-variant") + 1] == "UD-Q4_K_XL" + assert cmd[cmd.index("--context-length") + 1] == "8192" + assert "--tensor-parallel" in cmd + assert cmd[cmd.index("-p") + 1] == "8888" + assert start.LoadOptions().load_in_4bit is True and "--no-load-in-4bit" not in cmd + assert captured["kwargs"].get("start_new_session") is True # own process group + assert server.pid == 4321 + + +def test_auto_serves_when_no_server_then_tears_down(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model, load = load) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF:UD-Q4_K_XL"] + ) + assert result.exit_code == 0, result.output + # The `:QUANT` suffix is split off into the gguf_variant so `unsloth run` gets a valid + # repo id plus `--gguf-variant`, mirroring how `unsloth run` accepts either form. + assert started["model"] == "unsloth/Qwen3-1.7B-GGUF" + assert started["load"].gguf_variant == "UD-Q4_K_XL" + assert started["base"] == BASE + # Torn down after the agent session ended. + assert started.get("down") is fake + + +def test_codex_preflight_failure_tears_down_auto_served(fake_studio, monkeypatch): + # The Codex GGUF preflight runs after _connect may have auto-started a server but + # before _run's teardown finally, so a preflight rejection must not leave the server + # holding the port/GPU (waiting on the atexit backstop). + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started.update(base = base, model = model) + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr( + start, "_shutdown_server", lambda server: started.__setitem__("down", server) + ) + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/api/inference/status"): + return {"is_gguf": False, "model_identifier": "transformers-model"} + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + result = CliRunner().invoke( + start.start_app, ["codex", "--model", "unsloth/Qwen3-1.7B", "--launch"] + ) + assert result.exit_code != 0, result.output + assert "GGUF" in result.output + # Torn down at the point the preflight rejected the model, not only via atexit. + assert started.get("down") is fake + + +def test_no_serve_preserves_error(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-serve"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_launch_never_serves(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {"called": False} + monkeypatch.setattr( + start, "_start_studio_server", lambda *a, **k: started.__setitem__("called", True) + ) + result = CliRunner().invoke( + start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF", "--no-launch"] + ) + assert result.exit_code == 1 + assert "No running Studio server" in result.output + assert started["called"] is False + + +def test_no_server_no_model_hints_model_flag(fake_studio, monkeypatch): + monkeypatch.setattr(start, "find_studio_server", lambda: None) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 1 + assert "--model" in result.output + + +@pytest.mark.parametrize( + "base, expected", + [ + ("http://127.0.0.1", "http://127.0.0.1:8888"), # portless -> unsloth run's :8888 + ("http://127.0.0.1:8888", "http://127.0.0.1:8888"), # explicit port kept + ("http://127.0.0.1:9000", "http://127.0.0.1:9000"), + ("http://localhost", "http://localhost:8888"), + ("http://[::1]", "http://[::1]:8888"), # IPv6 literal stays bracketed + ("http://[::1]:8888", "http://[::1]:8888"), + # Paths are stripped: unsloth run serves at the root, so /studio would make the + # health poll hit /studio/api/health (404) until the startup timeout. + ("http://127.0.0.1:8888/studio", "http://127.0.0.1:8888"), + ("http://127.0.0.1/studio", "http://127.0.0.1:8888"), + ], +) +def test_effective_base(base, expected): + assert start._effective_base(base) == expected + + +def test_auto_serve_normalizes_portless_url(fake_studio, monkeypatch): + # A portless UNSLOTH_STUDIO_URL must launch AND poll :8888 (what unsloth run binds), + # not port 80, or readiness never matches and we hit the startup timeout. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://127.0.0.1") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + started = {} + fake = SimpleNamespace(pid = 999, poll = lambda: None) + + def fake_start(base, model, load): + started["base"] = base + start._auto_served_server = fake + return fake + + monkeypatch.setattr(start, "_start_studio_server", fake_start) + monkeypatch.setattr(start, "_shutdown_server", lambda server: None) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + + result = CliRunner().invoke(start.start_app, ["claude", "--model", "unsloth/Qwen3-1.7B-GGUF"]) + assert result.exit_code == 0, result.output + assert started["base"] == "http://127.0.0.1:8888" + + +def test_connect_explicit_api_key_skips_mint(fake_studio): + result = CliRunner().invoke( + start.start_app, + ["claude", "--no-launch", "--api-key", "sk-unsloth-deadbeefdeadbeef"], + ) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "ANTHROPIC_AUTH_TOKEN", "sk-unsloth-deadbeefdeadbeef") + assert not any(c[1].endswith("/api/auth/api-keys") for c in fake_studio) + + +# ── OpenClaw (Anthropic /v1/messages) ──────────────────────────────── + + +def test_write_openclaw_config_fresh(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["models"]["providers"]["unsloth"] + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + assert provider["api"] == "openai-completions" + assert provider["models"] == [ + {"id": MODEL["id"], "name": MODEL["id"], "contextWindow": MODEL["context_length"]} + ] + # The default model must be pinned or OpenClaw has nothing active. + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["gateway"]["mode"] == "local" + assert config["gateway"]["auth"]["mode"] == "none" # unauth loopback gateway + if os.name != "nt": # the file holds an API key + assert path.stat().st_mode & 0o777 == 0o600 + + +def test_write_openclaw_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "openclaw.json" + path.write_text( + json.dumps( + { + "theme": "dark", + "agents": {"defaults": {"temperature": 0.5}}, + "models": {"mode": "replace", "providers": {"openrouter": {"baseUrl": "x"}}}, + } + ) + ) + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "dark" + assert config["agents"]["defaults"]["temperature"] == 0.5 # other agent defaults kept + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert config["models"]["mode"] == "replace" # user's mode is left as-is + assert config["models"]["providers"]["openrouter"]["baseUrl"] == "x" + assert config["models"]["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_write_openclaw_config_corrupt_left_alone(tmp_path, capsys): + path = tmp_path / "openclaw.json" + path.write_text("{not json") + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == "{not json" + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_openclaw_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "openclaw" in result.output + config_path = tmp_path / "agents" / "openclaw" / "openclaw.json" + # Config + state are scoped to the session dir, not the user's ~/.openclaw. + _assert_env_set(result.output, "OPENCLAW_CONFIG_PATH", str(config_path)) + _assert_env_set(result.output, "OPENCLAW_STATE_DIR", str(tmp_path / "agents" / "openclaw")) + config = json.loads(config_path.read_text()) + assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + # OpenAI /v1/chat/completions works on either backend — no GGUF gate. + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── + + +def test_write_opencode_config_fresh(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["provider"]["unsloth"] + assert provider["npm"] == "@ai-sdk/openai-compatible" + assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} + # Context limit must be declared, or OpenCode treats it as 0 and disables compaction. + assert provider["models"] == { + MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}} + } + assert config["model"] == f"unsloth/{MODEL['id']}" + # Compaction buffer scaled to ~10% of the window (compact near 90%). + assert config["compaction"] == {"auto": True, "reserved": 131072 // 10} + + +def test_write_opencode_config_preserves_and_idempotent(tmp_path): + path = tmp_path / "opencode.json" + path.write_text( + json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) + ) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["theme"] == "tokyonight" + assert config["provider"]["anthropic"]["name"] == "Anthropic" + assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" + before = path.read_text() + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_opencode_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + assert "opencode" in result.output + config_path = tmp_path / "agents" / "opencode" / "opencode.json" + # OPENCODE_CONFIG overlay points at the session file, not the user's global config. + _assert_env_set(result.output, "OPENCODE_CONFIG", str(config_path)) + config = json.loads(config_path.read_text()) + assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["model"] == f"unsloth/{MODEL['id']}" + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Hermes (OpenAI /v1/chat/completions, key via env) ──────────────── + + +@pytest.fixture() +def hermes_config(tmp_path): + return tmp_path / "config.yaml" + + +def test_write_hermes_config_fresh(hermes_config): + yaml = pytest.importorskip("yaml") + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes only honors the key for a *named* custom provider, so the endpoint + # is registered under providers.* and model.provider points at it. + assert config["model"]["provider"] == "custom:unsloth" + assert config["model"]["default"] == MODEL["id"] + assert config["model"]["api_mode"] == "openai" + # Pin the real context window (top-level override) and compact at 90% of it. + assert config["model"]["context_length"] == MODEL["context_length"] + assert config["compression"] == {"enabled": True, "threshold": 0.9} + # Windows at or above Hermes' floor need no auxiliary compression override. + assert "auxiliary" not in config + provider = config["providers"]["unsloth"] + assert provider["base_url"] == f"{BASE}/v1" + assert provider["api_mode"] == "openai" + assert provider["key_env"] == "UNSLOTH_API_KEY" + # The key is resolved from the launch env, never written to disk. + assert "sk-unsloth" not in hermes_config.read_text() + + +def test_write_hermes_config_small_window_claims_floor(hermes_config): + yaml = pytest.importorskip("yaml") + small = {"id": "unsloth/Qwen3-1.7B-GGUF", "context_length": 40960} + start.write_hermes_config(BASE, small, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + # Hermes refuses to initialize below its 64,000-token floor, so the recipe + # claims the floor and scales the compaction threshold so it still fires at + # 90% of the REAL window: 0.9 * 40960 / 65536. + assert config["model"]["context_length"] == 65536 + assert config["compression"] == {"enabled": True, "threshold": 0.5625} + # The same floor check runs against the compression model mid-session. + assert config["auxiliary"]["compression"]["context_length"] == 65536 + + +def test_write_hermes_config_preserves_and_idempotent(hermes_config): + yaml = pytest.importorskip("yaml") + hermes_config.write_text( + yaml.safe_dump( + { + "terminal": {"backend": "local"}, + "model": {"temperature": 0.7}, + "providers": {"openrouter": {"base_url": "https://openrouter.ai/api/v1"}}, + } + ) + ) + start.write_hermes_config(BASE, MODEL, hermes_config) + config = yaml.safe_load(hermes_config.read_text()) + assert config["terminal"] == {"backend": "local"} # unrelated sections kept + assert config["model"]["temperature"] == 0.7 # unrelated model keys kept + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["openrouter"]["base_url"] == "https://openrouter.ai/api/v1" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + before = hermes_config.read_text() + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == before + + +def test_write_hermes_config_preserves_non_mapping_file(hermes_config, capsys): + pytest.importorskip("yaml") + original = "- just\n- a\n- list\n" # valid YAML, but not a mapping + hermes_config.write_text(original) + start.write_hermes_config(BASE, MODEL, hermes_config) + assert hermes_config.read_text() == original # user-managed file left untouched + assert "couldn't parse" in capsys.readouterr().err + + +def test_connect_hermes_no_launch(fake_studio, tmp_path): + yaml = pytest.importorskip("yaml") + result = CliRunner().invoke(start.start_app, ["hermes", "--no-launch"]) + assert result.exit_code == 0, result.output + _assert_env_set(result.output, "UNSLOTH_API_KEY", "sk-unsloth-feedfacefeedface") + # HERMES_HOME relocates the whole hermes home, so the user's ~/.hermes is untouched. + home = tmp_path / "agents" / "hermes" + _assert_env_set(result.output, "HERMES_HOME", str(home)) + assert "hermes" in result.output + config = yaml.safe_load((home / "config.yaml").read_text()) + assert config["model"]["provider"] == "custom:unsloth" + assert config["providers"]["unsloth"]["base_url"] == f"{BASE}/v1" + assert config["model"]["default"] == MODEL["id"] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +# ── Pi (OpenAI-compatible /v1, key in config, ~/.pi relocated via HOME) ── + + +def test_write_pi_config_fresh(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + provider = config["providers"]["unsloth"] + assert provider["api"] == "openai-completions" + assert provider["baseUrl"] == f"{BASE}/v1" + assert provider["apiKey"] == "sk-unsloth-abc" + # Pin the loaded window (and a sane output cap) so Pi compacts instead of + # overflowing; without it Pi assumes its 128000 default. + assert provider["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + + +def test_write_pi_config_preserves_and_idempotent(tmp_path): + path = tmp_path / ".pi" / "agent" / "models.json" + path.parent.mkdir(parents = True) + path.write_text(json.dumps({"providers": {"google": {"api": "gemini"}}})) + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["providers"]["google"] == {"api": "gemini"} # unrelated provider kept + assert config["providers"]["unsloth"]["baseUrl"] == f"{BASE}/v1" + before = path.read_text() + start.write_pi_config(BASE, "sk-unsloth-abc", MODEL, path) + assert path.read_text() == before + + +def test_connect_pi_no_launch(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + # Pi resolves its config dir from PI_CODING_AGENT_DIR first, so pin it at the session + # dir (and relocate HOME) to keep the user's real ~/.pi untouched and their own + # PI_CODING_AGENT_DIR from redirecting Pi away from our provider/key. + home = tmp_path / "agents" / "pi" + _assert_env_set(result.output, "HOME", str(home)) + _assert_env_set(result.output, "PI_CODING_AGENT_DIR", str(home / ".pi" / "agent")) + # Provider/model pinned on the command (Pi defaults to google otherwise). + assert f"pi --provider unsloth --model {MODEL['id']}" in result.output + config = json.loads((home / ".pi" / "agent" / "models.json").read_text()) + assert config["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["providers"]["unsloth"]["models"] == [ + {"id": MODEL["id"], "contextWindow": MODEL["context_length"], "maxTokens": 8192} + ] + assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) + + +def test_connect_pi_no_launch_windows_relocates_userprofile(fake_studio, tmp_path, monkeypatch): + # On native Windows Node resolves ~/.pi via USERPROFILE, not HOME, so the session + # must point USERPROFILE at the relocated home or Pi reads the user's real ~/.pi. + monkeypatch.setattr(start.os, "name", "nt") + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + home = tmp_path / "agents" / "pi" + assert f'$env:HOME = "{home}"' in result.output + assert f'$env:USERPROFILE = "{home}"' in result.output + + +# ── WSLENV path translation + PowerShell quoting (helper units) ── + + +def test_wsl_bridge_names_flags_paths_not_scalars(): + # WSLENV only translates a var to a Windows path when its entry carries /p. + # Path-valued vars must get it; scalar knobs and URLs must not, or WSLENV would + # mangle them when handing off to a Windows shim under /mnt. + env = { + "CODEX_HOME": "/tmp/sess/codex", + "HOME": "/tmp/sess/pi", + "CLAUDE_CODE_AUTO_COMPACT_WINDOW": "4096", + "ANTHROPIC_BASE_URL": "http://127.0.0.1:8888", + "USERPROFILE": r"C:\Users\x", + } + names = start._wsl_bridge_names(env, ("ANTHROPIC_API_KEY",)) + assert "CODEX_HOME/p" in names + assert "HOME/p" in names + assert "USERPROFILE/p" in names # drive-qualified Windows path + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW" in names # scalar: no /p + assert "CLAUDE_CODE_AUTO_COMPACT_WINDOW/p" not in names + assert "ANTHROPIC_BASE_URL" in names # URL is not a filesystem path + assert "ANTHROPIC_API_KEY" in names # cleared var carries no value to translate + + +def test_merge_wslenv_dedups_on_base_name(): + # An already-shared var must not be appended again just because the flag differs. + merged = start._merge_wslenv("CODEX_HOME/p:FOO", ("CODEX_HOME/p", "BAR/p")) + parts = merged.split(":") + assert parts.count("CODEX_HOME/p") == 1 + assert "FOO" in parts and "BAR/p" in parts + + +def test_merge_wslenv_upgrades_existing_unflagged_entry(): + # A user's pre-existing bare "HOME" must be upgraded to "HOME/p" (not left bare or + # duplicated), or the Windows shim gets the path without WSL translation. + merged = start._merge_wslenv("HOME:FOO", ("HOME/p", "CODEX_HOME/p")) + parts = merged.split(":") + assert "HOME/p" in parts and "HOME" not in parts # upgraded in place + assert parts.count("HOME/p") == 1 + assert "FOO" in parts # untouched user var preserved + assert "CODEX_HOME/p" in parts + + +def test_powershell_quote_single_quotes_json(): + # Bare flags/paths pass through; JSON payloads get single-quoted so PowerShell + # keeps the embedded double quotes literal (list2cmdline's backslashes would not). + assert start._powershell_quote("--settings") == "--settings" + assert start._powershell_quote("unsloth/gemma-4-26B") == "unsloth/gemma-4-26B" + quoted = start._powershell_quote(start._CLAUDE_SETTINGS_OVERLAY) + assert quoted == "'" + start._CLAUDE_SETTINGS_OVERLAY + "'" + assert "\\" not in quoted # no cmd.exe backslash escaping + assert start._powershell_quote("a'b") == "'a''b'" # embedded quote doubled + + +# ── --yolo: one switch routed to each agent's own auto-approve form ── + +# The native "run tools without prompting" CLI flag each agent should receive. +_NATIVE_YOLO = { + "claude": "--dangerously-skip-permissions", + "codex": "--dangerously-bypass-approvals-and-sandbox", + "hermes": "--yolo", + "pi": "--approve", +} + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_yolo_routes_to_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + assert native in result.output + + +@pytest.mark.parametrize("agent, native", sorted(_NATIVE_YOLO.items())) +def test_no_yolo_omits_native_flag(fake_studio, agent, native): + result = CliRunner().invoke(start.start_app, [agent, "--no-launch"]) + assert result.exit_code == 0, result.output + # pi's --approve is a real flag only added under --yolo; assert it's absent here. + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert native not in command + + +@pytest.mark.parametrize( + "alias", + ["--yolo", "--dangerously-skip-permissions", "--dangerously-bypass-approvals-and-sandbox"], +) +def test_yolo_aliases_are_interchangeable(fake_studio, alias): + # Any spelling on any agent routes to that agent's own flag, even the "wrong" one. + claude = CliRunner().invoke(start.start_app, ["claude", alias, "--no-launch"]) + assert claude.exit_code == 0, claude.output + assert "--dangerously-skip-permissions" in claude.output + # The codex spelling must not leak through to Claude's command line. + assert "--dangerously-bypass-approvals-and-sandbox" not in claude.output + + codex = CliRunner().invoke(start.start_app, ["codex", alias, "--no-launch"]) + assert codex.exit_code == 0, codex.output + assert "--dangerously-bypass-approvals-and-sandbox" in codex.output + assert "--dangerously-skip-permissions" not in codex.output + + +def test_yolo_opencode_writes_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_no_yolo_opencode_has_no_permission_block(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + config = json.loads((tmp_path / "agents" / "opencode" / "opencode.json").read_text()) + assert "permission" not in config + + +def test_yolo_openclaw_writes_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + # Both layers: the host approvals file in OPENCLAW_STATE_DIR must also be set, or + # OpenClaw can still prompt/deny despite the config. + approvals = json.loads((state / "exec-approvals.json").read_text()) + assert approvals["defaults"] == {"security": "full", "ask": "off", "askFallback": "full"} + + +def test_no_yolo_openclaw_has_no_exec_policy(fake_studio, tmp_path): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch"]) + assert result.exit_code == 0, result.output + state = tmp_path / "agents" / "openclaw" + config = json.loads((state / "openclaw.json").read_text()) + assert "exec" not in config.get("tools", {}) # no auto-approve policy without --yolo + assert not (state / "exec-approvals.json").exists() + + +def test_write_opencode_config_yolo_unit(tmp_path): + path = tmp_path / "opencode.json" + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} + + +def test_write_openclaw_config_yolo_unit(tmp_path): + path = tmp_path / "openclaw.json" + start.write_openclaw_config(BASE, "sk-unsloth-abc", MODEL, path, yolo = True) + config = json.loads(path.read_text()) + assert config["tools"]["exec"] == {"host": "gateway", "security": "full", "ask": "off"} + approvals = json.loads((path.parent / "exec-approvals.json").read_text()) + assert approvals == { + "version": 1, + "defaults": {"security": "full", "ask": "off", "askFallback": "full"}, + } + + +def test_yolo_command_flags_unmapped_agent_is_empty(): + # Config-based agents (and any typo) must yield no flag, not a KeyError. + assert start._yolo_command_flags("opencode", True) == [] + assert start._yolo_command_flags("openclaw", True) == [] + assert start._yolo_command_flags("claude", True) == ["--dangerously-skip-permissions"] + assert start._yolo_command_flags("claude", False) == [] + + +def test_yolo_config_agents_add_no_command_flag(fake_studio): + # opencode/openclaw auto-approve is config-only; nothing should leak onto argv. + for agent in ("opencode", "openclaw"): + result = CliRunner().invoke(start.start_app, [agent, "--yolo", "--no-launch"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command and command[0] == agent, result.output + assert not any("--yolo" in arg or "--dangerous" in arg for arg in command) + + +def test_pi_launch_clears_screen_first(fake_studio, monkeypatch): + # Pi paints inline from the current cursor position (no alternate screen, no + # clear on its first render), so the launcher hands it a clean screen. The + # clear must come BEFORE the exec, and only on the launch path. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/pi") + + def run(command, env): + calls.append("exec") + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + assert calls == ["clear", "exec"] + + +def test_pi_no_launch_does_not_clear(fake_studio, monkeypatch): + # The --no-launch recipe is meant to be read (and piped); never wipe it. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + result = CliRunner().invoke(start.start_app, ["pi", "--no-launch"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +def test_claude_launch_does_not_clear(fake_studio, monkeypatch): + # Alternate-screen agents manage the terminal themselves; leave it alone. + calls = [] + monkeypatch.setattr(start.click, "clear", lambda: calls.append("clear")) + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + monkeypatch.setattr(start.subprocess, "run", lambda command, env: SimpleNamespace(returncode = 0)) + result = CliRunner().invoke(start.start_app, ["claude"]) + assert result.exit_code == 0, result.output + assert calls == [] + + +@pytest.mark.skipif( + os.name == "nt", + reason = "WSL-from-Linux scenario: a Windows pi shim under /mnt called from WSL " + "(os.name is 'posix' under WSL), so this can't run on a native Windows runner.", +) +def test_connect_pi_wsl_windows_shim_relocates_userprofile(fake_studio, monkeypatch): + captured = {} + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr(start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/pi") + + def run(command, env): + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, ["pi"]) + assert result.exit_code == 0, result.output + home = captured["env"]["HOME"] + # A Windows pi shim resolves ~/.pi via USERPROFILE, so it must match the session + # HOME and ride the WSLENV bridge (with /p) so the path is translated for Windows. + assert captured["env"]["USERPROFILE"] == home + wslenv = captured["env"]["WSLENV"].split(":") + assert "HOME/p" in wslenv + assert "USERPROFILE/p" in wslenv + + +def test_agent_api_key_auto_started_rejected_env_key_falls_back(fake_studio, tmp_path, monkeypatch): + # UNSLOTH_API_KEY exported for some OTHER server must not fail the launch + # against a server this run just auto-started: validate, then fall back to + # the local mint path, and never remember the foreign key for this base. + inner = start._http_json + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + if url.endswith("/v1/models") and token == "sk-unsloth-other-server": + raise urllib.error.HTTPError(url, 401, "Unauthorized", None, None) + return inner(method, url, token, payload, timeout, error) + + monkeypatch.setattr(start, "_http_json", http_json) + key = start._agent_api_key(BASE, "sk-unsloth-other-server", auto_started = True) + assert key == "sk-unsloth-feedfacefeedface" # minted for the fresh server + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert "sk-unsloth-other-server" not in json.dumps(cached["servers"].get(BASE, {})) + + +def test_agent_api_key_auto_started_accepted_key_is_honored(fake_studio, tmp_path): + # An explicit key the fresh server accepts (e.g. persisted in this Studio + # home's auth db across restarts) keeps working exactly as before. + key = start._agent_api_key(BASE, "sk-unsloth-deadbeefdeadbeef", auto_started = True) + assert key == "sk-unsloth-deadbeefdeadbeef" + cached = json.loads((tmp_path / "agent_api_key.json").read_text()) + assert cached["servers"][BASE]["saved"] == ["sk-unsloth-deadbeefdeadbeef"] + + +def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path): + # A previously printed recipe may still be running an agent whose sessions + # or sqlite state live in the stable home; a re-run must not wipe it. + with start._session_config("codex", launch = False) as home: + marker = home / "sessions" / "live.sqlite" + marker.parent.mkdir(parents = True) + marker.write_text("state") + with start._session_config("codex", launch = False) as home2: + assert home2 == home + assert (home2 / "sessions" / "live.sqlite").read_text() == "state"