From 734cec9e7aef78088f3b39649f89c1353306d6d3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 03:15:45 -0700 Subject: [PATCH 001/161] Studio STT: only load safetensors weights for custom dictation models (RCE fix) (#7364) * Studio STT: only load safetensors weights for custom dictation models The STT sidecar accepts arbitrary Hugging Face owner/model repos for custom dictation models and, when safetensors were absent, downloaded and loaded pytorch_model.bin through WhisperForConditionalGeneration .from_pretrained. PyTorch checkpoints are pickles that execute code during deserialization, and this path does not run the malware gate the normal model loader applies, so an authenticated client on an exposed Studio instance could load a crafted Whisper-looking repo and run code in the backend. Restrict custom STT repos to safetensors: the snapshot selector no longer falls back to pytorch_model.bin(.index.json), the cached-snapshot completeness check ignores pickle weights, and the load forces use_safetensors so a stray cached pickle still cannot execute. The five curated Whisper defaults already ship safetensors only, so this changes nothing for the built-in models. * STT: reject safetensors indexes that reference non-safetensors shards A safetensors index (model.safetensors.index.json) is attacker-supplied JSON and can name pytorch_model-*.bin shards in its weight_map. Transformers dispatches shard loading per file by extension, so those .bin shards still load through torch.load (pickle) even with use_safetensors set. Require every weight_map value to end in .safetensors in both the snapshot selector and the completeness check so no pickle shard is downloaded or reused. --- studio/backend/core/inference/stt_sidecar.py | 62 +++++++++++-------- .../backend/tests/test_stt_review_fixes_2.py | 30 +++++++-- studio/backend/tests/test_stt_sidecar.py | 40 ++++++++++++ 3 files changed, 101 insertions(+), 31 deletions(-) diff --git a/studio/backend/core/inference/stt_sidecar.py b/studio/backend/core/inference/stt_sidecar.py index 7ac7e93c87..edf57c16e3 100644 --- a/studio/backend/core/inference/stt_sidecar.py +++ b/studio/backend/core/inference/stt_sidecar.py @@ -49,8 +49,11 @@ _MAX_AUDIO_SECONDS = 30 * 60 _TARGET_SAMPLE_RATE = 16000 # Non-weight files WhisperProcessor/WhisperForConditionalGeneration may load. -# Weight selection is built from pinned Hub metadata so repositories publishing -# both formats do not download the same checkpoint twice. +# Weight selection is built from pinned Hub metadata. A custom repo id is +# attacker-controllable, so only safetensors weights are accepted: a +# pytorch_model.bin is a pickle and executes code while Transformers +# deserializes it (see utils/security/file_security.py), and this path skips +# the malware gate the normal model loader applies. _STT_SNAPSHOT_SUPPORT_FILES = ( "config.json", "generation_config.json", @@ -65,9 +68,7 @@ _STT_SNAPSHOT_SUPPORT_FILES = ( "added_tokens.json", ) _STT_SAFETENSORS_INDEX = "model.safetensors.index.json" -_STT_PYTORCH_INDEX = "pytorch_model.bin.index.json" _STT_SAFETENSORS_WEIGHTS = "model.safetensors" -_STT_PYTORCH_WEIGHTS = "pytorch_model.bin" _STT_REVISION_RECORD_VERSION = 1 @@ -350,7 +351,9 @@ def _selected_file_from_sibling(sibling) -> _SelectedHubFile: def _select_snapshot_files(info, load_index) -> tuple[_SelectedHubFile, ...]: - """Select support files and exactly one complete Transformers weight format.""" + """Select support files and one complete safetensors weight set. Pickle + (pytorch_model.bin) weights are never selected: they are an RCE sink on a + custom repo id (see _STT_SNAPSHOT_SUPPORT_FILES).""" siblings = { sibling.rfilename: sibling for sibling in (getattr(info, "siblings", None) or []) @@ -363,12 +366,11 @@ def _select_snapshot_files(info, load_index) -> tuple[_SelectedHubFile, ...]: index_name = _STT_SAFETENSORS_INDEX elif _STT_SAFETENSORS_WEIGHTS in siblings: selected.add(_STT_SAFETENSORS_WEIGHTS) - elif _STT_PYTORCH_INDEX in siblings: - index_name = _STT_PYTORCH_INDEX - elif _STT_PYTORCH_WEIGHTS in siblings: - selected.add(_STT_PYTORCH_WEIGHTS) else: - raise SttModelCompatibilityError("The STT repository has no complete model weights.") + raise SttModelCompatibilityError( + "The STT repository has no safetensors model weights. Only safetensors " + "checkpoints are supported; convert the model with save_pretrained(safe_serialization=True)." + ) if index_name is not None: weight_map = load_index(index_name).get("weight_map") @@ -377,6 +379,14 @@ def _select_snapshot_files(info, load_index) -> tuple[_SelectedHubFile, ...]: shards = set(weight_map.values()) if not all(isinstance(shard, str) and shard in siblings for shard in shards): raise SttModelCompatibilityError(f"Checkpoint index '{index_name}' has missing shards.") + # The index JSON is attacker-controlled: a safetensors index can name + # pytorch_model-*.bin shards, which Transformers still loads through + # torch.load (pickle) since it dispatches per shard by file extension. + # Require every shard to be safetensors so no pickle file is selected. + if not all(shard.endswith(".safetensors") for shard in shards): + raise SttModelCompatibilityError( + f"Checkpoint index '{index_name}' references non-safetensors shards." + ) selected.add(index_name) selected.update(shards) @@ -441,24 +451,23 @@ def _snapshot_is_complete(snapshot: Path) -> bool: tokenizer, and weights directly. is_file() follows cache symlinks, so a link from an interrupted blob download does not count. """ - index = next( - ( - snapshot / name - for name in ("model.safetensors.index.json", "pytorch_model.bin.index.json") - if (snapshot / name).is_file() - ), - None, - ) - if index is not None: - # Sharded checkpoint (safetensors or PyTorch): every shard must exist. + # Safetensors only: a cached pytorch_model.bin is a pickle load path and is + # never treated as a usable snapshot (a repo shipping only pickle weights + # re-resolves and fails closed in _select_snapshot_files). + index = snapshot / _STT_SAFETENSORS_INDEX + if index.is_file(): + # Sharded safetensors checkpoint: every shard must exist and be + # safetensors (a safe index naming .bin shards would still pickle-load + # them, matching the _select_snapshot_files guard). weight_map = _read_json_object(index).get("weight_map") if not isinstance(weight_map, dict) or not weight_map: return False - has_weights = all((snapshot / shard).is_file() for shard in set(weight_map.values())) + shards = set(weight_map.values()) + if not all(isinstance(shard, str) and shard.endswith(".safetensors") for shard in shards): + return False + has_weights = all((snapshot / shard).is_file() for shard in shards) else: - has_weights = any( - (snapshot / name).is_file() for name in (_STT_SAFETENSORS_WEIGHTS, _STT_PYTORCH_WEIGHTS) - ) + has_weights = (snapshot / _STT_SAFETENSORS_WEIGHTS).is_file() # WhisperProcessor needs the tokenizer: either the fast tokenizer.json or # the slow vocab.json + merges.txt pair. has_tokenizer = (snapshot / "tokenizer.json").is_file() or ( @@ -880,8 +889,11 @@ class WhisperSttSidecar: try: processor = WhisperProcessor.from_pretrained(snapshot_path, local_files_only = True) self._raise_if_load_cancelled(cancel_event) + # use_safetensors forces the pickle-free load path even if a + # pytorch_model.bin somehow reached the cache; the selector and the + # completeness check already exclude pickle weights upstream. model = WhisperForConditionalGeneration.from_pretrained( - snapshot_path, torch_dtype = dtype, local_files_only = True + snapshot_path, torch_dtype = dtype, local_files_only = True, use_safetensors = True ) self._raise_if_load_cancelled(cancel_event) model.to(torch.device(device)) diff --git a/studio/backend/tests/test_stt_review_fixes_2.py b/studio/backend/tests/test_stt_review_fixes_2.py index 130c43c956..f0bdab42b5 100644 --- a/studio/backend/tests/test_stt_review_fixes_2.py +++ b/studio/backend/tests/test_stt_review_fixes_2.py @@ -6,8 +6,8 @@ 1. scripts/build_whisper_cpp.sh must not rm -rf a whisper.cpp/src tree under a custom Studio home unless Studio itself created it (ownership marker), the same policy studio/setup.sh applies before its destructive replacements. -2. _snapshot_is_complete must validate every shard of a sharded PyTorch - (pytorch_model.bin.index.json) checkpoint, like the safetensors path. +2. _snapshot_is_complete must reject pickle (pytorch_model.bin) checkpoints + outright; only safetensors weights count as a usable snapshot. 3. _snapshot_is_complete must require tokenizer assets (tokenizer.json or vocab.json + merges.txt); weights + config alone decode to blank text. 4. Custom-repo downloads must pin the revision validated beforehand and @@ -143,7 +143,10 @@ def _base_snapshot(tmp_path: Path) -> Path: return snap -def test_sharded_pytorch_snapshot_requires_every_shard(tmp_path): +def test_pickle_checkpoint_snapshot_is_never_complete(tmp_path): + # A cached pytorch_model.bin is a pickle RCE load path; the snapshot must + # read as incomplete no matter how many shards are present, so update + # re-resolves and _select_snapshot_files fails it closed. snap = _base_snapshot(tmp_path) index = { "weight_map": { @@ -153,14 +156,29 @@ def test_sharded_pytorch_snapshot_requires_every_shard(tmp_path): } (snap / "pytorch_model.bin.index.json").write_text(json.dumps(index)) (snap / "pytorch_model-00001-of-00002.bin").write_bytes(b"w" * 8) - - # One missing .bin shard must read as incomplete, like the safetensors path. + (snap / "pytorch_model-00002-of-00002.bin").write_bytes(b"w" * 8) assert stt_sidecar_module._snapshot_is_complete(snap) is False - (snap / "pytorch_model-00002-of-00002.bin").write_bytes(b"w" * 8) + # A single-file pickle checkpoint is likewise rejected; the safetensors + # equivalent in the same dir makes it complete. + (snap / "pytorch_model.bin").write_bytes(b"w" * 8) + assert stt_sidecar_module._snapshot_is_complete(snap) is False + (snap / "model.safetensors").write_bytes(b"w" * 8) assert stt_sidecar_module._snapshot_is_complete(snap) is True +def test_safe_index_naming_pickle_shards_is_not_complete(tmp_path): + # A safetensors index that references .bin shards would still pickle-load + # via Transformers' per-shard dispatch; the cached snapshot must read as + # incomplete so it re-resolves and fails closed at selection. + snap = _base_snapshot(tmp_path) + (snap / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"a": "pytorch_model-00001-of-00001.bin"}}) + ) + (snap / "pytorch_model-00001-of-00001.bin").write_bytes(b"w" * 8) + assert stt_sidecar_module._snapshot_is_complete(snap) is False + + def test_snapshot_without_tokenizer_assets_is_incomplete(tmp_path): snap = _base_snapshot(tmp_path) (snap / "model.safetensors").write_bytes(b"w" * 8) diff --git a/studio/backend/tests/test_stt_sidecar.py b/studio/backend/tests/test_stt_sidecar.py index 3f78845ac7..b138f46331 100644 --- a/studio/backend/tests/test_stt_sidecar.py +++ b/studio/backend/tests/test_stt_sidecar.py @@ -486,6 +486,9 @@ def test_load_uses_model_hub_cache_without_implicit_download(monkeypatch): } # Never fetch weights implicitly; the Model Hub owns downloads. assert all(kwargs.get("local_files_only") is True for _, _, kwargs in calls) + # The weight load forces safetensors so a pickle checkpoint cannot execute. + model_kwargs = next(kwargs for kind, _, kwargs in calls if kind == "model") + assert model_kwargs.get("use_safetensors") is True def test_model_cache_preflight_uses_shared_offline_resolver(monkeypatch): @@ -1046,6 +1049,43 @@ def test_snapshot_selection_includes_every_indexed_shard(): } +def test_snapshot_selection_rejects_pickle_only_weights(): + # A custom repo shipping only pytorch_model.bin (pickle) must fail closed: + # selecting it would download a checkpoint that runs code at load time. + info = SimpleNamespace( + siblings = [ + _sibling("config.json", 10, "config"), + _sibling("preprocessor_config.json", 20, "preprocessor"), + _sibling("tokenizer.json", 30, "tokenizer"), + _sibling("pytorch_model.bin", 110, "torch"), + ] + ) + + with pytest.raises(SttModelCompatibilityError, match = "safetensors"): + stt_sidecar_module._select_snapshot_files( + info, lambda _name: pytest.fail("pickle weights must not be selected") + ) + + +def test_snapshot_selection_rejects_safe_index_pointing_at_pickle_shards(): + # A safetensors index can name .bin shards; Transformers dispatches shard + # loading by extension, so those shards would still pickle-load. The index + # is attacker-controlled, so a non-safetensors shard must fail closed. + info = SimpleNamespace( + siblings = [ + _sibling("config.json", 10, "config"), + _sibling("model.safetensors.index.json", 5, "index"), + _sibling("pytorch_model-00001-of-00001.bin", 90, "shard"), + ] + ) + + with pytest.raises(SttModelCompatibilityError, match = "non-safetensors shards"): + stt_sidecar_module._select_snapshot_files( + info, + lambda _name: {"weight_map": {"a": "pytorch_model-00001-of-00001.bin"}}, + ) + + def test_progress_counts_only_selected_blobs_and_caps_incomplete_files(monkeypatch, tmp_path): monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path / "hub")) blobs = tmp_path / "hub" / "models--owner--whisper" / "blobs" From 6e868860bdb7602b65cca95634b3a6fdb5130415 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 03:17:25 -0700 Subject: [PATCH 002/161] Bump install.sh / install.ps1 pin to unsloth>=2026.7.5 (#7365) --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index c06f3a1120..9c91d4ba16 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2347,7 +2347,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2361,7 +2361,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2435,7 +2435,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2447,7 +2447,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2475,7 +2475,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index 44d51490f7..84c5be9742 100755 --- a/install.sh +++ b/install.sh @@ -3471,7 +3471,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -3488,7 +3488,7 @@ if [ "$_MIGRATED" = true ]; then run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" ${_MLX_LM_EXCLUDE_ARG:-} [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" _UNSLOTH_TORCH_OVERRIDES="" fi @@ -3712,7 +3712,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" + "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -3731,7 +3731,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ - --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" + --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -3759,7 +3759,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." From a26692612d361ad3ca53ebcca1606b6af406c168 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:54:09 +0100 Subject: [PATCH 003/161] Normalize PWD for POSIX agent launches (#7110) Keep the child process environment consistent with the cwd used to launch native POSIX coding agents. Some Node-based agents use PWD during project-root discovery, so inheriting a stale PWD can make them edit files in a parent or unrelated directory even when the wrapper process cwd is correct. Only apply this normalization for native POSIX launches. WSL-launched Windows shims stay on the existing WSLENV bridge path so path translation behavior is unchanged. Add regression coverage that launches an agent with a deliberately stale inherited PWD and asserts the child environment is normalized to os.getcwd(). Co-authored-by: Leo Borcherding --- unsloth_cli/commands/start.py | 5 +++++ unsloth_cli/tests/test_start.py | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 01dffa0634..7c768d16f5 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -2110,6 +2110,11 @@ def _launch( for name in unset_env: child_env.pop(name, None) child_env.update(env) + if os.name != "nt" and not wsl_env_bridge: + # Keep POSIX child processes from seeing a stale inherited PWD when + # subprocess cwd was changed by the caller. Some Node CLIs use PWD for + # project-root discovery instead of process.cwd(). + child_env["PWD"] = os.getcwd() # Ctrl+C cancels a turn inside the agent; don't let it kill this wrapper. previous = signal.signal(signal.SIGINT, signal.SIG_IGN) try: diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index e76ae24f8b..36a5e51938 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -874,6 +874,27 @@ def test_connect_claude_compact_window_omitted_without_context(fake_studio, monk assert "CLAUDE_AUTOCOMPACT_PCT_OVERRIDE" not in result.output +def test_launch_native_posix_child_gets_current_pwd(fake_studio, monkeypatch, tmp_path): + captured = {} + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("PWD", "/stale/outer/repo") + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode") + + 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, ["opencode"]) + + assert result.exit_code == 0, result.output + assert captured["command"][0] == "/usr/local/bin/opencode" + if os.name != "nt": + assert captured["env"]["PWD"] == os.getcwd() + + def test_connect_claude_launch_scrubs_conflicting_auth_env(fake_studio, monkeypatch): captured = {} monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-anthropic-stale") From 447f16f49ae6a326087c978476b992cbc4f8de4f Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Fri, 24 Jul 2026 04:59:13 +0530 Subject: [PATCH 004/161] Studio: fix composer reset after failed send (#7377) * fix composer reset * Studio: clear composer draft on send * Studio: cancel the pending draft save when clearing on send --- .../src/components/assistant-ui/thread.tsx | 28 ++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 43a0c43950..68f96c46c8 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -1556,6 +1556,7 @@ const Composer: FC<{ const draftThreadId = referenceThreadId; const draftKey = draftThreadId ? composerDraftKey(draftThreadId) : null; const lastDraftKeyRef = useRef(draftKey); + const draftSaveTimerRef = useRef | null>(null); useEffect(() => { const draft = draftKey ? (readComposerDraft(draftKey) ?? "") : ""; const composer = aui.composer(); @@ -1574,8 +1575,25 @@ const Composer: FC<{ return; } const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300); + draftSaveTimerRef.current = t; return () => clearTimeout(t); }, [composerText, draftKey]); + // Without this the restore effect above puts the sent text back when the + // runtime rebinds on the first message. + const draftKeyRef = useRef(draftKey); + useEffect(() => { + draftKeyRef.current = draftKey; + }, [draftKey]); + const clearStoredDraft = useCallback(() => { + if (draftSaveTimerRef.current !== null) { + clearTimeout(draftSaveTimerRef.current); + draftSaveTimerRef.current = null; + } + const key = draftKeyRef.current; + if (key) { + writeComposerDraft(key, ""); + } + }, []); // react-textarea-autosize re-measures only on value change or window resize, // not on the width swap from expanding, so it keeps the taller height and // leaves a stray blank row. Nudge a resize whenever input width changes. @@ -1726,9 +1744,10 @@ const Composer: FC<{ setPendingSend(false); dismissWaitToast(); if (text.trim().length > 0 || attachments.length > 0) { + clearStoredDraft(); aui.composer().send(); } - }, [pendingSend, indexingActive, aui, dismissWaitToast]); + }, [pendingSend, indexingActive, aui, clearStoredDraft, dismissWaitToast]); // Drop any queued send + toast on unmount (e.g. thread switch). useEffect( @@ -1771,6 +1790,7 @@ const Composer: FC<{ flushResourcesSync(() => { aui.composer().setText(""); }); + clearStoredDraft(); startPromptQueue( [queuedPrompt], createPromptQueueTarget(), @@ -1804,6 +1824,7 @@ const Composer: FC<{ closeOverlay(); return; } + clearStoredDraft(); setImageToolsEnabled(true); setPendingImageEditReference({ threadId: overlay.threadId ?? referenceThreadId, @@ -1821,11 +1842,15 @@ const Composer: FC<{ ); }); closeOverlay(); + return; } + + clearStoredDraft(); }, [ aui, canQueueCurrentPrompt, + clearStoredDraft, closeOverlay, composerText, createPromptQueueTarget, @@ -1947,6 +1972,7 @@ const Composer: FC<{ flushResourcesSync(() => { aui.composer().setText(""); }); + clearStoredDraft(); startPromptQueue([queuedPrompt], createPromptQueueTarget(), true); }} onSendClick={interceptSend} From e2ccf4d376200a6ea1542a2d99492a0be553d4f0 Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:41:05 +0530 Subject: [PATCH 005/161] fix(studio): show chat sidebar menu on touch devices (#7297) * fix(studio): show chat sidebar menu on touch devices Recents/Pinned chat row actions were hidden until hover, so iPad users could not open the kebab menu to delete chats. Reveal actions on coarse pointers using the same pattern as hub model rows. Fixes #7276 * Fix coarse-pointer sidebar row action visibility (#7276) Move the touch-device override into index.css after .sidebar-row-action so it wins the cascade. Arbitrary Tailwind media utilities on the element had equal specificity and were overridden by the base rule. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scope coarse-pointer sidebar actions to chat rows (#7276) Only chat kebabs/unpin buttons that reserve touch padding get sidebar-touch-reveal, so project/run/nav rows stay hover-revealed. * Tighten comments * Reserve full kebab hit area on coarse-pointer unpinned rows --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../frontend/src/components/app-sidebar.tsx | 15 ++++++------ studio/frontend/src/index.css | 6 +++++ ...t_desktop_reliability_frontend_contract.py | 24 +++++++++++++++++++ 3 files changed, 38 insertions(+), 7 deletions(-) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index a6f64243ad..c95112c748 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -901,8 +901,8 @@ export function AppSidebar() { : "group/recent-item relative"; const actionClass = variant === "project" - ? "sidebar-row-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" - : "sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; + ? "sidebar-row-action sidebar-touch-reveal group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" + : "sidebar-row-action sidebar-touch-reveal group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"; const buttonClass = cn( "sidebar-nav-btn h-[33px] cursor-pointer rounded-full pr-4 text-ui-14p5 leading-ui-19 tracking-nav font-medium", // pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the @@ -912,13 +912,14 @@ export function AppSidebar() { isPinned && variant !== "project" && "gap-[8.5px]", variant === "project" ? // Room for the hover pin quick-action plus the kebab. - "group-hover/project-chat-item:pr-14 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8" + "group-hover/project-chat-item:pr-14 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8 [@media(pointer:coarse)]:pr-14" : isPinned ? // Pinned rows show an extra unpin button on hover, so reserve more room // (pr-8 when the menu is open keeps the unpin button clear of the title). - "group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8" + "group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8 [@media(pointer:coarse)]:pr-16" : // Hover room for the kebab only; title keeps one more character. - "group-hover/recent-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-6", + // Touch rows clear the full always-visible kebab hit area (pr-10). + "group-hover/recent-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-6 [@media(pointer:coarse)]:pr-10", ); const isRenamingThis = @@ -987,7 +988,7 @@ export function AppSidebar() { togglePinnedChat(item.id); }} aria-label={isPinned ? "Unpin chat" : "Pin chat"} - className="sidebar-row-action is-unpin-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" + className="sidebar-row-action sidebar-touch-reveal is-unpin-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" > @@ -1002,7 +1003,7 @@ export function AppSidebar() { togglePinnedChat(item.id); }} aria-label="Unpin chat" - className="sidebar-row-action is-unpin-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" + className="sidebar-row-action sidebar-touch-reveal is-unpin-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto" > diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css index 0cba3407fe..ba112d5f53 100644 --- a/studio/frontend/src/index.css +++ b/studio/frontend/src/index.css @@ -1018,6 +1018,12 @@ html[data-chat-font] .aui-root { .sidebar-row-action { @apply absolute top-0 bottom-0 right-0 inline-flex cursor-pointer items-center justify-end pl-2 pr-1.5 opacity-0 pointer-events-none outline-none; } + @media (pointer: coarse) { + /* Only chat rows reserve touch padding (#7276); other rows stay hover-revealed to avoid clipped labels. */ + .sidebar-row-action.sidebar-touch-reveal { + @apply opacity-100 pointer-events-auto; + } + } .sidebar-row-action[data-state="open"] { @apply opacity-100 pointer-events-auto; } diff --git a/tests/studio/test_desktop_reliability_frontend_contract.py b/tests/studio/test_desktop_reliability_frontend_contract.py index 868895b8f0..2552b6c1d8 100644 --- a/tests/studio/test_desktop_reliability_frontend_contract.py +++ b/tests/studio/test_desktop_reliability_frontend_contract.py @@ -14,6 +14,7 @@ DATA_TAB = FRONTEND / "features/settings/tabs/data-tab.tsx" PROMPT_STORAGE = FRONTEND / "features/chat/prompt-storage/prompt-storage-dialog.tsx" APP_SIDEBAR = FRONTEND / "components/app-sidebar.tsx" +INDEX_CSS = FRONTEND / "index.css" THREAD = FRONTEND / "components/assistant-ui/thread.tsx" THREAD_SIDEBAR = FRONTEND / "features/chat/thread-sidebar.tsx" SHARED_COMPOSER = FRONTEND / "features/chat/shared-composer.tsx" @@ -129,3 +130,26 @@ def test_expanded_titlebar_button_and_corner_match_sidebar_edge(): 'className="pointer-events-none absolute top-full size-3 -translate-x-px rounded-tl-[12px] border-l border-t border-sidebar-border bg-background"' in source ) + + +def test_chat_sidebar_row_actions_visible_on_coarse_pointers(): + """unslothai/unsloth#7276: Recents chat kebab must be tappable on iPad.""" + sidebar_source = APP_SIDEBAR.read_text(encoding = "utf-8") + css_source = INDEX_CSS.read_text(encoding = "utf-8") + assert "renderChatSidebarItem" in sidebar_source + block = sidebar_source.split("function renderChatSidebarItem", 1)[1].split("\n function ", 1)[ + 0 + ] + assert "[@media(pointer:coarse)]:pr-10" in block + assert "sidebar-touch-reveal" in block + # Coarse-pointer visibility must come after .sidebar-row-action { opacity-0 }. + coarse_idx = css_source.index("@media (pointer: coarse)") + base_idx = css_source.index(".sidebar-row-action {") + assert coarse_idx > base_idx + coarse_block = css_source[coarse_idx : coarse_idx + 280] + assert "sidebar-touch-reveal" in coarse_block + assert "opacity-100" in coarse_block + assert "pointer-events-auto" in coarse_block + # Must not reveal every sidebar-row-action (project/run/nav rows lack padding). + assert ".sidebar-row-action {\n\t\t\t@apply opacity-100" not in coarse_block + assert ".sidebar-row-action.sidebar-touch-reveal" in coarse_block From b448fb5de02a7996f4738cc7de2600d368006221 Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:41:50 +0530 Subject: [PATCH 006/161] fix(studio): persist connection model selections for remote clients (#7298) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(studio): persist connection model selections server-side Remote Studio clients could see saved connections but not their enabled model lists because models lived only in browser localStorage. Store models and available_models in llm_providers and sync them through the providers API so alternate clients inherit the same catalog state. Fixes #7281 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Hydrate external connections on chat startup (#7281) Extract provider sync logic into sync-external-providers.ts and call it from chat-page on mount so persisted model selections appear in the Connected picker without opening Settings → Connections first. * fix(studio): backfill connection models and preserve local options (#7298) Address Codex P2 on remote connection persistence: - Backfill localStorage model selections to /api/providers when backend rows still have empty models_json (legacy upgrades) - Carry promptCacheTtl and openaiContainerTtlMinutes through startup sync - Await hydratePersistedSettings before syncing on ChatPage mount Contract tests: 7 passed; npm run typecheck passed. * Tighten comments * Tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/models/providers.py | 21 ++ studio/backend/routes/providers.py | 51 ++-- studio/backend/storage/providers_db.py | 91 +++++++- .../backend/tests/test_providers_db_models.py | 70 ++++++ .../assistant-ui/tool-ui-python.tsx | 6 +- studio/frontend/src/config/training.ts | 6 +- .../src/features/chat/api/providers-api.ts | 12 + .../frontend/src/features/chat/chat-page.tsx | 15 +- .../features/chat/chat-providers-dialog.tsx | 137 ++--------- .../features/chat/sync-external-providers.ts | 221 ++++++++++++++++++ .../studio/sections/params-section.tsx | 18 +- .../studio/sections/progress-section.tsx | 5 +- .../test_remote_connection_models_contract.py | 61 +++++ 13 files changed, 530 insertions(+), 184 deletions(-) create mode 100644 studio/backend/tests/test_providers_db_models.py create mode 100644 studio/frontend/src/features/chat/sync-external-providers.ts create mode 100644 tests/studio/test_remote_connection_models_contract.py diff --git a/studio/backend/models/providers.py b/studio/backend/models/providers.py index 5a75246c07..4238403e00 100644 --- a/studio/backend/models/providers.py +++ b/studio/backend/models/providers.py @@ -47,6 +47,14 @@ class ProviderCreate(BaseModel): None, description = "Custom base URL (overrides registry default). Omit to use the default.", ) + models: list[str] = Field( + default_factory = list, + description = "Enabled model IDs for this connection", + ) + available_models: list[str] = Field( + default_factory = list, + description = "Discovered catalog model IDs last fetched for this connection", + ) class ProviderUpdate(BaseModel): @@ -55,6 +63,11 @@ class ProviderUpdate(BaseModel): display_name: Optional[str] = Field(None, description = "New display name") base_url: Optional[str] = Field(None, description = "New base URL") is_enabled: Optional[bool] = Field(None, description = "Enable or disable this provider") + models: Optional[list[str]] = Field(None, description = "Enabled model IDs for this connection") + available_models: Optional[list[str]] = Field( + None, + description = "Discovered catalog model IDs last fetched for this connection", + ) class ProviderResponse(BaseModel): @@ -65,6 +78,14 @@ class ProviderResponse(BaseModel): display_name: str = Field(..., description = "User-chosen label") base_url: str = Field(..., description = "API base URL") is_enabled: bool = Field(True, description = "Whether this provider is enabled") + models: list[str] = Field( + default_factory = list, + description = "Enabled model IDs for this connection", + ) + available_models: list[str] = Field( + default_factory = list, + description = "Discovered catalog model IDs last fetched for this connection", + ) created_at: str = Field(..., description = "ISO 8601 creation timestamp") updated_at: str = Field(..., description = "ISO 8601 last-update timestamp") diff --git a/studio/backend/routes/providers.py b/studio/backend/routes/providers.py index 5a55c9b0bb..4e7e53f2f0 100644 --- a/studio/backend/routes/providers.py +++ b/studio/backend/routes/providers.py @@ -47,6 +47,20 @@ logger = structlog.get_logger(__name__) router = APIRouter() +def _provider_response(row: dict) -> ProviderResponse: + return ProviderResponse( + id = row["id"], + provider_type = row["provider_type"], + display_name = row["display_name"], + base_url = row["base_url"], + is_enabled = bool(row["is_enabled"]), + models = row.get("models") or [], + available_models = row.get("available_models") or [], + created_at = row["created_at"], + updated_at = row["updated_at"], + ) + + # ── Public key for API key encryption ───────────────────────────── @@ -89,18 +103,7 @@ async def get_pricing_snapshot(current_subject: str = Depends(get_current_subjec async def list_provider_configs(current_subject: str = Depends(get_current_subject)): """List all saved provider configurations.""" rows = providers_db.list_providers() - return [ - ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) - for row in rows - ] + return [_provider_response(row) for row in rows] @router.post("/", response_model = ProviderResponse, status_code = 201) @@ -124,18 +127,12 @@ async def create_provider_config( provider_type = payload.provider_type, display_name = payload.display_name, base_url = base_url, + models = payload.models, + available_models = payload.available_models, ) row = providers_db.get_provider(provider_id) - return ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) + return _provider_response(row) @router.put("/{provider_id}", response_model = ProviderResponse) @@ -154,20 +151,14 @@ async def update_provider_config( display_name = payload.display_name, base_url = payload.base_url, is_enabled = payload.is_enabled, + models = payload.models, + available_models = payload.available_models, ) if not updated: raise HTTPException(status_code = 400, detail = "No fields to update") row = providers_db.get_provider(provider_id) - return ProviderResponse( - id = row["id"], - provider_type = row["provider_type"], - display_name = row["display_name"], - base_url = row["base_url"], - is_enabled = bool(row["is_enabled"]), - created_at = row["created_at"], - updated_at = row["updated_at"], - ) + return _provider_response(row) @router.delete("/{provider_id}", status_code = 204) diff --git a/studio/backend/storage/providers_db.py b/studio/backend/storage/providers_db.py index 07165cbe70..e6f40c5030 100644 --- a/studio/backend/storage/providers_db.py +++ b/studio/backend/storage/providers_db.py @@ -6,8 +6,12 @@ Same pattern as studio_db.py (module-level functions, raw sqlite3, WAL, per-function connections). API keys are NOT stored here: they live only in the browser (localStorage) and are sent encrypted per-request. + +Enabled model selections and discovered catalog IDs are stored server-side so +remote Studio clients see the same connection state (#7281). """ +import json import logging import sqlite3 import threading @@ -22,6 +26,33 @@ _schema_lock = threading.Lock() _schema_ready = False +def _encode_models_json(models: Optional[list[str]]) -> str: + if not models: + return "[]" + return json.dumps([str(model).strip() for model in models if str(model).strip()]) + + +def _decode_models_json(raw: Optional[str]) -> list[str]: + if not raw: + return [] + try: + parsed = json.loads(raw) + except json.JSONDecodeError: + return [] + if not isinstance(parsed, list): + return [] + return [str(model).strip() for model in parsed if str(model).strip()] + + +def _row_models(row: sqlite3.Row) -> tuple[list[str], list[str]]: + return ( + _decode_models_json(row["models_json"] if "models_json" in row.keys() else None), + _decode_models_json( + row["available_models_json"] if "available_models_json" in row.keys() else None + ), + ) + + def _ensure_schema(conn: sqlite3.Connection) -> None: """Create the llm_providers table if absent. Called once per process.""" conn.execute("PRAGMA journal_mode=WAL") @@ -38,6 +69,13 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: ) """ ) + existing_cols = {row[1] for row in conn.execute("PRAGMA table_info(llm_providers)").fetchall()} + if "models_json" not in existing_cols: + conn.execute("ALTER TABLE llm_providers ADD COLUMN models_json TEXT NOT NULL DEFAULT '[]'") + if "available_models_json" not in existing_cols: + conn.execute( + "ALTER TABLE llm_providers ADD COLUMN available_models_json TEXT NOT NULL DEFAULT '[]'" + ) def get_connection() -> sqlite3.Connection: @@ -59,17 +97,37 @@ def get_connection() -> sqlite3.Connection: return conn -def create_provider(id: str, provider_type: str, display_name: str, base_url: str) -> None: +def create_provider( + id: str, + provider_type: str, + display_name: str, + base_url: str, + models: Optional[list[str]] = None, + available_models: Optional[list[str]] = None, +) -> None: """Insert a new provider configuration.""" now = datetime.now(timezone.utc).isoformat() conn = get_connection() try: conn.execute( """ - INSERT INTO llm_providers (id, provider_type, display_name, base_url, created_at, updated_at) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO llm_providers ( + id, provider_type, display_name, base_url, + models_json, available_models_json, + created_at, updated_at + ) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, - (id, provider_type, display_name, base_url, now, now), + ( + id, + provider_type, + display_name, + base_url, + _encode_models_json(models), + _encode_models_json(available_models), + now, + now, + ), ) conn.commit() finally: @@ -81,6 +139,8 @@ def update_provider( display_name: Optional[str] = None, base_url: Optional[str] = None, is_enabled: Optional[bool] = None, + models: Optional[list[str]] = None, + available_models: Optional[list[str]] = None, ) -> bool: """Update fields on an existing provider. Returns True if a row was updated.""" updates = [] @@ -94,6 +154,12 @@ def update_provider( if is_enabled is not None: updates.append("is_enabled = ?") params.append(1 if is_enabled else 0) + if models is not None: + updates.append("models_json = ?") + params.append(_encode_models_json(models)) + if available_models is not None: + updates.append("available_models_json = ?") + params.append(_encode_models_json(available_models)) if not updates: return False updates.append("updated_at = ?") @@ -128,7 +194,13 @@ def get_provider(id: str) -> Optional[dict]: conn = get_connection() try: row = conn.execute("SELECT * FROM llm_providers WHERE id = ?", (id,)).fetchone() - return dict(row) if row else None + if not row: + return None + data = dict(row) + models, available_models = _row_models(row) + data["models"] = models + data["available_models"] = available_models + return data finally: conn.close() @@ -138,6 +210,13 @@ def list_providers() -> list[dict]: conn = get_connection() try: rows = conn.execute("SELECT * FROM llm_providers ORDER BY created_at").fetchall() - return [dict(row) for row in rows] + providers: list[dict] = [] + for row in rows: + data = dict(row) + models, available_models = _row_models(row) + data["models"] = models + data["available_models"] = available_models + providers.append(data) + return providers finally: conn.close() diff --git a/studio/backend/tests/test_providers_db_models.py b/studio/backend/tests/test_providers_db_models.py new file mode 100644 index 0000000000..ca9dffbd70 --- /dev/null +++ b/studio/backend/tests/test_providers_db_models.py @@ -0,0 +1,70 @@ +# 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 provider model persistence (unslothai/unsloth#7281).""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import storage.providers_db as providers_db + + +@pytest.fixture() +def isolated_providers_db(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + db_path = tmp_path / "studio.db" + monkeypatch.setattr(providers_db, "studio_db_path", lambda: db_path) + monkeypatch.setattr(providers_db, "ensure_dir", lambda _path: None) + providers_db._schema_ready = False + yield db_path + providers_db._schema_ready = False + + +def test_create_and_list_provider_models(isolated_providers_db: Path): + providers_db.create_provider( + id = "ollama1", + provider_type = "ollama", + display_name = "Home Ollama", + base_url = "http://127.0.0.1:11434", + models = ["llama3.2", "qwen2.5"], + available_models = ["llama3.2", "qwen2.5", "mistral"], + ) + + row = providers_db.get_provider("ollama1") + assert row is not None + assert row["models"] == ["llama3.2", "qwen2.5"] + assert row["available_models"] == ["llama3.2", "qwen2.5", "mistral"] + + listed = providers_db.list_providers() + assert len(listed) == 1 + assert listed[0]["models"] == ["llama3.2", "qwen2.5"] + + +def test_update_provider_models(isolated_providers_db: Path): + providers_db.create_provider( + id = "vllm1", + provider_type = "vllm", + display_name = "Remote vLLM", + base_url = "http://studio-host:8000/v1", + models = ["meta-llama/Llama-3.2-1B-Instruct"], + available_models = ["meta-llama/Llama-3.2-1B-Instruct"], + ) + + assert providers_db.update_provider( + id = "vllm1", + models = ["meta-llama/Llama-3.2-3B-Instruct"], + available_models = [ + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B-Instruct", + ], + ) + + row = providers_db.get_provider("vllm1") + assert row is not None + assert row["models"] == ["meta-llama/Llama-3.2-3B-Instruct"] + assert row["available_models"] == [ + "meta-llama/Llama-3.2-1B-Instruct", + "meta-llama/Llama-3.2-3B-Instruct", + ] diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx index 1e816132b8..34e51b9d5a 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-python.tsx @@ -83,7 +83,7 @@ function CopyBtn({ text }: { text: string }) { ); } -/** Save the executed script as a .py file via a client-side Blob (no server file serving). */ +/** Save the script as a .py file via a client-side Blob. */ function DownloadBtn({ code, name = "script.py" }: { code: string; name?: string }) { const download = useCallback(() => { if (typeof document === "undefined") { @@ -229,8 +229,8 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({ const authToken = getAuthToken(); return ( - // Run status and output collapse from history, but the script source is - // rendered outside ToolFallbackContent so it stays visible on reopen (#7165). + // Status/output collapse from history; the script source renders outside + // ToolFallbackContent so it stays visible on reopen (#7165). { value: "adamw_torch_fused", label: "AdamW (PyTorch Fused)" }, ]; -// Optimizers the MLX trainer actually supports on Apple Silicon. Values must -// match SUPPORTED_MLX_OPTIMIZERS in unsloth-zoo's mlx/trainer.py; on MLX the -// bitsandbytes/torch names above have no meaning and are remapped to plain -// AdamW, so Studio offers this list instead when running on a Mac. +// MLX trainer optimizers (Apple Silicon); must match SUPPORTED_MLX_OPTIMIZERS in +// unsloth-zoo's mlx/trainer.py. The CUDA/torch names above are remapped to AdamW on MLX. export const MLX_OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ { value: "adamw", label: "AdamW" }, { value: "adam", label: "Adam" }, diff --git a/studio/frontend/src/features/chat/api/providers-api.ts b/studio/frontend/src/features/chat/api/providers-api.ts index 4ad996d54f..c6c5613272 100644 --- a/studio/frontend/src/features/chat/api/providers-api.ts +++ b/studio/frontend/src/features/chat/api/providers-api.ts @@ -23,6 +23,8 @@ export interface ProviderConfig { display_name: string; base_url: string; is_enabled: boolean; + models?: string[]; + available_models?: string[]; created_at: string; updated_at: string; } @@ -123,6 +125,8 @@ export async function createProviderConfig(payload: { providerType: string; displayName: string; baseUrl?: string | null; + models?: string[]; + availableModels?: string[]; }): Promise { const response = await authFetch("/api/providers/", { method: "POST", @@ -131,6 +135,8 @@ export async function createProviderConfig(payload: { provider_type: payload.providerType, display_name: payload.displayName, base_url: payload.baseUrl ?? null, + models: payload.models ?? [], + available_models: payload.availableModels ?? [], }), }); return parseJsonOrThrow(response); @@ -158,6 +164,8 @@ export async function updateProviderConfig( displayName?: string; baseUrl?: string | null; isEnabled?: boolean; + models?: string[]; + availableModels?: string[]; }, ): Promise { const response = await authFetch(`/api/providers/${providerId}`, { @@ -167,6 +175,10 @@ export async function updateProviderConfig( ...(payload.displayName === undefined ? {} : { display_name: payload.displayName }), ...(payload.baseUrl === undefined ? {} : { base_url: payload.baseUrl }), ...(payload.isEnabled === undefined ? {} : { is_enabled: payload.isEnabled }), + ...(payload.models === undefined ? {} : { models: payload.models }), + ...(payload.availableModels === undefined + ? {} + : { available_models: payload.availableModels }), }), }); return parseJsonOrThrow(response); diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index e46ea0ac46..c241607e28 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -175,6 +175,7 @@ import { } from "./stores/chat-runtime-store"; import { useChatPreferencesStore } from "./stores/chat-preferences-store"; import { useExternalProvidersStore } from "./stores/external-providers-store"; +import { syncExternalProvidersFromBackend } from "./sync-external-providers"; import { buildChatTourSteps } from "./tour"; import type { ChatView, MessageRecord } from "./types"; import { @@ -1762,8 +1763,18 @@ export function ChatPage({ const externalProvidersForChat = connectionsEnabled ? externalProviders : []; useEffect(() => { - void hydratePersistedSettings(); - }, [hydratePersistedSettings]); + void (async () => { + await hydratePersistedSettings(); + try { + const synced = await syncExternalProvidersFromBackend( + useExternalProvidersStore.getState().providers, + ); + setExternalProviders(synced); + } catch { + // Silent on startup; Connections settings still surfaces load errors. + } + })(); + }, [hydratePersistedSettings, setExternalProviders]); useEffect(() => { // Skip while off-route: ChatPage stays mounted, and toast+navigate here would diff --git a/studio/frontend/src/features/chat/chat-providers-dialog.tsx b/studio/frontend/src/features/chat/chat-providers-dialog.tsx index bfe7c71918..4ab560e5ad 100644 --- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx @@ -41,7 +41,6 @@ import { type ProviderRegistryEntry, createProviderConfig, deleteProviderConfig, - listProviderConfigs, listProviderModels, listProviderRegistry, testProviderConnection, @@ -49,7 +48,6 @@ import { } from "./api/providers-api"; import type { ExternalProviderConfig } from "./external-providers"; import { - CUSTOM_BACKEND_PROVIDER_TYPE, CUSTOM_PROVIDER_PRESETS, allowsManualModelIdsWithCatalog, customProviderBaseUrlPlaceholder, @@ -68,6 +66,10 @@ import { toExternalBackendProviderType, } from "./external-providers"; import { useExternalProvidersStore } from "./stores/external-providers-store"; +import { + pruneProviderModelIds, + syncExternalProvidersFromBackend, +} from "./sync-external-providers"; /** Matches navbar / thread layout easing (see index.css --ease-out-quart) */ const PROVIDER_FORM_EASE: [number, number, number, number] = [ @@ -76,58 +78,7 @@ const PROVIDER_FORM_EASE: [number, number, number, number] = [ const PROVIDER_FORM_DURATION = 0.2; const CUSTOM_PROVIDER_MISSING_KEY_MESSAGE = "No API key found. Add a valid API key for this connection."; -const ANTHROPIC_DATED_SNAPSHOT_SUFFIX = /-\d{8}$/; -const OPENAI_DEPRECATED_MODELS = new Set(["gpt-5.3"]); const HIDDEN_PROVIDER_TYPES = new Set(["qwen"]); -const OPENROUTER_EXCLUDED_MODELS = new Set([ - "google/chirp-3", - "kwaivgi/kling-v3.0-pro", - "openai/whisper-1", - "openai/gpt-4o-mini-transcribe", - "recraft/recraft-v4-pro", -]); - -function normalizeUrl(input: string): string { - return input.trim().replace(/\/+$/, ""); -} - -function resolveUiProviderTypeFromConfig( - configProviderType: string, - configDisplayName: string | null | undefined, - configBaseUrl: string | null | undefined, - registryRows: ProviderRegistryEntry[], - existingProviderType: string | undefined, -): string { - if (existingProviderType && isCustomProviderType(existingProviderType)) { - return existingProviderType; - } - if (configProviderType !== CUSTOM_BACKEND_PROVIDER_TYPE) { - return configProviderType; - } - const displayName = (configDisplayName ?? "").trim().toLowerCase(); - const matchingCustomPreset = CUSTOM_PROVIDER_PRESETS.find( - (preset) => preset.displayName.toLowerCase() === displayName, - ); - if (matchingCustomPreset) { - return matchingCustomPreset.providerType; - } - const openAiRegistry = registryRows.find( - (entry) => entry.provider_type === CUSTOM_BACKEND_PROVIDER_TYPE, - ); - if (!openAiRegistry) { - return configProviderType; - } - const openAiDisplayName = openAiRegistry.display_name.trim().toLowerCase(); - if (displayName.length > 0 && displayName !== openAiDisplayName) { - return LEGACY_CUSTOM_PROVIDER_TYPE; - } - const configUrl = normalizeUrl(configBaseUrl ?? ""); - const defaultUrl = normalizeUrl(openAiRegistry.base_url ?? ""); - if (configUrl.length > 0 && configUrl !== defaultUrl) { - return LEGACY_CUSTOM_PROVIDER_TYPE; - } - return configProviderType; -} function parseManualModelIds(text: string): string[] { const seen = new Set(); @@ -182,19 +133,6 @@ function shouldAppendOpenAiVersionPath(providerType: string): boolean { ); } -function pruneProviderModelIds(providerType: string, modelIds: string[]): string[] { - if (providerType === "anthropic") { - return modelIds.filter((id) => !ANTHROPIC_DATED_SNAPSHOT_SUFFIX.test(id)); - } - if (providerType === "openai") { - return modelIds.filter((id) => !OPENAI_DEPRECATED_MODELS.has(id)); - } - if (providerType === "openrouter") { - return modelIds.filter((id) => !OPENROUTER_EXCLUDED_MODELS.has(id)); - } - return modelIds; -} - function formatModelSummary(models: string[]): string { if (models.length === 0) { return "No models enabled"; @@ -360,9 +298,9 @@ export function ChatProvidersSettings({ } let syncSucceeded = false; try { - const [registryRows, configRows] = await Promise.all([ + const [registryRows, syncedProviders] = await Promise.all([ listProviderRegistry(), - listProviderConfigs(), + syncExternalProvidersFromBackend(providersRef.current), ]); if (!isMounted) return; syncSucceeded = true; @@ -377,61 +315,6 @@ export function ChatProvidersSettings({ } return registryRows[0]?.provider_type ?? ""; }); - const existingById = new Map(); - for (const provider of providersRef.current) { - existingById.set(provider.id, provider); - } - const syncedProviders: ExternalProviderConfig[] = configRows - .filter((config) => config.is_enabled) - .map((config) => { - const existing = existingById.get(config.id); - const uiProviderType = resolveUiProviderTypeFromConfig( - config.provider_type, - config.display_name, - config.base_url, - registryRows, - existing?.providerType, - ); - const createdAt = Number.isFinite(Date.parse(config.created_at)) - ? Date.parse(config.created_at) - : Date.now(); - const updatedAt = Number.isFinite(Date.parse(config.updated_at)) - ? Date.parse(config.updated_at) - : Date.now(); - const registryEntry = - registryRows.find((entry) => entry.provider_type === uiProviderType) ?? - registryRows.find((entry) => entry.provider_type === config.provider_type); - const defaultModels = pruneProviderModelIds( - uiProviderType, - registryEntry?.default_models ?? [], - ); - const savedModels = existing?.models ?? []; - const savedAvailableModels = existing?.availableModels ?? []; - const existingModels = pruneProviderModelIds( - uiProviderType, - savedModels.length > 0 ? savedModels : defaultModels, - ); - const existingAvailableModels = pruneProviderModelIds( - uiProviderType, - savedAvailableModels.length > 0 ? savedAvailableModels : defaultModels, - ); - return { - id: config.id, - providerType: uiProviderType, - name: config.display_name, - baseUrl: config.base_url ?? "", - models: existingModels, - availableModels: existingAvailableModels, - enablePromptCaching: supportsProviderPromptCaching(uiProviderType) - ? (existing?.enablePromptCaching ?? true) - : undefined, - isReasoningModel: supportsProviderReasoningToggle(uiProviderType) - ? existing?.isReasoningModel === true - : undefined, - createdAt: existing?.createdAt ?? createdAt, - updatedAt, - }; - }); // Trust the backend response. An empty array means every connection was // removed (often from another tab); mirror that locally, else stale // entries become un-removable here until localStorage is cleared. @@ -699,6 +582,10 @@ export function ChatProvidersSettings({ providerType: backendProviderType, displayName, baseUrl, + models: modelsToSave, + availableModels: manualOnly + ? [] + : pruneProviderModelIds(providerType, availableModels), }); const createdAt = Number.isFinite(Date.parse(created.created_at)) ? Date.parse(created.created_at) @@ -814,6 +701,10 @@ export function ChatProvidersSettings({ customProviderDisplayName(existing.providerType) : existing.name, baseUrl, + models: modelsToSave, + availableModels: manualOnly + ? [] + : pruneProviderModelIds(existing.providerType, availableModels), }); if (apiKey.trim()) { setExternalProviderApiKey(editingProviderId, apiKey.trim()); diff --git a/studio/frontend/src/features/chat/sync-external-providers.ts b/studio/frontend/src/features/chat/sync-external-providers.ts new file mode 100644 index 0000000000..cd966f7fd4 --- /dev/null +++ b/studio/frontend/src/features/chat/sync-external-providers.ts @@ -0,0 +1,221 @@ +// 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 { + type ProviderRegistryEntry, + listProviderConfigs, + listProviderRegistry, + updateProviderConfig, +} from "./api/providers-api"; +import { + CUSTOM_BACKEND_PROVIDER_TYPE, + CUSTOM_PROVIDER_PRESETS, + type ExternalProviderConfig, + isCustomProviderType, + isPromptCacheTtl, + LEGACY_CUSTOM_PROVIDER_TYPE, + supportsProviderPromptCaching, + supportsProviderPromptCacheTtl, + supportsProviderReasoningToggle, +} from "./external-providers"; + +const ANTHROPIC_DATED_SNAPSHOT_SUFFIX = /-\d{8}$/; +const OPENAI_DEPRECATED_MODELS = new Set(["gpt-5.3"]); +const OPENROUTER_EXCLUDED_MODELS = new Set([ + "google/chirp-3", + "kwaivgi/kling-v3.0-pro", + "openai/whisper-1", + "openai/gpt-4o-mini-transcribe", + "recraft/recraft-v4-pro", +]); + +function normalizeUrl(input: string): string { + return input.trim().replace(/\/+$/, ""); +} + +export function resolveUiProviderTypeFromConfig( + configProviderType: string, + configDisplayName: string | null | undefined, + configBaseUrl: string | null | undefined, + registryRows: ProviderRegistryEntry[], + existingProviderType: string | undefined, +): string { + if (existingProviderType && isCustomProviderType(existingProviderType)) { + return existingProviderType; + } + if (configProviderType !== CUSTOM_BACKEND_PROVIDER_TYPE) { + return configProviderType; + } + const displayName = (configDisplayName ?? "").trim().toLowerCase(); + const matchingCustomPreset = CUSTOM_PROVIDER_PRESETS.find( + (preset) => preset.displayName.toLowerCase() === displayName, + ); + if (matchingCustomPreset) { + return matchingCustomPreset.providerType; + } + const openAiRegistry = registryRows.find( + (entry) => entry.provider_type === CUSTOM_BACKEND_PROVIDER_TYPE, + ); + if (!openAiRegistry) { + return configProviderType; + } + const openAiDisplayName = openAiRegistry.display_name.trim().toLowerCase(); + if (displayName.length > 0 && displayName !== openAiDisplayName) { + return LEGACY_CUSTOM_PROVIDER_TYPE; + } + const configUrl = normalizeUrl(configBaseUrl ?? ""); + const defaultUrl = normalizeUrl(openAiRegistry.base_url ?? ""); + if (configUrl.length > 0 && configUrl !== defaultUrl) { + return LEGACY_CUSTOM_PROVIDER_TYPE; + } + return configProviderType; +} + +export function pruneProviderModelIds( + providerType: string, + modelIds: string[], +): string[] { + if (providerType === "anthropic") { + return modelIds.filter((id) => !ANTHROPIC_DATED_SNAPSHOT_SUFFIX.test(id)); + } + if (providerType === "openai") { + return modelIds.filter((id) => !OPENAI_DEPRECATED_MODELS.has(id)); + } + if (providerType === "openrouter") { + return modelIds.filter((id) => !OPENROUTER_EXCLUDED_MODELS.has(id)); + } + return modelIds; +} + +/** Carry browser-local provider knobs through a backend sync rebuild. */ +export function mergeLocalProviderOptions( + existing: ExternalProviderConfig | undefined, + synced: ExternalProviderConfig, +): ExternalProviderConfig { + if (!existing) { + return synced; + } + const providerType = synced.providerType; + return { + ...synced, + enablePromptCaching: supportsProviderPromptCaching(providerType) + ? (existing.enablePromptCaching ?? synced.enablePromptCaching ?? true) + : undefined, + promptCacheTtl: + supportsProviderPromptCacheTtl(providerType) && + isPromptCacheTtl(existing.promptCacheTtl) + ? existing.promptCacheTtl + : synced.promptCacheTtl, + isReasoningModel: supportsProviderReasoningToggle(providerType) + ? (existing.isReasoningModel ?? synced.isReasoningModel) + : undefined, + openaiContainerTtlMinutes: + providerType === "openai" && + typeof existing.openaiContainerTtlMinutes === "number" && + existing.openaiContainerTtlMinutes >= 1 + ? Math.min(existing.openaiContainerTtlMinutes, 20) + : synced.openaiContainerTtlMinutes, + }; +} + +/** Merge enabled backend provider configs with local store state. */ +export async function syncExternalProvidersFromBackend( + existingProviders: ExternalProviderConfig[], +): Promise { + const [registryRows, configRows] = await Promise.all([ + listProviderRegistry(), + listProviderConfigs(), + ]); + + const existingById = new Map(); + for (const provider of existingProviders) { + existingById.set(provider.id, provider); + } + + const backfillTasks: Promise[] = []; + const syncedProviders = configRows + .filter((config) => config.is_enabled) + .map((config) => { + const existing = existingById.get(config.id); + const uiProviderType = resolveUiProviderTypeFromConfig( + config.provider_type, + config.display_name, + config.base_url, + registryRows, + existing?.providerType, + ); + const createdAt = Number.isFinite(Date.parse(config.created_at)) + ? Date.parse(config.created_at) + : Date.now(); + const updatedAt = Number.isFinite(Date.parse(config.updated_at)) + ? Date.parse(config.updated_at) + : Date.now(); + const registryEntry = + registryRows.find((entry) => entry.provider_type === uiProviderType) ?? + registryRows.find((entry) => entry.provider_type === config.provider_type); + const defaultModels = pruneProviderModelIds( + uiProviderType, + registryEntry?.default_models ?? [], + ); + const serverModels = pruneProviderModelIds( + uiProviderType, + config.models ?? [], + ); + const serverAvailableModels = pruneProviderModelIds( + uiProviderType, + config.available_models ?? [], + ); + const savedModels = existing?.models ?? []; + const savedAvailableModels = existing?.availableModels ?? []; + const resolvedModels = pruneProviderModelIds( + uiProviderType, + serverModels.length > 0 + ? serverModels + : savedModels.length > 0 + ? savedModels + : defaultModels, + ); + const resolvedAvailableModels = pruneProviderModelIds( + uiProviderType, + serverAvailableModels.length > 0 + ? serverAvailableModels + : savedAvailableModels.length > 0 + ? savedAvailableModels + : defaultModels, + ); + const needsModelBackfill = + serverModels.length === 0 && savedModels.length > 0; + const needsAvailableBackfill = + serverAvailableModels.length === 0 && savedAvailableModels.length > 0; + if (needsModelBackfill || needsAvailableBackfill) { + backfillTasks.push( + updateProviderConfig(config.id, { + models: resolvedModels, + availableModels: resolvedAvailableModels, + }), + ); + } + const synced: ExternalProviderConfig = { + id: config.id, + providerType: uiProviderType, + name: config.display_name, + baseUrl: config.base_url ?? "", + models: resolvedModels, + availableModels: resolvedAvailableModels, + enablePromptCaching: supportsProviderPromptCaching(uiProviderType) + ? (existing?.enablePromptCaching ?? true) + : undefined, + isReasoningModel: supportsProviderReasoningToggle(uiProviderType) + ? existing?.isReasoningModel === true + : undefined, + createdAt: existing?.createdAt ?? createdAt, + updatedAt, + }; + return mergeLocalProviderOptions(existing, synced); + }); + + if (backfillTasks.length > 0) { + await Promise.allSettled(backfillTasks); + } + return syncedProviders; +} diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 3270eb4e3d..f2323dcc6b 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -205,26 +205,19 @@ export function ParamsSection(): ReactElement { setCtxInput(String(store.contextLength)); }, [store.contextLength]); - // On Apple Silicon the MLX trainer supports a different optimizer set than - // the CUDA/bitsandbytes list, so offer the MLX names there. + // Apple Silicon (MLX) supports a different optimizer set than the CUDA list. const isMac = platformDeviceType === "mac"; const optimizerOptions = isMac ? MLX_OPTIMIZER_OPTIONS : OPTIMIZER_OPTIONS; - // On Mac, the MLX backend normalizes every CUDA/bitsandbytes optimizer in - // OPTIMIZER_OPTIONS (including the shared default) to plain AdamW, so show - // AdamW for those to keep the control truthful and non-blank. Any other - // value -- an MLX optimizer the user picked, or an unrecognized/non-canonical - // imported one -- is shown as-is rather than mislabeled as AdamW, since the - // backend would run or reject it on its own terms. Non-Mac display unchanged. + // On Mac the MLX backend remaps CUDA optimizers to AdamW, so label those as + // AdamW; other values (MLX or imported) show as-is. Non-Mac unchanged. const isCudaAliasOptimizer = OPTIMIZER_OPTIONS.some( (o) => o.value === store.optimizerType, ); const selectedOptimizer = isMac && isCudaAliasOptimizer ? "adamw" : store.optimizerType; - // LoftQ is not supported on MLX (the backend rejects it), so clear a stale - // selection to lora on Apple Silicon -- whether persisted, applied from a - // model default, or imported -- so the backend never receives it. + // LoftQ is unsupported on MLX; clear a stale selection to lora on Apple Silicon. const setLoraVariant = store.setLoraVariant; useEffect(() => { if (isMac && store.loraVariant === "loftq") { @@ -232,8 +225,7 @@ export function ParamsSection(): ReactElement { } }, [isMac, store.loraVariant, setLoraVariant]); - // Packing is not supported on MLX (the backend forces it off), so clear it on - // Apple Silicon -- the checkbox is disabled and the flag is never sent. + // Packing is unsupported on MLX; clear it on Apple Silicon (checkbox disabled). const setPacking = store.setPacking; useEffect(() => { if (isMac && store.packing) { diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index ad813e888e..a2cf711296 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -189,9 +189,8 @@ export function ProgressSection({ const cfgLoraDropout = cfg?.loraDropout; const cfgLoraVariant = cfg?.loraVariant; - // Mirror the training form: on Mac the CUDA/bitsandbytes optimizer names run - // as plain AdamW (the MLX backend normalizes them), so label them AdamW here - // too rather than by the requested, unnormalized name. + // Mirror the training form: on Mac the MLX backend runs CUDA optimizers as + // AdamW, so label them AdamW here too. const effectiveOptimizer = platformDeviceType === "mac" && OPTIMIZER_OPTIONS.some((o) => o.value === cfgOptimizerType) diff --git a/tests/studio/test_remote_connection_models_contract.py b/tests/studio/test_remote_connection_models_contract.py new file mode 100644 index 0000000000..47124541e0 --- /dev/null +++ b/tests/studio/test_remote_connection_models_contract.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Static contracts for remote connection model persistence (#7281).""" + +from pathlib import Path + + +REPO = Path(__file__).resolve().parents[2] +FRONTEND = REPO / "studio/frontend/src" +PROVIDERS_API = FRONTEND / "features/chat/api/providers-api.ts" +SYNC_PROVIDERS = FRONTEND / "features/chat/sync-external-providers.ts" +CHAT_PAGE = FRONTEND / "features/chat/chat-page.tsx" +PROVIDERS_DB = REPO / "studio/backend/storage/providers_db.py" +PROVIDERS_MODELS = REPO / "studio/backend/models/providers.py" + + +def test_providers_db_stores_model_json_columns(): + source = PROVIDERS_DB.read_text(encoding = "utf-8") + assert "models_json" in source + assert "available_models_json" in source + assert "ALTER TABLE llm_providers ADD COLUMN models_json" in source + + +def test_provider_api_schemas_expose_models(): + source = PROVIDERS_MODELS.read_text(encoding = "utf-8") + assert "models: list[str]" in source + assert "available_models: list[str]" in source + + +def test_frontend_sync_prefers_server_models_on_remote_clients(): + source = SYNC_PROVIDERS.read_text(encoding = "utf-8") + assert "config.models" in source + assert "config.available_models" in source + assert "serverModels.length > 0" in source + + +def test_frontend_sync_backfills_local_models_to_backend(): + source = SYNC_PROVIDERS.read_text(encoding = "utf-8") + assert "updateProviderConfig" in source + assert "needsModelBackfill" in source + assert "Promise.allSettled(backfillTasks)" in source + + +def test_frontend_sync_preserves_local_provider_options(): + source = SYNC_PROVIDERS.read_text(encoding = "utf-8") + assert "mergeLocalProviderOptions" in source + assert "promptCacheTtl" in source + assert "openaiContainerTtlMinutes" in source + + +def test_chat_page_hydrates_connections_on_startup(): + source = CHAT_PAGE.read_text(encoding = "utf-8") + assert "syncExternalProvidersFromBackend" in source + assert "await hydratePersistedSettings()" in source + + +def test_providers_api_sends_models_to_backend(): + source = PROVIDERS_API.read_text(encoding = "utf-8") + assert "available_models: payload.availableModels" in source + assert "models: payload.models" in source From 8c975fcbaf28361faa601b8cca6e4613368cde79 Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:42:52 +0530 Subject: [PATCH 007/161] fix: pin torchcodec for torch 2.10 and warn on ABI mismatch (#7299) * fix: pin torchcodec for torch 2.10 and warn on ABI mismatch Add unsloth[audio] extra with torchcodec>=0.10.0,<0.11.0 and emit a clear warning when installed torchcodec minors disagree with torch (unslothai/unsloth#7225). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(packaging): address Codex review on torchcodec/torch 2.10 compat (#7299) - Postpone annotations so import_fixes loads on Python 3.9 - Align TORCH_TORCHCODEC matrix with upstream (2.9: 0.8/0.9, 2.8: 0.6/0.7) - Fix mismatch hint upper bound (<0.11.0) and gate audio-torch210 suggestion - Split audio extra per torch minor; gate torch210 pin behind python>=3.10 - Bundle audio-torch210 only in *-torch2100 install extras * fix(security): refresh openai CRITICAL scan baseline hashes (#7299) openai package code drift reopened five CRITICAL findings in the extras pip-scan-packages shard (C2 loop body hashes + IMDS/network evidence). Update the reviewed allowlist evidence/hashes so CI gates on new findings only, not benign SDK churn. * chore: retrigger CI after baseline refresh (#7299) * chore: touch scan baseline comment to retrigger security audit (#7299) * Guard torchcodec version parsing so bad version strings cannot break import * Bundle audio pin into intel-gpu-torch210 and guard the mismatch warning * Tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- pyproject.toml | 22 +++- scripts/notebook_validator.py | 4 +- scripts/scan_packages_baseline.json | 22 ++-- tests/python/test_torchcodec_torch_compat.py | 129 +++++++++++++++++++ unsloth/import_fixes.py | 64 +++++++++ 5 files changed, 227 insertions(+), 14 deletions(-) create mode 100644 tests/python/test_torchcodec_torch_compat.py diff --git a/pyproject.toml b/pyproject.toml index fe0ebd13b9..0f57ecf4df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,6 +93,17 @@ huggingfacenotorch = [ "trl>=0.18.2,!=0.19.0,<=0.24.0", "sentence-transformers", ] +# torchcodec backend for Gemma audio / datasets>=4 (#7225). +# Pick the audio-torch* pin matching your torch minor (see TORCH_TORCHCODEC). +audio-torch210 = [ + "torchcodec>=0.10.0,<0.11.0 ; python_version >= '3.10'", +] +audio-torch290 = [ + "torchcodec>=0.8.0,<0.10.0 ; python_version >= '3.10'", +] +audio-torch280 = [ + "torchcodec>=0.6.0,<0.8.0 ; python_version >= '3.9'", +] huggingface = [ "unsloth[huggingfacenotorch]", "unsloth_zoo>=2026.7.6", @@ -532,16 +543,19 @@ cu126-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] kaggle = [ "unsloth[huggingface]", @@ -831,16 +845,19 @@ cu126-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu126onlytorch2100]", + "unsloth[audio-torch210]", ] cu128-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu128onlytorch2100]", + "unsloth[audio-torch210]", ] cu130-ampere-torch2100 = [ "unsloth[huggingface]", "bitsandbytes>=0.45.5,!=0.46.0,!=0.48.0", "unsloth[cu130onlytorch2100]", + "unsloth[audio-torch210]", ] flashattentiontorch260abiFALSEcu12x = [ "flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'", @@ -1125,7 +1142,8 @@ intelgputorch210 = [ "torchvision @ https://download.pytorch.org/whl/xpu/torchvision-0.25.0%2Bxpu-cp313-cp313-win_amd64.whl#sha256=1c4b44b36a557f7381e3076fb8843366742238648441d607c8d049c6da0f8886 ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", ] intel-gpu-torch210 = [ - "unsloth[intelgputorch210]" + "unsloth[intelgputorch210]", + "unsloth[audio-torch210]", ] intelgputorch2110 = [ "unsloth_zoo[intelgpu]", @@ -1279,6 +1297,7 @@ rocm72-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2/torchvision-0.25.0%2Brocm7.2.0.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] rocm711-torch2100 = [ "unsloth[amd]", @@ -1297,6 +1316,7 @@ rocm711-torch2100 = [ "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp311-cp311-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp312-cp312-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'", "torchvision @ https://repo.radeon.com/rocm/manylinux/rocm-rel-7.1.1/torchvision-0.25.0%2Brocm7.1.1.git82df5f59-cp313-cp313-linux_x86_64.whl ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'", + "unsloth[audio-torch210]", ] [project.urls] diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index c1be7a63a4..7bcee47c66 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -95,8 +95,8 @@ COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-i # Source: pytorch/torchcodec compatibility matrix on its README. TORCH_TORCHCODEC: dict[str, set[str]] = { "2.10": {"0.10"}, - "2.9": {"0.7", "0.8", "0.9"}, - "2.8": {"0.6"}, + "2.9": {"0.8", "0.9"}, + "2.8": {"0.6", "0.7"}, "2.7": {"0.3", "0.4", "0.5"}, "2.6": {"0.2", "0.3"}, "2.5": {"0.1", "0.2"}, diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 65b8d2b11c..1c21f8da86 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -1,5 +1,5 @@ { - "_comment": "scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", + "_comment": "scan_packages.py allowlist (reviewed). Each entry is a CRITICAL/HIGH finding manually judged benign. Matched on (package, package-relative file, check, evidence_hash); evidence_hash is over the matched code with L: markers stripped, so version bumps and line shifts do not reopen an entry but changed code does. severity and evidence are for review only. Regenerate with --write-baseline AFTER reviewing every line.", "version": 1, "entries": [ { @@ -303,8 +303,8 @@ "file": "openai/_base_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L264: while True: sha256:95ca67e46d42354ae650abbdc5b0d97df8b0ed43187800bf40f5690c3901b94b", - "evidence_hash": "a57d8d15fed0bf04f9967dcc18a18b80bb19f4095675bccbb78ac0450d7fce14" + "evidence": "L274: while True: sha256:90a38e5c1e26893c7c273354143612640e9a9c0f079d3e2b60612d79f24e80a6", + "evidence_hash": "1022e8e8649436ec64a98a9d9141d085452c49549fd2157b0278fc369a83ac66" }, { "package": "openai", @@ -319,8 +319,8 @@ "file": "openai/auth/_workload.py", "check": "Accesses cloud metadata/IMDS AND makes network calls", "severity": "CRITICAL", - "evidence": "IMDS: L96: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L149: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L77: http_client: httpx.Client | None = None, | L108: with httpx.Client() as client: | L133: http_client: httpx.Client | None = None, | L155: with httpx.Client() as client: | L248: with httpx.Client() as client:", - "evidence_hash": "1581d9f4a23393e9af23fbe5ef9f66807b22c5b5a3f1fe167254c9ebee108567" + "evidence": "IMDS: L97: url = \"http://169.254.169.254/metadata/identity/oauth2/token\" | L150: url = \"http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity\"\nNetwork: L78: http_client: httpx.Client | None = None, | L109: with httpx.Client() as client: | L134: http_client: httpx.Client | None = None, | L156: with httpx.Client() as client: | L251: exchange_client = DefaultHttpx2Client(follow_redirects=False) if self._use_httpx2 else httpx.Client()", + "evidence_hash": "9717e51cb961dc14c458955d91a1e48e3753997346ecea0106bded3a8d64bfe0" }, { "package": "openai", @@ -343,8 +343,8 @@ "file": "openai/resources/beta/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", - "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" + "evidence": "L4000: while True: sha256:f8ab538118daba9ec06e27399dbdc90a4521c3390e6a47a6348a1f180a83effd", + "evidence_hash": "31481ea83c687acc27144d72d3832d4fb98dd1c79fb5e0ddd85080de95997b9f" }, { "package": "openai", @@ -359,16 +359,16 @@ "file": "openai/resources/realtime/realtime.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L310: while True: sha256:458198ff3d3f05870bf98c9564cbfd68c739e57b9bbe4120ed81e3eb6af74a05", - "evidence_hash": "a3165d21e46b3ce553795daeae53e8f80e8e89c5cb228e68e6dcaff54bca5a89" + "evidence": "L311: while True: sha256:5b63313072aae9ca28677e03426513ccf12221e4f4e0ea6c31efbe09790633b5", + "evidence_hash": "05e1af469d651b51673763a7c4cdf759af9472fb627b7b470adc28cc237bd650" }, { "package": "openai", "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", - "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" + "evidence": "L3951: while True: sha256:d68ef896bf0743ca430cfacb9a3353da1f3b9c51c3a21b6450a07a32b55aa2ac", + "evidence_hash": "160eecdd79b521bffbe8476f782b69a0724c35d1b19376a7600807165fd54f9f" }, { "package": "openai", diff --git a/tests/python/test_torchcodec_torch_compat.py b/tests/python/test_torchcodec_torch_compat.py new file mode 100644 index 0000000000..6ad16a73f4 --- /dev/null +++ b/tests/python/test_torchcodec_torch_compat.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""torch / torchcodec ABI guardrails (unslothai/unsloth#7225).""" + +from __future__ import annotations + +import importlib.util +import re +import sys +import types +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[2] +PYPROJECT = REPO_ROOT / "pyproject.toml" +IMPORT_FIXES_PATH = REPO_ROOT / "unsloth" / "import_fixes.py" + + +def _load_import_fixes_module(): + spec = importlib.util.spec_from_file_location( + "unsloth_import_fixes_under_test", + IMPORT_FIXES_PATH, + ) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_pyproject_declares_torch210_audio_extra_with_python_gate(): + text = PYPROJECT.read_text(encoding = "utf-8") + assert "audio-torch210 = [" in text + assert "torchcodec>=0.10.0,<0.11.0" in text + assert "python_version >= '3.10'" in text + assert "audio-torch290 = [" in text + assert "audio-torch280 = [" in text + assert "\naudio = [" not in text + + +def _stub_torch(monkeypatch, version: str): + torch_mod = types.ModuleType("torch") + torch_mod.__version__ = version + monkeypatch.setitem(sys.modules, "torch", torch_mod) + + +def test_torch210_extras_bundle_audio_torch210(): + text = PYPROJECT.read_text(encoding = "utf-8") + for extra in ( + "cu128-torch2100", + "cu126-ampere-torch2100", + "rocm72-torch2100", + ): + match = re.search(rf"^{extra} = \[(.*?)^\]", text, re.MULTILINE | re.DOTALL) + assert match is not None, extra + assert "unsloth[audio-torch210]" in match.group(1) + + +def test_torchcodec_matrix_matches_notebook_validator(): + from scripts import notebook_validator as nv + fixes = _load_import_fixes_module() + assert fixes._TORCH_TORCHCODEC_MINORS == nv.TORCH_TORCHCODEC + + +def test_torchcodec_exclusive_upper_bound(): + fixes = _load_import_fixes_module() + assert fixes._torchcodec_exclusive_upper("0.10") == "<0.11.0" + assert fixes._torchcodec_exclusive_upper("0.9") == "<0.10.0" + + +def test_torch290_rejects_torchcodec_07(monkeypatch): + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.9.0+cu128") + monkeypatch.setattr(importlib.metadata, "version", lambda _name: "0.7.0") + + hint = fixes._torchcodec_version_mismatch_hint() + assert hint is not None + assert "audio-torch210" not in hint + + +def test_torch280_accepts_torchcodec_07(monkeypatch): + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.8.0+cu128") + monkeypatch.setattr(importlib.metadata, "version", lambda _name: "0.7.0") + + assert fixes._torchcodec_version_mismatch_hint() is None + + +def test_torch210_rejects_torchcodec_011(monkeypatch): + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.10.0+cu128") + monkeypatch.setattr( + importlib.metadata, + "version", + lambda _name: "0.11.0", + ) + + hint = fixes._torchcodec_version_mismatch_hint() + assert hint is not None + assert "torchcodec 0.11.0" in hint + assert "audio-torch210" in hint + assert "<0.11.0" in hint + assert "<11.0" not in hint + + +def test_torch210_accepts_torchcodec_010(monkeypatch): + import importlib.metadata + + fixes = _load_import_fixes_module() + _stub_torch(monkeypatch, "2.10.0+cu128") + monkeypatch.setattr( + importlib.metadata, + "version", + lambda _name: "0.10.0+cu128", + ) + + assert fixes._torchcodec_version_mismatch_hint() is None + + +def test_import_fixes_loads_on_python39_syntax(): + """Regression: module must import on 3.9 (postponed annotations for str | None).""" + fixes = _load_import_fixes_module() + assert callable(fixes._torchcodec_version_mismatch_hint) diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 5d54815705..9cd5e7243a 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + import os import importlib.abc import importlib.machinery @@ -1525,6 +1527,59 @@ def patch_torchcodec_audio_decoder(): pass +# torch.minor -> compatible torchcodec.minor strings (see notebook_validator.py). +_TORCH_TORCHCODEC_MINORS: dict[str, set[str]] = { + "2.10": {"0.10"}, + "2.9": {"0.8", "0.9"}, + "2.8": {"0.6", "0.7"}, + "2.7": {"0.3", "0.4", "0.5"}, + "2.6": {"0.2", "0.3"}, + "2.5": {"0.1", "0.2"}, +} + + +def _torchcodec_exclusive_upper(pin: str) -> str: + """Next torchcodec minor as an exclusive pip upper bound (0.10 -> <0.11.0).""" + major, minor = pin.split(".", 1) + return f"<{major}.{int(minor) + 1}.0" + + +def _torchcodec_version_mismatch_hint() -> str | None: + """Return a user-facing hint when installed torchcodec mismatches torch.""" + try: + import importlib.metadata as importlib_metadata + import torch + from packaging.version import Version + + torchcodec_version = importlib_metadata.version("torchcodec") + except Exception: + return None + + def _minor(version: str) -> str: + parts = Version(version.split("+", 1)[0]).release + return ".".join(str(p) for p in parts[:2]) + + try: + torch_minor = _minor(torch.__version__) + codec_minor = _minor(torchcodec_version) + except Exception: + # Non-PEP440 version strings must never break `import unsloth`. + return None + allowed = _TORCH_TORCHCODEC_MINORS.get(torch_minor) + if allowed is None or codec_minor in allowed: + return None + + pin = sorted(allowed)[-1] + upper = _torchcodec_exclusive_upper(pin) + install_hint = f"`pip install 'torchcodec>={pin},{upper}'`" + if torch_minor == "2.10": + install_hint += " or `pip install 'unsloth[audio-torch210]'`" + return ( + f"torchcodec {torchcodec_version} is incompatible with torch {torch.__version__}; " + f"install a matching build with {install_hint}." + ) + + def disable_torchcodec_if_broken(): """Make broken torchcodec behave as if uninstalled (#5446). @@ -1533,6 +1588,15 @@ def disable_torchcodec_if_broken(): flags and seat a sys.modules sentinel so downstream imports fall through their existing except ImportError handlers cleanly. """ + mismatch_hint = _torchcodec_version_mismatch_hint() + if mismatch_hint is not None: + try: + import warnings + warnings.warn(mismatch_hint, stacklevel = 2) + except Exception: + # Warning filters promoted to errors must not abort the disable + # fallback below (e.g. PYTHONWARNINGS=error, pytest -W error). + pass try: import importlib.util if importlib.util.find_spec("torchcodec") is None: From 09b6bf6c3937fd6f42dcfaf30d34237394a6dd5d Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:43:54 +0530 Subject: [PATCH 008/161] fix(studio): opt-in source-build GPU smoke validation (#7322) * fix(studio): opt-in source-build GPU smoke validation (#5854) Gap 1 (empty CUDA arch -> CPU) already landed in #6481. Wire gap 2: after a GPU source build, optionally run the same staged llama-server smoke test as the prebuilt path, then CPU-fallback on failure. Gated by UNSLOTH_LLAMA_STAGED_VALIDATION (default off) to avoid Blackwell JIT stalls. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(install): normalize staged validation env in setup.sh (#7322) Strip and lowercase UNSLOTH_LLAMA_STAGED_VALIDATION before the shell gate so values like True and surrounding whitespace match the Python staged_validation_enabled() helper. * Rebuild visual server after staged-validation CPU fallback (#5854) Mirror the primary source-build path by best-effort building llama-diffusion-gemma-visual-server after smoke-failure CPU fallback. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/studio-backend-ci.yml | 1 + studio/install_llama_prebuilt.py | 97 ++++++++++++++- studio/setup.sh | 63 ++++++++++ tests/run_all.sh | 1 + tests/sh/test_staged_validation_enabled.sh | 116 ++++++++++++++++++ .../test_install_llama_prebuilt_logic.py | 60 +++++++++ 6 files changed, 334 insertions(+), 4 deletions(-) create mode 100755 tests/sh/test_staged_validation_enabled.sh diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index b8f587b63e..d926d5c3e4 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -228,6 +228,7 @@ jobs: tests/sh/test_system_node_readonly.sh \ tests/sh/test_nvcc_meets_llama_minimum.sh \ tests/sh/test_resolve_cuda_archs.sh \ + tests/sh/test_staged_validation_enabled.sh \ tests/sh/test_tauri_install_exit_order.sh \ tests/sh/test_torch_constraint.sh \ tests/sh/test_torch_flavor.sh \ diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 6026962478..6ea850139e 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -186,8 +186,24 @@ VALIDATION_MODEL_CACHE_FILENAME = "stories260K.gguf" # in validate_prebuilt_choice. Disabled for now: the llama-server GPU forward pass # JIT-compiles CUDA kernels on first load and stalls every install and update by # minutes on Blackwell (sm_100). The check and the source-build fallback it triggers -# are kept intact -- set this to True to re-enable them. +# are kept intact -- set this to True, or set UNSLOTH_LLAMA_STAGED_VALIDATION=1, to +# re-enable them (#5854 gap 2). _RUN_STAGED_PREBUILT_VALIDATION = False + + +def staged_validation_enabled() -> bool: + """True when the expensive llama-server GPU smoke test should run. + + Default off (Blackwell CUDA JIT stalls installs). Opt in via the module + constant or ``UNSLOTH_LLAMA_STAGED_VALIDATION`` (1/true/yes/on). Used by both + the prebuilt path and setup.sh's source-build post-check (#5854). + """ + if _RUN_STAGED_PREBUILT_VALIDATION: + return True + raw = os.environ.get("UNSLOTH_LLAMA_STAGED_VALIDATION", "").strip().lower() + return raw in ("1", "true", "yes", "on") + + INSTALL_LOCK_TIMEOUT_SECONDS = 300 INSTALL_STAGING_ROOT_NAME = ".staging" GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"} @@ -5868,9 +5884,10 @@ def validate_prebuilt_choice( # so they are always validated. For an approved bundle the sha256 manifest # already proves integrity, so its runtime smoke test -- a cold CUDA-JIT pass # costing minutes on Blackwell sm_100 -- is gated behind - # _RUN_STAGED_PREBUILT_VALIDATION, disabled for now. The check and the - # source-build fallback it triggers are kept intact; flip the flag to restore it. - if choice.expected_sha256 is None or _RUN_STAGED_PREBUILT_VALIDATION: + # staged_validation_enabled() (constant or UNSLOTH_LLAMA_STAGED_VALIDATION), + # disabled for now. The check and the source-build fallback it triggers are + # kept intact; flip the flag / env to restore it (#5854). + if choice.expected_sha256 is None or staged_validation_enabled(): validate_quantize( quantize_path, probe_path, @@ -5891,6 +5908,49 @@ def validate_prebuilt_choice( return server_path, quantize_path +def validate_existing_install( + install_dir: Path, + *, + install_kind: str | None = None, + host: HostInfo | None = None, +) -> None: + """Run the staged smoke test against an already-built llama.cpp tree (#5854). + + Used by setup.sh after a GPU source build when ``UNSLOTH_LLAMA_STAGED_VALIDATION`` + is set. Raises ``PrebuiltFallback`` on failure so the caller can retry CPU. + """ + host = host or detect_host() + bin_dir = install_dir / "build" / "bin" + server_name = "llama-server.exe" if host.is_windows else "llama-server" + quantize_name = "llama-quantize.exe" if host.is_windows else "llama-quantize" + server_path = bin_dir / server_name + quantize_path = bin_dir / quantize_name + if not server_path.is_file(): + raise PrebuiltFallback(f"llama-server not found at {server_path}") + + with tempfile.TemporaryDirectory(prefix = "unsloth-llama-source-validate-") as tmp: + work_dir = Path(tmp) + probe_path = work_dir / "stories260K.gguf" + quantized_path = work_dir / "stories260K-q4.gguf" + download_validation_model(probe_path, validation_model_cache_path(install_dir)) + if quantize_path.is_file(): + validate_quantize( + quantize_path, + probe_path, + quantized_path, + install_dir, + host, + ) + validate_server( + server_path, + probe_path, + host, + install_dir, + install_kind = install_kind, + ) + log(f"staged source-build validation succeeded for {install_dir}") + + def validate_prebuilt_attempts( attempts: Iterable[AssetChoice], host: HostInfo, @@ -6345,6 +6405,24 @@ def parse_args() -> argparse.Namespace: "fork). Use --output-format json." ), ) + resolve_group.add_argument( + "--validate-install", + metavar = "DIR", + help = ( + "Run the staged llama-server smoke test against an existing build " + "tree (setup.sh source-build post-check, #5854). Exit 2 on failure. " + "Normally gated by UNSLOTH_LLAMA_STAGED_VALIDATION; this flag always " + "runs the check." + ), + ) + parser.add_argument( + "--install-kind", + default = None, + help = ( + "Install kind for --validate-install GPU offload (e.g. linux-cuda, " + "linux-rocm, macos-arm64). When omitted, host detection decides." + ), + ) parser.add_argument( "--output-format", choices = ("plain", "json"), @@ -6381,6 +6459,17 @@ def emit_resolver_output(payload: dict[str, Any], *, output_format: str) -> None def main() -> int: args = parse_args() + if args.validate_install is not None: + try: + validate_existing_install( + Path(args.validate_install), + install_kind = args.install_kind, + ) + except PrebuiltFallback as exc: + print(str(exc), file = sys.stderr) + raise SystemExit(EXIT_FALLBACK) from exc + return EXIT_SUCCESS + if args.resolve_llama_tag is not None: resolved = resolve_requested_llama_tag( args.resolve_llama_tag, diff --git a/studio/setup.sh b/studio/setup.sh index 3c97f77065..f6a6bc346b 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -251,6 +251,38 @@ _resolve_cuda_archs() { printf '%s' "$_archs" } +# Opt-in staged GPU smoke test after a source build (#5854 gap 2). Default off: +# llama-server's first GPU forward pass JIT-compiles CUDA kernels and stalls +# installs for minutes on Blackwell. Same env as install_llama_prebuilt.py. +_staged_validation_enabled() { + local _raw="${UNSLOTH_LLAMA_STAGED_VALIDATION:-}" + # Match install_llama_prebuilt.py staged_validation_enabled(): strip + lowercase. + _raw="$(printf '%s' "$_raw" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//' | tr '[:upper:]' '[:lower:]')" + case "$_raw" in + 1|true|yes|on) return 0 ;; + *) return 1 ;; + esac +} + +# Map the source-build GPU backend to install_llama_prebuilt --install-kind so +# validate_server enables --n-gpu-layers for the right backends. +_source_smoke_install_kind() { + if [ "${_TRY_METAL_CPU_FALLBACK:-false}" = true ]; then + printf '%s' "macos-arm64" + return 0 + fi + case "${GPU_BACKEND:-}" in + cuda) + case "$(uname -m 2>/dev/null || true)" in + aarch64|arm64) printf '%s' "linux-arm64-cuda" ;; + *) printf '%s' "linux-cuda" ;; + esac + ;; + rocm) printf '%s' "linux-rocm" ;; + *) printf '%s' "" ;; + esac +} + # Run a GPU probe under a 10s timeout when `timeout` is available so a wedged # NVIDIA driver cannot hang setup; fall back to a bare call where it is not. _setup_run_smi() { @@ -1900,6 +1932,37 @@ else run_quiet_no_exit "build diffusion visual server" cmake --build "$_BUILD_TMP/build" --config Release --target llama-diffusion-gemma-visual-server -j"$NCPU" || true fi + # Opt-in post-build GPU smoke test (#5854 gap 2). Default off (Blackwell + # CUDA JIT stalls). On failure, reuse the CPU fallback path so the user + # still gets a working llama-server. Runs before the install swap. + if [ "$BUILD_OK" = true ] && _staged_validation_enabled; then + _FB_LABEL="$(_gpu_fallback_label)" + _SMOKE_KIND="$(_source_smoke_install_kind)" + if [ -n "$_FB_LABEL" ]; then + _SMOKE_CMD=( + python "$SCRIPT_DIR/install_llama_prebuilt.py" + --validate-install "$_BUILD_TMP" + ) + [ -n "$_SMOKE_KIND" ] && _SMOKE_CMD+=(--install-kind "$_SMOKE_KIND") + if ! run_quiet_no_exit "validate source llama.cpp" "${_SMOKE_CMD[@]}"; then + substep "$_FB_LABEL source build failed smoke test; retrying CPU build..." "$C_WARN" + _TRY_METAL_CPU_FALLBACK=false + rm -rf "$_BUILD_TMP/build" + if run_quiet_no_exit "cmake llama.cpp (cpu fallback)" cmake $CMAKE_GENERATOR_ARGS -S "$_BUILD_TMP" -B "$_BUILD_TMP/build" $CPU_FALLBACK_CMAKE_ARGS; then + _BUILD_DESC="building (CPU fallback after $_FB_LABEL smoke failed)" + GPU_BACKEND="" + run_quiet_no_exit "build llama-server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false + if [ "$BUILD_OK" = true ]; then + run_quiet_no_exit "build llama-quantize (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-quantize -j"$NCPU" || true + run_quiet_no_exit "build diffusion visual server (cpu fallback)" cmake --build "$_BUILD_TMP/build" --config Release --target llama-diffusion-gemma-visual-server -j"$NCPU" || true + fi + else + BUILD_OK=false + fi + fi + fi + fi + # Swap only after build succeeds -- preserves existing install on failure if [ "$BUILD_OK" = true ]; then _assert_studio_owned_or_absent "$LLAMA_CPP_DIR" "llama.cpp install" diff --git a/tests/run_all.sh b/tests/run_all.sh index eaa726f73c..6eccffc75f 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -12,6 +12,7 @@ sh "$TESTS_DIR/sh/test_mac_intel_compat.sh" sh "$TESTS_DIR/sh/test_torch_constraint.sh" sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh" sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh" +sh "$TESTS_DIR/sh/test_staged_validation_enabled.sh" sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh" sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh" sh "$TESTS_DIR/sh/test_torch_flavor.sh" diff --git a/tests/sh/test_staged_validation_enabled.sh b/tests/sh/test_staged_validation_enabled.sh new file mode 100755 index 0000000000..da6a0bdd27 --- /dev/null +++ b/tests/sh/test_staged_validation_enabled.sh @@ -0,0 +1,116 @@ +#!/bin/bash +# 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 setup.sh staged-validation helpers (#5854 gap 2). +# Opt-in GPU smoke after a source build; default off (Blackwell JIT stall). +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh" +PASS=0 +FAIL=0 + +_FUNC_FILE=$(mktemp) +{ + sed -n '/^_staged_validation_enabled()/,/^}/p' "$SETUP_SH" + sed -n '/^_source_smoke_install_kind()/,/^}/p' "$SETUP_SH" +} > "$_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_FUNC_FILE" +rm -f "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1)) + fi +} + +assert_rc() { + _label="$1"; _expected="$2" + shift 2 + set +e + "$@" >/dev/null 2>&1 + _rc=$? + set -e + if [ "$_rc" -eq "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected rc $_expected, got $_rc)"; FAIL=$((FAIL + 1)) + fi +} + +echo "=== _staged_validation_enabled ===" +unset UNSLOTH_LLAMA_STAGED_VALIDATION +assert_rc "default off" 1 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=0 +assert_rc "0 is off" 1 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=1 +assert_rc "1 is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=true +assert_rc "true is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=yes +assert_rc "yes is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=on +assert_rc "on is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=True +assert_rc "True is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=' yes ' +assert_rc "whitespace yes is on" 0 _staged_validation_enabled + +UNSLOTH_LLAMA_STAGED_VALIDATION=maybe +assert_rc "maybe is off" 1 _staged_validation_enabled +unset UNSLOTH_LLAMA_STAGED_VALIDATION + +echo "=== _source_smoke_install_kind ===" +_TRY_METAL_CPU_FALLBACK=true +GPU_BACKEND="" +assert_eq "metal" "macos-arm64" "$(_source_smoke_install_kind)" + +_TRY_METAL_CPU_FALLBACK=false +GPU_BACKEND=cuda +_kind="$(_source_smoke_install_kind)" +case "$(uname -m)" in + aarch64|arm64) assert_eq "cuda arm" "linux-arm64-cuda" "$_kind" ;; + *) assert_eq "cuda x86" "linux-cuda" "$_kind" ;; +esac + +GPU_BACKEND=rocm +assert_eq "rocm" "linux-rocm" "$(_source_smoke_install_kind)" + +GPU_BACKEND="" +assert_eq "cpu empty" "" "$(_source_smoke_install_kind)" + +echo "=== setup.sh source smoke contract ===" +assert_contains() { + _label="$1"; _hay="$2"; _needle="$3" + case "$_hay" in + *"$_needle"*) echo " PASS: $_label"; PASS=$((PASS + 1)) ;; + *) echo " FAIL: $_label (missing '$_needle')"; FAIL=$((FAIL + 1)) ;; + esac +} +_src=$(cat "$SETUP_SH") +assert_contains "env gate present" "$_src" "UNSLOTH_LLAMA_STAGED_VALIDATION" +assert_contains "calls validate-install" "$_src" "--validate-install" +assert_contains "smoke fail retries CPU" "$_src" "source build failed smoke test; retrying CPU build" +# Smoke must run before the install swap. +_smoke_pos=$(printf '%s' "$_src" | awk '/validate source llama.cpp/{print NR; exit}') +_swap_pos=$(printf '%s' "$_src" | awk '/mv "\$_BUILD_TMP" "\$LLAMA_CPP_DIR"/{print NR; exit}') +if [ -n "$_smoke_pos" ] && [ -n "$_swap_pos" ] && [ "$_smoke_pos" -lt "$_swap_pos" ]; then + echo " PASS: smoke before install swap"; PASS=$((PASS + 1)) +else + echo " FAIL: smoke before install swap (smoke=$_smoke_pos swap=$_swap_pos)"; FAIL=$((FAIL + 1)) +fi + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] diff --git a/tests/studio/install/test_install_llama_prebuilt_logic.py b/tests/studio/install/test_install_llama_prebuilt_logic.py index 0964c047d5..3eaf56d15c 100644 --- a/tests/studio/install/test_install_llama_prebuilt_logic.py +++ b/tests/studio/install/test_install_llama_prebuilt_logic.py @@ -3324,6 +3324,66 @@ def test_validate_prebuilt_choice_approved_validation_runs_when_flag_enabled(tmp assert calls == {"quantize": 1, "server": 1} +def test_staged_validation_enabled_default_off(monkeypatch): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", False) + monkeypatch.delenv("UNSLOTH_LLAMA_STAGED_VALIDATION", raising = False) + assert INSTALL_LLAMA_PREBUILT.staged_validation_enabled() is False + + +@pytest.mark.parametrize("value", ["1", "true", "YES", "on"]) +def test_staged_validation_enabled_env_opt_in(monkeypatch, value): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", False) + monkeypatch.setenv("UNSLOTH_LLAMA_STAGED_VALIDATION", value) + assert INSTALL_LLAMA_PREBUILT.staged_validation_enabled() is True + + +def test_validate_prebuilt_choice_approved_validation_runs_when_env_enabled(tmp_path, monkeypatch): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "_RUN_STAGED_PREBUILT_VALIDATION", False) + monkeypatch.setenv("UNSLOTH_LLAMA_STAGED_VALIDATION", "1") + calls = _run_validate_prebuilt_choice(monkeypatch, tmp_path, expected_sha256 = "ab" * 32) + assert calls == {"quantize": 1, "server": 1} + + +def test_validate_existing_install_runs_server_smoke(tmp_path, monkeypatch): + # setup.sh --validate-install path: exercise smoke helpers without a real GPU. + install_dir = tmp_path / "llama.cpp" + bin_dir = install_dir / "build" / "bin" + bin_dir.mkdir(parents = True) + (bin_dir / "llama-server").write_text("#!/bin/sh\n", encoding = "utf-8") + (bin_dir / "llama-quantize").write_text("#!/bin/sh\n", encoding = "utf-8") + calls: dict[str, int] = {"quantize": 0, "server": 0, "download": 0} + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "download_validation_model", + lambda path, cache = None: calls.__setitem__("download", calls["download"] + 1), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "validate_quantize", + lambda *a, **k: calls.__setitem__("quantize", calls["quantize"] + 1), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "validate_server", + lambda *a, **k: calls.__setitem__("server", calls["server"] + 1), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "detect_host", + lambda: linux_host(), + ) + + INSTALL_LLAMA_PREBUILT.validate_existing_install(install_dir, install_kind = "linux-cuda") + assert calls == {"quantize": 1, "server": 1, "download": 1} + + +def test_validate_existing_install_missing_server_raises(tmp_path, monkeypatch): + monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "detect_host", lambda: linux_host()) + with pytest.raises(INSTALL_LLAMA_PREBUILT.PrebuiltFallback, match = "llama-server not found"): + INSTALL_LLAMA_PREBUILT.validate_existing_install(tmp_path / "missing") + + def test_diffusion_visual_server_uses_approved_checksum_download(monkeypatch, tmp_path: Path): asset_name = "llama-diffusion-gemma-visual-server-linux-x64" expected_sha = "a" * 64 From f5a0c2226b73f7225e7d9759c94159b081584695 Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:45:01 +0530 Subject: [PATCH 009/161] fix(studio): resolve bare git on Windows sandbox PATH (#7323) * fix(studio): resolve bare git on Windows sandbox PATH Sandboxed terminal tools rebuilt PATH as venv + System32 only, so user-installed Git under Program Files never resolved by bare name. Append absolute host PATH dirs after the curated prefix and inherit PATHEXT on Windows (#7317). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): restrict sandbox PATH inheritance to Windows Git dirs (#7323) Only append Git-for-Windows install directories from the host PATH on Windows, instead of every absolute entry. This fixes bare `git` resolution (#7317) without letting user-writable dirs (venv, node_modules/.bin) shadow auto-safe terminal commands. * Pin sandbox PATHEXT to block cwd script hijacks (#7317) Use a fixed .EXE;.COM list instead of inheriting the host PATHEXT so cmd cannot resolve auto-approved bare names from workdir .BAT/.CMD stubs. * Resolve sandbox git dir via shutil.which and disable cwd exe lookup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep non-exe git launchers resolvable under restricted PATHEXT * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restrict inherited sandbox git dir to system install roots * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop SystemRoot trust and canonicalize short paths for sandbox git * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve Program Files via known-folder API and append canonical git dir * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trust native Program Files on 32-bit Windows and stub program roots in tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Scan PATH for a trusted git and derive native Program Files root * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop ProgramFiles env from the trusted-root fallback * Fail closed when trusted Program Files root cannot be resolved --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/tools.py | 138 ++++++++++++- studio/backend/tests/test_sandbox_tools.py | 216 +++++++++++++++++++++ 2 files changed, 352 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index bc9ffe85c2..0ef6dd46cf 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -18,6 +18,7 @@ import queue import random import re import shlex +import shutil import ssl import subprocess import sys @@ -328,6 +329,7 @@ def _find_blocked_commands(command: str) -> set[str]: # Directory holding the sandbox ``sitecustomize.py`` shim (code-interpreter # path remap); placed on the sandboxed child's PYTHONPATH in _build_safe_env. _SANDBOX_SITE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sandbox_site") + # ── "Approve for me" (permission_mode="auto") safety detection ────────────── # Auto mode pauses only calls classified here as potentially unsafe. The sandbox # and hard blocks (blocklist, rlimits) still apply at run time; this gate only @@ -2491,15 +2493,124 @@ def is_potentially_unsafe_tool_call(name: str, arguments: dict) -> bool: return True +def _canon_win_path(p: str) -> str: + """Canonical form for trust comparison: realpath (expands 8.3 aliases and + resolves junctions/symlinks) + normcase/normpath.""" + return os.path.normcase(os.path.normpath(os.path.realpath(p))) + + +def _augment_native_program_roots(roots: list[str]) -> list[str]: + """Add the native Program Files sibling for any x86 root by stripping the + `` (x86)`` suffix, so a 32-bit process (whose known-folder ids map only to + the x86 root) still trusts a 64-bit Git install.""" + out = list(roots) + for root in roots: + base = root.rstrip("\\/") + if base.lower().endswith(" (x86)"): + native = base[: -len(" (x86)")] + if native and native not in out: + out.append(native) + return out + + +def _windows_program_roots() -> list[str]: + """Program Files install roots, resolved ONLY from the Windows known-folder + API (SHGetKnownFolderPath). Fails closed (returns ``[]``) if the API is + unavailable: env vars (%ProgramFiles%, even %SystemDrive%) are caller- + overrideable and could relocate the trust boundary, so we never derive a + trusted root from them. On any real Windows host shell32 is present, so + this only returns empty in a broken/non-Windows environment where the + sandbox git-PATH feature is not needed anyway (#7317). + """ + roots: list[str] = [] + try: + import ctypes + from ctypes import wintypes + + # FOLDERID_ProgramFiles, _ProgramFilesX86, _ProgramFilesX64. The X64 + # id (Win10 1703+) yields the native root even from a 32-bit process, + # where the first two both map to Program Files (x86). + folder_ids = ( + "{905e63b6-c1bf-494e-b29c-65b732d3d21a}", + "{7C5A40EF-A0FB-4BFC-874A-C0F2E0B9FA8E}", + "{6D809377-6AF0-444b-8957-A3773F02200E}", + ) + _SHGet = ctypes.windll.shell32.SHGetKnownFolderPath + _CoTaskMemFree = ctypes.windll.ole32.CoTaskMemFree + for fid in folder_ids: + guid = ctypes.create_string_buffer(16) + ctypes.windll.ole32.CLSIDFromString(wintypes.LPCWSTR(fid), ctypes.byref(guid)) + ptr = ctypes.c_wchar_p() + if _SHGet(ctypes.byref(guid), 0, None, ctypes.byref(ptr)) == 0: + if ptr.value: + roots.append(ptr.value) + _CoTaskMemFree(ptr) + except Exception: + return [] + return _augment_native_program_roots(roots) + + +def _resolve_trusted_windows_git() -> tuple[str, str]: + """Find a git launcher in a TRUSTED Program Files dir. Returns + ``(canonical_dir, ext)`` or ``("", "")``. + + ``shutil.which`` returns only the first PATH match, which may be an + untrusted user shim; scan the remaining PATH entries for a later trusted + Git so bare ``git`` still resolves (#7317). + """ + exts = [e for e in (os.environ.get("PATHEXT") or ".EXE;.CMD;.BAT;.COM").split(os.pathsep)] + candidates: list[str] = [] + primary = shutil.which("git") + if primary: + candidates.append(primary) + for entry in (os.environ.get("PATH") or "").split(os.pathsep): + entry = entry.strip().strip('"') + if not entry or not os.path.isabs(entry): + continue + for ext in exts: + cand = os.path.join(entry, "git" + ext) + if os.path.isfile(cand): + candidates.append(cand) + for git_exe in candidates: + git_dir = os.path.dirname(git_exe) + if os.path.isabs(git_dir) and _is_trusted_windows_program_dir(git_dir): + return os.path.realpath(git_dir), os.path.splitext(git_exe)[1].upper() + return "", "" + + +def _is_trusted_windows_program_dir(path: str) -> bool: + """True when ``path`` sits under a system-managed Program Files root. + + Only the Program Files roots are trusted (admin-writable only), resolved + via the known-folder API so an overridden env var cannot relocate them, + never ``%SystemRoot%`` (Git does not install there and it holds + world-writable subdirs like ``Windows\\Temp``). Per-user managers + (Scoop/Choco shims under the profile) are refused. Paths are canonicalized + so 8.3 aliases and junctions still resolve to their real root (#7317). + """ + norm = _canon_win_path(path) + for root in _windows_program_roots(): + root_norm = _canon_win_path(root) + if norm == root_norm or norm.startswith(root_norm + os.sep): + return True + return False + + def _build_safe_env(workdir: str) -> dict[str, str]: """Build a minimal, credential-free environment for sandboxed subprocesses. Whitelist-built from scratch (parent env NOT inherited): only PATH/HOME/ TMPDIR/LANG/TERM/PYTHONIOENCODING/PYTHONPATH (+VIRTUAL_ENV or Windows - SystemRoot) reach the child; all credential vars (HF_TOKEN, AWS_*, etc.) - are absent. HOME points at the sandbox workdir so SDKs can't read the + SystemRoot and a minimal PATHEXT) reach the child; all credential vars + (HF_TOKEN, AWS_*, etc.) are absent. HOME points at the sandbox workdir so SDKs can't read the operator's cached creds. PYTHONPATH carries only the sandbox sitecustomize shim directory. + + PATH starts with the Studio interpreter / venv and OS system dirs so + ``python``/``pip`` stay pinned. On Windows only, Git-for-Windows install + dirs from the host PATH are appended so bare ``git`` resolves (#7317). + User-writable host PATH entries (venv, ``node_modules/.bin``, etc.) are + never inherited — they could shadow auto-safe terminal commands. """ # Start from the running interpreter's dir so 'python'/'pip' resolve to the # same environment the Unsloth server runs in. @@ -2519,6 +2630,20 @@ def _build_safe_env(workdir: str) -> dict[str, str]: else: path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"]) + # Windows Git installs live outside System32; inherit the dir of the git + # the HOST shell resolves, but ONLY when it sits under a system install + # root (Program Files, windir). A user-writable dir (Scoop/Choco shims) + # is refused: it would let an attacker drop rg.exe/jq.exe beside git and + # have an auto-approved bare command execute it (#7317). + git_ext = "" + if sys.platform == "win32": + # Append the CANONICAL (realpath) trusted git dir, scanning past any + # untrusted user shim that sorts first on PATH; the canonical path + # cannot be retargeted via a junction after the trust check. + _trusted_git_dir, git_ext = _resolve_trusted_windows_git() + if _trusted_git_dir: + path_entries.append(_trusted_git_dir) + # Deduplicate, preserving order. deduped = list(dict.fromkeys(p for p in path_entries if p)) @@ -2538,6 +2663,15 @@ def _build_safe_env(workdir: str) -> dict[str, str]: # Windows needs SystemRoot for Python/subprocess to work. if sys.platform == "win32": env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows") + # Restrict PATHEXT so cwd .BAT/.CMD cannot hijack bare names (#7317). + pathext = ".EXE;.COM" + if git_ext and git_ext not in (".EXE", ".COM"): + # Keep the host git launcher (e.g. a .CMD shim) resolvable. + pathext += ";" + git_ext + env["PATHEXT"] = pathext + # cmd/CreateProcess search cwd before PATH for bare names; disable so + # a workdir rg.exe/git.exe cannot shadow auto-approved commands. + env["NoDefaultCurrentDirectoryInExePath"] = "1" return env diff --git a/studio/backend/tests/test_sandbox_tools.py b/studio/backend/tests/test_sandbox_tools.py index 2970b1a6bb..64201477e3 100644 --- a/studio/backend/tests/test_sandbox_tools.py +++ b/studio/backend/tests/test_sandbox_tools.py @@ -297,6 +297,8 @@ class TestSandboxEnvIsolation: "PYTHONPATH", "VIRTUAL_ENV", "SystemRoot", + "PATHEXT", # Windows only; minimal list so cwd scripts cannot hijack + "NoDefaultCurrentDirectoryInExePath", # Windows only; no cwd-first lookup } extras = set(env.keys()) - allowed assert not extras, f"sandbox env added unexpected keys: {extras}" @@ -305,6 +307,220 @@ class TestSandboxEnvIsolation: assert env["PYTHONPATH"].endswith("sandbox_site") assert "leak-me" not in env["PYTHONPATH"] + def test_host_git_dir_appended_after_curated(self, monkeypatch, tmp_path): + # #7317: Windows Git lives under Program Files, not System32. Sandbox + # PATH resolves bare `git` by appending the dir of the git the HOST + # shell resolves (shutil.which), after the curated prefix. + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + git_dir = prog / "Git" / "cmd" + git_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_dir / "git.exe")) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(git_dir) in parts + # Curated prefix stays ahead of host Git so Studio python/pip win. + assert parts.index(str(git_dir)) > 0 + + def test_host_path_dirs_not_inherited(self, monkeypatch, tmp_path): + """Host PATH dirs (user-writable, git-lookalike) are never inherited; + only the resolved git dir is. No git resolved -> nothing appended.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + venv_scripts = tmp_path / "venv" / "Scripts" + venv_scripts.mkdir(parents = True) + fake_git = tmp_path / "scratch" / "Git" / "cmd" + fake_git.mkdir(parents = True) + monkeypatch.setenv( + "PATH", + os.pathsep.join([str(venv_scripts), str(fake_git), os.environ.get("PATH", "")]), + ) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: None) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(venv_scripts) not in parts + # A git-suffixed but unresolved (user-writable) dir is NOT trusted. + assert str(fake_git) not in parts + + def test_git_cmd_shim_extension_added_to_pathext(self, monkeypatch, tmp_path): + """A host git resolved as a .cmd shim under a trusted root stays + resolvable under the restricted PATHEXT (cwd lookup disabled).""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + git_dir = prog / "Git" / "cmd" + git_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_dir / "git.cmd")) + env = _build_safe_env(str(tmp_path)) + assert str(git_dir) in env["PATH"].split(os.pathsep) + assert env["PATHEXT"] == ".EXE;.COM;.CMD" + + def test_user_writable_git_dir_refused(self, monkeypatch, tmp_path): + """Git resolved from a per-user manager (Scoop shims) is NOT trusted: + an attacker could drop rg.exe beside it and hit the auto-approve gate.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr( + tools_mod, "_windows_program_roots", lambda: [str(tmp_path / "Program Files")] + ) + shim_dir = tmp_path / "users" / "alice" / "scoop" / "shims" + shim_dir.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(shim_dir / "git.exe")) + env = _build_safe_env(str(tmp_path)) + assert str(shim_dir) not in env["PATH"].split(os.pathsep) + # No trusted git launcher -> PATHEXT stays minimal. + assert env["PATHEXT"] == ".EXE;.COM" + + def test_trust_uses_known_folder_not_env_override(self, monkeypatch, tmp_path): + """Trust is driven by the resolved Program Files roots, so a git under + an attacker-overridden %ProgramFiles% env value is still refused.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "RealProgramFiles" + (real_prog).mkdir() + evil = tmp_path / "attacker" + (evil / "Git" / "cmd").mkdir(parents = True) + # Resolver returns the genuine root; env is overridden to the evil dir. + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + monkeypatch.setenv("ProgramFiles", str(evil)) + monkeypatch.setattr( + tools_mod.shutil, "which", lambda name: str(evil / "Git" / "cmd" / "git.exe") + ) + env = _build_safe_env(str(tmp_path)) + assert str(evil / "Git" / "cmd") not in env["PATH"].split(os.pathsep) + + def test_canonical_git_dir_appended(self, monkeypatch, tmp_path): + """The PATH entry is the realpath of the trusted dir, not a junction + alias, so it cannot be retargeted after the trust check.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "Program Files" + real_git = real_prog / "Git" / "cmd" + real_git.mkdir(parents = True) + link = tmp_path / "link" + try: + link.symlink_to(real_prog, target_is_directory = True) + except (OSError, NotImplementedError): + pytest.skip("symlink unsupported in this environment") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + monkeypatch.setattr( + tools_mod.shutil, + "which", + lambda name: str(link / "Git" / "cmd" / "git.exe"), + ) + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(real_git) in parts # canonical, not the `link/...` alias + + def test_windows_temp_git_dir_refused(self, monkeypatch, tmp_path): + """A git under a world-writable %SystemRoot% subdir (Windows\\Temp) is + NOT trusted, even though it sits under the Windows root.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr( + tools_mod, "_windows_program_roots", lambda: [str(tmp_path / "Program Files")] + ) + temp_git = tmp_path / "Windows" / "Temp" / "Git" / "cmd" + temp_git.mkdir(parents = True) + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(temp_git / "git.exe")) + env = _build_safe_env(str(tmp_path)) + assert str(temp_git) not in env["PATH"].split(os.pathsep) + + def test_trusted_program_dir_matches_via_realpath(self, monkeypatch, tmp_path): + """The trust check canonicalizes paths, so a symlinked/short alias of + Program Files still matches (stand-in for 8.3 PROGRA~1 on Windows).""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + real_prog = tmp_path / "Program Files" + (real_prog / "Git" / "cmd").mkdir(parents = True) + alias = tmp_path / "PROGRA~1" + try: + alias.symlink_to(real_prog, target_is_directory = True) + except (OSError, NotImplementedError): + pytest.skip("symlink unsupported in this environment") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(real_prog)]) + git_via_alias = alias / "Git" / "cmd" / "git.exe" + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(git_via_alias)) + env = _build_safe_env(str(tmp_path)) + parts = [os.path.normcase(os.path.realpath(p)) for p in env["PATH"].split(os.pathsep)] + assert os.path.normcase(str(real_prog / "Git" / "cmd")) in parts + + def test_scan_past_untrusted_git_shim(self, monkeypatch, tmp_path): + """When an untrusted shim sorts first on PATH, the scan still finds a + later trusted Program Files git.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + prog = tmp_path / "Program Files" + trusted_git = prog / "Git" / "cmd" + trusted_git.mkdir(parents = True) + (trusted_git / "git.EXE").write_text("") # match PATHEXT case on this FS + shim = tmp_path / "scoop" / "shims" + shim.mkdir(parents = True) + (shim / "git.EXE").write_text("") + monkeypatch.setattr(tools_mod, "_windows_program_roots", lambda: [str(prog)]) + # shutil.which returns the untrusted shim first. + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: str(shim / "git.EXE")) + monkeypatch.setenv("PATH", os.pathsep.join([str(shim), str(trusted_git)])) + monkeypatch.setenv("PATHEXT", ".EXE") + env = _build_safe_env(str(tmp_path)) + parts = env["PATH"].split(os.pathsep) + assert str(trusted_git) in parts + assert str(shim) not in parts + + def test_program_roots_fails_closed_without_known_folder_api(self, monkeypatch): + """When the known-folder API is unavailable, no roots are trusted: env + vars (even %SystemDrive%) are caller-overrideable, so we never derive a + trusted root from them.""" + import core.inference.tools as tools_mod + + # ctypes fails on this Linux host, so the API path raises and we fail + # closed. Any attacker override of these env vars must be irrelevant. + monkeypatch.setenv("ProgramFiles", r"D:\attacker-writable") + monkeypatch.setenv("ProgramW6432", r"D:\attacker-writable") + monkeypatch.setenv("SystemDrive", "D:") + assert tools_mod._windows_program_roots() == [] + + def test_augment_native_program_roots_derives_native_sibling(self): + """A 32-bit process only sees the x86 root; the native sibling is + derived by stripping the ` (x86)` suffix.""" + import core.inference.tools as tools_mod + + roots = tools_mod._augment_native_program_roots([r"C:\Program Files (x86)"]) + lowered = [r.lower() for r in roots] + assert r"c:\program files (x86)" in lowered + assert r"c:\program files" in lowered + + def test_no_default_current_directory_in_exe_path_set_on_windows(self, monkeypatch, tmp_path): + """cmd/CreateProcess must not search cwd for bare names in the sandbox.""" + import core.inference.tools as tools_mod + from core.inference.tools import _build_safe_env + + monkeypatch.setattr(sys, "platform", "win32") + monkeypatch.setattr(tools_mod.shutil, "which", lambda name: None) + env = _build_safe_env(str(tmp_path)) + assert env["NoDefaultCurrentDirectoryInExePath"] == "1" + def test_home_points_at_sandbox_workdir(self, tmp_path): from core.inference.tools import _build_safe_env From 0807d03ed00541ab1cae6efc7e9bc82edd538abb Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:46:00 +0530 Subject: [PATCH 010/161] fix(install): show detected distro in sudo apt Accept prompt (#7324) * fix(install): show detected distro in sudo apt Accept prompt Make the package-install elevation prompt name the detected distro and state that packages come from official apt repos, so users know we are not installing a tarball outside their package manager (#6207). * fix(install): avoid case/;; inside $() for bash 3.2 macOS CI uses bash 3.2, which misparses case arms inside command substitution and fails install.sh at the apt distro helper. Use a plain subshell so the Accept? prompt still works everywhere. --- install.sh | 35 ++++++++++- tests/sh/test_apt_distro_prompt.sh | 94 ++++++++++++++++++++++++++++++ 2 files changed, 128 insertions(+), 1 deletion(-) create mode 100755 tests/sh/test_apt_distro_prompt.sh diff --git a/install.sh b/install.sh index 84c5be9742..dface28918 100755 --- a/install.sh +++ b/install.sh @@ -625,6 +625,36 @@ _is_pkg_installed() { esac } +# ── Helper: human-readable apt distro label for the sudo package prompt (#6207) ── +# Reads /etc/os-release so the Accept? prompt can say which distro we detected and +# that packages come from that distro's official apt repos (not a tarball). +_apt_distro_description() { + # Plain ( ... ) subshell — not $() — so case/;; stays bash-3.2-safe on macOS. + # Bash 3.2 misparses case arms inside command substitution and errors on `;;`. + ( + if [ ! -r /etc/os-release ]; then + printf 'a debian-like system' + exit 0 + fi + # shellcheck disable=SC1091 + . /etc/os-release 2>/dev/null || true + if [ -n "${NAME:-}" ] && [ -n "${VERSION_ID:-}" ]; then + _ad_label="$NAME $VERSION_ID" + elif [ -n "${PRETTY_NAME:-}" ]; then + _ad_label="$PRETTY_NAME" + elif [ -n "${NAME:-}" ]; then + _ad_label="$NAME" + else + printf 'a debian-like system' + exit 0 + fi + case " ${ID:-} ${ID_LIKE:-} " in + *" debian "*|*" ubuntu "*) _ad_label="${_ad_label} (debian-like)" ;; + esac + printf '%s' "$_ad_label" + ) +} + # ── Helper: install packages via apt, escalating to sudo only if needed ── # Usage: _smart_apt_install pkg1 pkg2 pkg3 ... _smart_apt_install() { @@ -655,11 +685,14 @@ _smart_apt_install() { # Step 3: Escalate -- need elevated permissions for remaining packages if command -v sudo >/dev/null 2>&1; then + _ad_desc="$(_apt_distro_description)" echo "" echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo " WARNING: We require sudo elevated permissions to install:" echo " $_STILL_MISSING" - echo " If you accept, we'll run sudo now, and it'll prompt your password." + echo " Detected ${_ad_desc}." + echo " If you accept, we'll run sudo apt-get to install these packages" + echo " from your distro's official repositories (not a third-party tarball)." echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" echo "" printf " Accept? [Y/n] " diff --git a/tests/sh/test_apt_distro_prompt.sh b/tests/sh/test_apt_distro_prompt.sh new file mode 100755 index 0000000000..19601b0065 --- /dev/null +++ b/tests/sh/test_apt_distro_prompt.sh @@ -0,0 +1,94 @@ +#!/bin/bash +# 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 install.sh's _apt_distro_description helper (#6207). +# The sudo Accept? prompt should name the detected distro and say packages come +# from official apt repos. Hermetic: extract the helper and rewrite +# /etc/os-release to per-test fixtures (same pattern as test_strixhalo_wsl_reroute.sh). +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +_TMP_ROOT=$(mktemp -d) +trap 'rm -rf "$_TMP_ROOT"' EXIT + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1)) + fi +} + +assert_contains() { + _label="$1"; _hay="$2"; _needle="$3" + case "$_hay" in + *"$_needle"*) echo " PASS: $_label"; PASS=$((PASS + 1)) ;; + *) echo " FAIL: $_label (missing '$_needle' in: $_hay)"; FAIL=$((FAIL + 1)) ;; + esac +} + +# Extract helper with /etc/os-release rewritten to $1. +build_func() { + _fix="$1" + _f=$(mktemp -p "$_TMP_ROOT") + sed -n '/^_apt_distro_description()/,/^}/p' "$INSTALL_SH" \ + | sed -e "s#/etc/os-release#$_fix/os-release#g" \ + > "$_f" + echo "$_f" +} + +run_desc() { + _os="$1" + _d=$(mktemp -d -p "$_TMP_ROOT") + printf '%s\n' "$_os" > "$_d/os-release" + _f=$(build_func "$_d") + # shellcheck disable=SC1090 + . "$_f" + _apt_distro_description +} + +echo "=== _apt_distro_description ===" + +assert_eq "ubuntu name+version debian-like" \ + "Ubuntu 24.04 (debian-like)" \ + "$(run_desc "$(printf 'NAME=\"Ubuntu\"\nVERSION_ID=\"24.04\"\nID=ubuntu\nID_LIKE=debian\n')")" + +assert_eq "debian name+version debian-like" \ + "Debian GNU/Linux 12 (debian-like)" \ + "$(run_desc "$(printf 'NAME=\"Debian GNU/Linux\"\nVERSION_ID=\"12\"\nID=debian\n')")" + +assert_eq "pretty_name fallback when name/version missing" \ + "Linux Mint 22 (debian-like)" \ + "$(run_desc "$(printf 'PRETTY_NAME=\"Linux Mint 22\"\nID=linuxmint\nID_LIKE=\"ubuntu debian\"\n')")" + +# NAME alone (no VERSION_ID) — still prefer NAME over PRETTY_NAME. +assert_eq "name only" \ + "Pop!_OS (debian-like)" \ + "$(run_desc "$(printf 'NAME=\"Pop!_OS\"\nID=pop\nID_LIKE=\"ubuntu debian\"\n')")" + +assert_eq "missing os-release file" \ + "a debian-like system" \ + "$( + _d=$(mktemp -d -p "$_TMP_ROOT") + _f=$(build_func "$_d") + # shellcheck disable=SC1090 + . "$_f" + _apt_distro_description + )" + +echo "=== _smart_apt_install prompt contract ===" +_smart=$(sed -n '/^_smart_apt_install()/,/^}/p' "$INSTALL_SH") +assert_contains "calls distro helper" "$_smart" '_apt_distro_description' +assert_contains "names detected distro" "$_smart" 'Detected ${_ad_desc}' +assert_contains "mentions apt-get" "$_smart" 'sudo apt-get' +assert_contains "mentions official repos" "$_smart" "official repositories" +assert_contains "rejects tarball worry" "$_smart" "not a third-party tarball" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] From c2114d64dd962f94dd54ccf27376e37cb112ec39 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 20:06:30 -0700 Subject: [PATCH 011/161] Studio: fail closed on index-referenced nested pickle shards in the offline embedding gate (#7366) * Studio: fail closed on index-referenced nested pickle shards in the offline embedding gate The offline embedding security gate (HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE) only scanned the direct files of each SentenceTransformer load root and never parsed local weight indexes, so a cached snapshot whose pytorch_model.bin.index.json maps a weight to a nested shard (e.g. shards/pytorch_model-00001-of-00001.bin) was treated as inert and allowed. The loader then follows the index into the subdir and unpickles the shard. The online gate already blocks index-referenced subdir pickles, so the offline path was strictly weaker. Parse each local weight index in a load root and follow weight_map into nested dirs, flagging any referenced pickle-extension shard. Paths resolve lexically (normpath), never Path.resolve(), since HF cache snapshot files symlink into blobs/ and resolving would leave the snapshot dir and false-block every sharded model offline. An absolute path, a .. traversal that escapes the snapshot, or an unreadable/invalid index fails closed. The existing safetensors-sibling suppression is kept. * Studio: classify offline indexed shards by torch.load path, not pickle extension load_state_dict picks safetensors vs torch.load per shard by the shard's own suffix, so two offline-gate gaps remained: - A model.safetensors.index.json whose weight_map points at a .bin shard was suppressed by has_base_safetensors (the index file itself matches the base safetensors regex), yet Transformers still torch.loads that shard. Only the pytorch index is superseded by a base safetensors now; a safetensors index is the chosen archive, so its non-safetensors targets are always flagged. - A pytorch index can map weights to arbitrary names (shards/payload, weights.data); the loader torch.loads any target not ending in .safetensors. Flag indexed shards by that rule instead of a pickle-extension allowlist. Restrict the scan to the two torch-family indexes (tf/flax load via non-pickle loaders). Add regression tests for both cases. * Studio: match offline weight-index filenames case-insensitively The index-name check compared the on-disk filename exactly, while the surrounding weight and safetensors matches use case-insensitive rules. On a case-insensitive volume (Windows or macOS) from_pretrained opens an oddly-cased cache file such as PYTORCH_MODEL.BIN.INDEX.JSON when it requests the canonical lowercase name, so the exact-case check skipped it and a nested pickle shard it referenced was allowed through. Lower-case the index name before matching, as the rest of the gate does, and add a regression test. * Studio: match load_state_dict format/selection exactly in the offline index scan Two edge cases in the offline weight-index scan: - load_state_dict decides safetensors vs torch.load with a case-sensitive endswith(".safetensors"), so a shard named payload.SAFETENSORS still deserializes via torch.load. Classify indexed shard suffixes case-sensitively to match, instead of lower-casing (which treated such a shard as inert). - A complete direct model.safetensors is selected before either sharded index, so a stale model.safetensors.index.json referencing a .bin shard never loads. Skip both indexes when a direct model.safetensors is present, so an otherwise loadable model is not over-blocked. Add regression tests for both. * Studio: read the offline weight index as UTF-8 Path.read_text() uses the locale default, which is cp1252 on Windows, so a UTF-8 weight index with non-ASCII bytes raised UnicodeDecodeError and the gate blocked an otherwise loadable model. JSON is UTF-8 by spec (and how the loader reads it), so pin the encoding. * Studio: resolve safetensors alternatives via the loader's own filename lookup The offline gate decided a safetensors alternative existed by case-folding the directory listing. On a case-sensitive filesystem that let an uppercase decoy such as MODEL.SAFETENSORS suppress the pickle scan, yet from_pretrained asks for the canonical lowercase model.safetensors, does not find the decoy, and selects the pickle (a direct pytorch_model.bin or the pytorch index) and deserializes it. Probe each alternative with (root / name).is_file() instead, mirroring the loader: is_file() honors the platform's case rules, so a decoy suppresses only where the loader would truly open it. Suppression must never fail open; detection stays case-insensitive (fail closed). Add regression tests for the direct and indexed pickle decoys (skipped on case-insensitive volumes, where no bypass exists). * Studio: resolve indexes and shards exactly as from_pretrained does Two more loader-fidelity gaps in the offline index scan: - Shard lookup normalized backslashes to forward slashes. On POSIX a backslash is a literal filename character, so an index naming dir\payload.bin matches a real pickle of that exact name that Transformers joins and deserializes, while the normalized dir/payload.bin missed it. Join the raw weight_map value with os.path.join so the probe mirrors the loader on each platform. - Index detection case-folded the directory listing, so on a case-sensitive filesystem an uppercase PYTORCH_MODEL.BIN.INDEX.JSON artifact the loader never opens was treated as live and its shard blocked. Probe the canonical name with the loader's own is_file lookup instead, so an index counts only where from_pretrained would actually load it. Update the uppercase-index tests to assert the correct per-filesystem behavior and add a POSIX backslash-shard regression test. --------- Co-authored-by: danielhanchen --- .../tests/test_offline_embedding_minimal.py | 344 ++++++++++++++++++ .../backend/utils/security/file_security.py | 113 +++++- 2 files changed, 438 insertions(+), 19 deletions(-) diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py index 7ff580c36d..ccc6b5f76a 100644 --- a/studio/backend/tests/test_offline_embedding_minimal.py +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -51,6 +51,27 @@ def _modules_json(*paths): _COMMIT = "0123456789abcdef0123456789abcdef01234567" +def _fs_case_sensitive(root): + """Whether root's filesystem is case-sensitive (Linux yes; macOS/Windows usually no). The gate + mirrors the loader, whose file lookups follow the same rule, so some cases only exist on one.""" + probe = Path(root) / "_case_probe" + probe.write_text("x") + try: + return not (Path(root) / "_CASE_PROBE").exists() + finally: + probe.unlink() + + +def _requires_case_sensitive_fs(root): + if not _fs_case_sensitive(root): + pytest.skip("requires a case-sensitive filesystem") + + +def _requires_case_insensitive_fs(root): + if _fs_case_sensitive(root): + pytest.skip("requires a case-insensitive filesystem") + + def _make_cache( root, repo_id, @@ -382,6 +403,329 @@ def test_gate_blocks_sharded_pickle(hf_cache): assert _offline_decision("org/shard").blocked is True +def test_gate_blocks_indexed_pickle_shard_in_subdirectory(hf_cache): + # from_pretrained follows weight_map paths relative to the root index, so these nested shards + # are deserialized even though they are not direct children of the load root (iterdir misses + # them). The online gate blocks index-referenced subdir pickles; the offline gate must too. + _make_cache( + hf_cache, + "org/indexed-shard", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"layer.weight": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-shard") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_pickle_shard_with_nonstandard_stem(hf_cache): + # The index tells the loader to deserialize this file, so a pickle EXTENSION is enough -- the + # shard's stem need not match the on-disk weight-name heuristic (which only guesses bare files). + _make_cache( + hf_cache, + "org/indexed-odd", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/evil-00001-of-00001.bin"}}', + "shards/evil-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-odd") + assert decision.blocked is True + assert any(u["path"] == "shards/evil-00001-of-00001.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_safetensors_index_pointing_to_pickle_shard(hf_cache): + # load_state_dict picks safetensors vs torch.load by each shard's own suffix, so a + # model.safetensors.index.json that maps a weight to a .bin shard still deserializes it. The + # index's own existence must not suppress the shard it names. + _make_cache( + hf_cache, + "org/st-index-pickle", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/st-index-pickle") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_shard_with_no_pickle_extension(hf_cache): + # Transformers torch.loads any indexed shard not ending in .safetensors, so an unconventional + # extensionless name is still a deserialization target. + _make_cache( + hf_cache, + "org/indexed-noext", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload"}}', + "shards/payload": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-noext") + assert decision.blocked is True + assert any(u["path"] == "shards/payload" for u in decision.unsafe_files) + + +_UPPER_INDEX_FILES = { + "PYTORCH_MODEL.BIN.INDEX.JSON": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", +} + + +def test_gate_blocks_uppercase_index_on_case_insensitive_fs(hf_cache): + # On a case-insensitive volume (Windows/macOS) from_pretrained opens an oddly-cased index when it + # requests the canonical lowercase name, so the loader-mirror lookup resolves it and blocks. + _requires_case_insensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + decision = _offline_decision("org/upper-index") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_allows_uppercase_index_on_case_sensitive_fs(hf_cache): + # On a case-sensitive FS from_pretrained's os.path.isfile of the canonical lowercase name misses + # the uppercase artifact and never loads its shard, so the gate must not over-block it. + _requires_case_sensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + assert _offline_decision("org/upper-index").blocked is False + + +def test_gate_blocks_indexed_shard_named_with_backslash(hf_cache): + # On POSIX a backslash is a literal filename char, so from_pretrained joins the raw weight_map + # value and deserializes a file actually named "dir\payload.bin"; the gate must probe it verbatim. + import os + + if os.sep != "/": + pytest.skip("backslash is a path separator off POSIX") + _make_cache( + hf_cache, + "org/backslash", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "dir\\\\payload.bin"}}', + "dir\\payload.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/backslash") + assert decision.blocked is True + assert any(u["path"] == "dir\\payload.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_shard_with_uppercase_safetensors_suffix(hf_cache): + # load_state_dict's endswith(".safetensors") is case-sensitive, so a shard named payload.SAFETENSORS + # falls to torch.load. The gate must classify shard suffixes case-sensitively to match it. + _make_cache( + hf_cache, + "org/upper-suffix", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload.SAFETENSORS"}}', + "shards/payload.SAFETENSORS": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-suffix") + assert decision.blocked is True + assert any(u["path"] == "shards/payload.SAFETENSORS" for u in decision.unsafe_files) + + +def test_gate_allows_stale_safetensors_index_beside_direct_safetensors(hf_cache): + # A complete direct model.safetensors is selected before either index, so a stale + # model.safetensors.index.json referencing a .bin shard never deserializes -> must not block. + _make_cache( + hf_cache, + "org/direct-plus-stale-index", + { + "model.safetensors": "tensors", + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + assert _offline_decision("org/direct-plus-stale-index").blocked is False + + +def test_gate_blocks_pytorch_index_with_uppercase_safetensors_decoy(hf_cache): + # On a case-sensitive FS, from_pretrained asks for the canonical lowercase model.safetensors, does + # not find an uppercase decoy, and selects the pytorch index instead. The decoy must not suppress. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy", + { + "MODEL.SAFETENSORS": "decoy", + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_direct_pickle_with_uppercase_safetensors_decoy(hf_cache): + # Same decoy against a direct pytorch_model.bin: the loader selects the pickle, so the uppercase + # safetensors must not suppress it on a case-sensitive FS. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy-direct", + {"MODEL.SAFETENSORS": "decoy", "pytorch_model.bin": "pickle"}, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy-direct") + assert decision.blocked is True + assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_pickle_shard_in_module_subdir(hf_cache): + # A weight index inside a sentence-transformers module load root points at a nested pickle shard. + _make_cache( + hf_cache, + "org/mod-indexed", + { + "modules.json": _modules_json("0_Transformer"), + "0_Transformer/pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "0_Transformer/shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/mod-indexed") + assert decision.blocked is True + assert any( + u["path"] == "0_Transformer/shards/pytorch_model-00001-of-00001.bin" + for u in decision.unsafe_files + ) + + +def test_gate_allows_indexed_pickle_shard_with_safetensors_sibling(hf_cache): + # A base model.safetensors makes the loader ignore the pickle index entirely, so it must not + # block (mirrors the direct-file safetensors-sibling suppression). + _make_cache( + hf_cache, + "org/indexed-both", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + "model.safetensors": "y", + }, + ) + with _no_network(): + assert _offline_decision("org/indexed-both").blocked is False + + +def test_gate_allows_indexed_safetensors_shard_in_subdirectory(hf_cache): + # A safetensors index lists inert shards -- following it must never block (guards against a + # scanner that flags every indexed shard regardless of format). + _make_cache( + hf_cache, + "org/st-indexed", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}' + ), + "shards/model-00001-of-00001.safetensors": "tensors", + }, + ) + with _no_network(): + assert _offline_decision("org/st-indexed").blocked is False + + +def test_gate_blocks_on_index_path_traversal(hf_cache): + # A weight_map entry escaping the snapshot via ".." is abnormal/hostile -> fail closed. + _make_cache( + hf_cache, + "org/escape", + {"pytorch_model.bin.index.json": '{"weight_map": {"w": "../../../../etc/evil.bin"}}'}, + ) + with _no_network(): + assert _offline_decision("org/escape").blocked is True + + +def test_gate_allows_symlinked_sharded_safetensors(tmp_path, monkeypatch): + # Real HF caches store snapshot files as symlinks into blobs/. A resolve()-based containment + # check would escape the snapshot and false-block every sharded model; the lexical gate must not. + import hashlib + import os + + from huggingface_hub.file_download import repo_folder_name + + root = tmp_path / "hub" + root.mkdir() + monkeypatch.setenv("HF_HOME", str(tmp_path)) + monkeypatch.setenv("HF_HUB_CACHE", str(root)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = root), + ) + repo_dir = root / repo_folder_name(repo_id = "org/sym", repo_type = "model") + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text(_COMMIT) + blobs = repo_dir / "blobs" + blobs.mkdir() + snapshot = repo_dir / "snapshots" / _COMMIT + (snapshot / "shards").mkdir(parents = True) + + def _blobbed(rel, content): + digest = hashlib.sha256(content.encode()).hexdigest() + (blobs / digest).write_text(content) + target = snapshot / rel + target.parent.mkdir(parents = True, exist_ok = True) + target.symlink_to(os.path.relpath(blobs / digest, target.parent)) + + _blobbed("config.json", "{}") + _blobbed( + "model.safetensors.index.json", + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}', + ) + _blobbed("shards/model-00001-of-00001.safetensors", "tensors") + with _no_network(): + assert _offline_decision("org/sym").blocked is False + + +def test_gate_allows_index_without_weight_map(hf_cache): + # An index whose top-level JSON has no dict weight_map lets the loader resolve no shards, so it + # must not crash or block on its own (only inert safetensors are cached here). + _make_cache( + hf_cache, + "org/no-wm", + {"model.safetensors.index.json": "[]", "model.safetensors": "x"}, + ) + with _no_network(): + assert _offline_decision("org/no-wm").blocked is False + + def test_gate_allows_nothing_cached(hf_cache): with _no_network(): assert _offline_decision("org/missing").blocked is False diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 91d7ad8f0e..892f7862a9 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -46,17 +46,6 @@ _PICKLE_WEIGHT_RE = re.compile( r"\.(bin|pt|pth|ckpt|pkl|pickle)$", re.IGNORECASE, ) -# Base-model safetensors set: HF names the base pickle pytorch_model.bin but the safetensors -# model.safetensors (stems differ), so a base pickle is replaced only by these, not an adapter's. -_BASE_SAFETENSORS_RE = re.compile( - r"^(model(-\d+-of-\d+)?\.safetensors|model\.safetensors\.index\.json)$", - re.IGNORECASE, -) -# Adapter (PEFT) safetensors set: adapter_model.safetensors, its shards, or index. -_ADAPTER_SAFETENSORS_RE = re.compile( - r"^(adapter_model(-\d+-of-\d+)?\.safetensors|adapter_model\.safetensors\.index\.json)$", - re.IGNORECASE, -) # Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/ # malicious or a future label) blocks, so Hub schema drift fails CLOSED. @@ -94,6 +83,13 @@ _INERT_SUFFIXES = frozenset( _SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"}) +# Torch-family weight indexes: from_pretrained feeds each shard they name to load_state_dict, which +# torch.load()s (pickle) any shard whose name does not end in .safetensors, whatever its stem. A +# pytorch index is superseded when a base safetensors is present (the loader prefers it); a +# safetensors index IS the chosen archive, so a non-safetensors target it names still loads. tf/flax +# indexes load via non-pickle loaders, so they are not a torch.load vector here. +_TORCH_INDEX_FILES = ("pytorch_model.bin.index.json", "model.safetensors.index.json") + # Root weight-index files. from_pretrained reads these to find sharded weights, so a # flagged subdir pickle is a load vector iff a root index references it. _TRANSFORMERS_INDEX_FILES = ( @@ -313,13 +309,72 @@ def _st_load_roots(snapshot: Path) -> list: return roots +def _indexed_pickle_shards(index_path: Path, root: Path, snapshot: Path) -> list: + """Shards a torch weight index points a ``from_pretrained`` load at that load_state_dict would + torch.load (pickle): every ``weight_map`` target NOT ending in ``.safetensors``, whatever its + stem (an arbitrary name like ``shards/payload`` still deserializes). Resolved relative to the + index dir (``root``) like the loader, so a shard in a nested dir is followed (iterdir misses it). + Lexical only, never ``Path.resolve()`` (HF snapshot files symlink into ``blobs/``, so resolving + escapes the snapshot and false-blocks every shard). Raises OSError -> caller fails CLOSED on an + unreadable/invalid index or a target escaping the snapshot.""" + import json + import os + + try: + # JSON is UTF-8 by spec; pin it so a non-ASCII index is not misdecoded (and needlessly + # blocked) under Windows' cp1252 default. + parsed = json.loads(index_path.read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + raise OSError(f"unreadable weight index: {index_path}") from exc + weight_map = parsed.get("weight_map") if isinstance(parsed, dict) else None + if not isinstance(weight_map, dict): + return [] # no dict weight_map -> the loader resolves no shards from this index + snapshot_norm = os.path.normpath(str(snapshot)) + shards = [] + for shard in weight_map.values(): + raw = str(shard) + if not raw: + continue + # Join the RAW weight_map value like from_pretrained's os.path.join: on POSIX a backslash is a + # literal filename char (not a separator), so normalizing it would probe a different path than + # the loader opens. normpath + containment stay platform-aware (os.sep) to block "..". + joined = os.path.normpath(os.path.join(str(root), raw)) + if joined != snapshot_norm and not joined.startswith(snapshot_norm + os.sep): + raise OSError(f"weight index escapes the snapshot: {index_path}") + shard_path = Path(joined) + # Case-SENSITIVE, mirroring load_state_dict's own endswith(".safetensors"): a shard named + # payload.SAFETENSORS is not treated as safetensors by the loader and falls to torch.load. + if not shard_path.name.endswith(".safetensors") and shard_path.is_file(): + shards.append(shard_path) + return shards + + +def _loader_resolves(root: Path, name: str) -> bool: + """True iff from_pretrained would open ``name`` under ``root``. ``is_file()`` honors the platform + (case-sensitive on Linux, case-insensitive on Windows/macOS), so it mirrors the loader's own + lookup: an oddly-cased decoy counts as an alternative only where the loader would truly open it. + A name-fold instead would let an uppercase MODEL.SAFETENSORS suppress the scan on Linux while the + loader, asking for the canonical lowercase name, silently falls through to a pickle index.""" + return (root / name).is_file() + + def _cached_pickle_weight_files(snapshot: Path) -> list: - """Pickle weight files in snapshot's ST load roots, EXCLUDING those whose weight family also - ships an inert safetensors in the same dir (the loader prefers it): a base pickle is suppressed - only by a base model.safetensors, an adapter pickle only by adapter_model.safetensors -- an - unrelated safetensors is no substitute. Load roots only. Raises OSError if the snapshot root is - unreadable (caller blocks).""" + """Pickle weight files a SentenceTransformer/Transformers load deserializes from snapshot's ST + load roots, EXCLUDING those whose weight family also ships an inert safetensors in the same dir + (the loader prefers it): a base pickle is suppressed only by a base model.safetensors, an adapter + pickle only by adapter_model.safetensors -- an unrelated safetensors is no substitute. Covers + both direct-child pickles AND pickle shards referenced by a local weight index (which the loader + follows into nested dirs, matching the online gate). Raises OSError -- caller fails CLOSED -- if + the snapshot root or a weight index is unreadable, or an index reference escapes the snapshot.""" blocked = [] + seen = set() + + def _add(path: Path): + key = str(path) + if key not in seen: + seen.add(key) + blocked.append(path) + for root in _st_load_roots(snapshot): try: entries = [p for p in root.iterdir() if p.is_file()] @@ -327,15 +382,35 @@ def _cached_pickle_weight_files(snapshot: Path) -> list: if root == snapshot: raise # top-level unreadable -> fail closed continue # unreadable module subdir: nothing loadable to attest here - has_base_safetensors = any(_BASE_SAFETENSORS_RE.match(p.name) for p in entries) - has_adapter_safetensors = any(_ADAPTER_SAFETENSORS_RE.match(p.name) for p in entries) + # Safetensors alternatives the loader would actually resolve (never a bare name-fold, which + # fails OPEN: see _loader_resolves). A base pickle is replaced only by a base safetensors, an + # adapter pickle only by an adapter one. A single model.safetensors also outranks BOTH indexes. + has_direct_base_safetensors = _loader_resolves(root, "model.safetensors") + has_base_safetensors = has_direct_base_safetensors or _loader_resolves( + root, "model.safetensors.index.json" + ) + has_adapter_safetensors = _loader_resolves(root, "adapter_model.safetensors") for path in entries: if not _PICKLE_WEIGHT_RE.match(path.name): continue is_adapter = path.name.lower().startswith("adapter_model") has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors if not has_alternative: - blocked.append(path) + _add(path) + # A torch weight index makes from_pretrained load nested shards iterdir never sees; the loader + # torch.loads any not ending in .safetensors. Probe the canonical index name with the loader's + # own lookup (_loader_resolves), so an oddly-cased artifact it would never open does not block. + # A direct model.safetensors wins over BOTH indexes; failing that a base safetensors still + # outranks the pytorch index, while a safetensors index is itself the chosen archive. + for index_name in _TORCH_INDEX_FILES: + if not _loader_resolves(root, index_name): + continue + if has_direct_base_safetensors: + continue + if index_name == "pytorch_model.bin.index.json" and has_base_safetensors: + continue + for shard_path in _indexed_pickle_shards(root / index_name, root, snapshot): + _add(shard_path) return blocked From 6f60bf4f82aa2f192e7163c8d1eeb7c5bc8938d7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 20:18:36 -0700 Subject: [PATCH 012/161] Studio whisper: pair slim bundles on the ggml commit, not the full llama tag (#7381) The slim whisper bundle is ggml-less and links the ggml runtime out of the installed llama.cpp prebuilt, so each whisper release pins a paired llama tag. The gate required an exact tag match, but llama fork tags are b-mix- and the build number tracks upstream llama and fork PRs that live outside ggml. When llama republishes a newer build with the same ggml commit (a frequent event), the installed llama advances past the whisper pin and curated dictation goes unavailable until whisper is republished, even though the ggml runtime is ABI-identical. Key the pairing gate on the ggml commit after -mix- instead of the full tag, in all three comparison sites (slim_pairing_for_artifact, _slim_release_incompatibility, resolve_selection). requires_ggml_sonames stays the real per-file ABI gate, and a genuine ggml skew still fails closed. Tags without a -mix- marker fall back to exact matching. --- studio/install_whisper_prebuilt.py | 36 +++++++++++--- .../test_install_whisper_prebuilt_logic.py | 49 ++++++++++++++++++- 2 files changed, 77 insertions(+), 8 deletions(-) diff --git a/studio/install_whisper_prebuilt.py b/studio/install_whisper_prebuilt.py index 1c81ad6105..f1a7fce5e0 100644 --- a/studio/install_whisper_prebuilt.py +++ b/studio/install_whisper_prebuilt.py @@ -335,6 +335,31 @@ def artifacts_for_host( # ── Slim selection (paired with the installed llama.cpp ggml runtime) ── +def _llama_ggml_commit(tag: str) -> str | None: + """The ggml commit a llama.cpp fork tag was built against. Fork tags are + "b-mix-"; the ggml commit after "-mix-" fixes + the ggml ABI the slim whisper bundle links against, while the build number + only tracks upstream llama / fork PRs outside ggml. None when the tag has no + "-mix-" marker (then only an exact tag pairs).""" + marker = "-mix-" + idx = tag.rfind(marker) + end = idx + len(marker) + return tag[end:] if idx >= 0 and end < len(tag) else None + + +def llama_runtime_pairs(installed_tag: str, required_tag: Any) -> bool: + """Whether an installed llama tag can back a slim bundle needing required_tag. + An exact tag always pairs; so does a shared ggml commit, since a newer llama + build with the same ggml ships an ABI-identical runtime. requires_ggml_sonames + stays the real per-file ABI gate.""" + if not isinstance(required_tag, str): + return False + if installed_tag == required_tag: + return True + commit = _llama_ggml_commit(installed_tag) + return commit is not None and commit == _llama_ggml_commit(required_tag) + + def slim_pairing_for_artifact( artifact: dict[str, Any], host: HostInfo, backend: str ) -> tuple[Path, str] | None: @@ -348,10 +373,10 @@ def slim_pairing_for_artifact( return None llama_bin_dir, llama_tag, _profile = runtime requires_tag = artifact.get("requires_llama_tag") - if not isinstance(requires_tag, str) or requires_tag != llama_tag: + if not llama_runtime_pairs(llama_tag, requires_tag): log( f"slim_selection: {asset} skipped: installed llama tag {llama_tag!r} " - f"!= required {requires_tag!r}" + f"does not pair with required {requires_tag!r}" ) return None sonames = artifact.get("requires_ggml_sonames") @@ -466,11 +491,10 @@ def _slim_release_incompatibility(manifest: dict[str, Any], host: HostInfo) -> s for artifact in os_compatible if isinstance(artifact.get("requires_llama_tag"), str) } - if required_tags and installed_tag not in required_tags: + if required_tags and not any(llama_runtime_pairs(installed_tag, tag) for tag in required_tags): required_tag = sorted(required_tags)[0] return ( - f"slim bundle requires llama.cpp {required_tag}; " - f"installed llama.cpp is {installed_tag}" + f"slim bundle requires llama.cpp {required_tag}; installed llama.cpp is {installed_tag}" ) return None @@ -820,7 +844,7 @@ def selection_from_artifact( # A slim selection carries its pairing so the install wiring and marker know # which llama runtime provides the ggml libraries. runtime = installed_llama_runtime() - if runtime is None or runtime[1] != artifact.get("requires_llama_tag"): + if runtime is None or not llama_runtime_pairs(runtime[1], artifact.get("requires_llama_tag")): raise PrebuiltFallback( "the paired llama.cpp runtime changed underneath the slim whisper selection" ) diff --git a/tests/studio/install/test_install_whisper_prebuilt_logic.py b/tests/studio/install/test_install_whisper_prebuilt_logic.py index 7364ea454d..0a5541b27c 100644 --- a/tests/studio/install/test_install_whisper_prebuilt_logic.py +++ b/tests/studio/install/test_install_whisper_prebuilt_logic.py @@ -441,8 +441,9 @@ def test_main_forwards_requested_whisper_tags(tmp_path, monkeypatch): monkeypatch.setattr( M, "resolve_prebuilt", - lambda host, **kwargs: seen.update(kwargs) - or {"prebuilt_available": False, "repo": "unslothai/whisper.cpp"}, + lambda host, **kwargs: ( + seen.update(kwargs) or {"prebuilt_available": False, "repo": "unslothai/whisper.cpp"} + ), ) assert M.main(["--resolve-prebuilt", "v1.8.0", "--output-format", "json"]) == 0 assert seen["whisper_tag"] == "v1.8.0" @@ -807,6 +808,50 @@ def test_slim_release_tag_skew_has_distinct_compatibility_error(tmp_path, monkey M.select_artifact_with_fallback(manifest, _cuda_host(), "cuda") +# A newer llama build that keeps the same ggml commit as SLIM_LLAMA_TAG. +NEWER_LLAMA_TAG = "b10079-mix-fb3d4ca" + + +@pytest.mark.parametrize( + "installed,required,pairs", + [ + (SLIM_LLAMA_TAG, SLIM_LLAMA_TAG, True), # exact tag + (NEWER_LLAMA_TAG, SLIM_LLAMA_TAG, True), # newer build, same ggml commit + ("b10069-mix-0000000", SLIM_LLAMA_TAG, False), # same build, different ggml + (SLIM_LLAMA_TAG, None, False), # no requirement recorded + ("b10069", "b10069", True), # tag without -mix-, exact only + ("b10070", "b10069", False), # tag without -mix-, no shared key + ], +) +def test_llama_runtime_pairs_keys_on_ggml_commit(installed, required, pairs): + assert M.llama_runtime_pairs(installed, required) is pairs + + +def test_slim_pairs_across_llama_build_bump_with_same_ggml(tmp_path, monkeypatch): + # The live failure: the llama installer advances to a newer build that keeps + # the same ggml commit, so the slim bundle's paired runtime is ABI-identical + # and must still select rather than degrade to CPU or report unavailable. + bin_dir = _fake_llama_bin(tmp_path) + monkeypatch.setattr( + M, "installed_llama_runtime", lambda: (bin_dir, NEWER_LLAMA_TAG, "cuda13-newer") + ) + artifact, backend, used_fallback = M.select_artifact_with_fallback( + _slim_manifest(), _cuda_host(), "cuda" + ) + assert artifact["asset"] == SLIM_ASSET + assert backend == "cuda" and used_fallback is False + + +def test_slim_build_bump_same_ggml_is_not_a_compatibility_error(tmp_path, monkeypatch): + # A same-ggml build bump must not surface as a release incompatibility (the + # update path reports that as unavailable); only a real ggml skew does. + bin_dir = _fake_llama_bin(tmp_path) + monkeypatch.setattr( + M, "installed_llama_runtime", lambda: (bin_dir, NEWER_LLAMA_TAG, "cuda13-newer") + ) + assert M._slim_release_incompatibility(_slim_manifest(), _cuda_host()) is None + + def test_link_ggml_runtime_hardlinks_every_ggml_library(tmp_path): bin_dir = _fake_llama_bin(tmp_path) whisper_bin = tmp_path / "whisper.cpp" / "build" / "bin" From a0f58c1128d59d653f3c28c7c714ce36079ccb3b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 20:47:29 -0700 Subject: [PATCH 013/161] Unsloth start: keep Claude subagents on the local model (#7333) Add CLAUDE_CODE_SUBAGENT_MODEL=inherit to the session-only claude settings overlay so built-in subagents stay on the loaded local model. --- unsloth_cli/commands/start.py | 14 +++++++++++--- unsloth_cli/tests/test_start.py | 7 ++++++- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 7c768d16f5..3d407576f8 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -1422,10 +1422,18 @@ _DYNAMIC_SECTIONS_FLAG = "--exclude-dynamic-system-prompt-sections" def _claude_settings_overlay(model_id: str) -> str: # Session-only `claude --settings` overlay (command-line tier, no ~/.claude write): - # suppress the attribution header, and pin availableModels to the served model so a - # user allowlist can't reject it. The pin must be non-empty; [] is ignored. + # suppress the attribution header, keep every subagent on the served model (a user + # CLAUDE_CODE_SUBAGENT_MODEL pin would otherwise route delegated work off the local + # endpoint), and pin availableModels to the served model so a user allowlist can't + # reject it. The pin must be non-empty; [] is ignored. return json.dumps( - {"env": {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}, "availableModels": [model_id]} + { + "env": { + "CLAUDE_CODE_ATTRIBUTION_HEADER": "0", + "CLAUDE_CODE_SUBAGENT_MODEL": "inherit", + }, + "availableModels": [model_id], + } ) diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 36a5e51938..8ce0bcc0d4 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -127,6 +127,8 @@ def test_claude_settings_overlay_pins_served_model(): assert overlay["availableModels"] == [MODEL["id"]] # The attribution-header suppression is preserved alongside it. assert overlay["env"]["CLAUDE_CODE_ATTRIBUTION_HEADER"] == "0" + # Subagents fall through to the served model instead of a user's opus/sonnet pin. + assert overlay["env"]["CLAUDE_CODE_SUBAGENT_MODEL"] == "inherit" def test_install_agent_prompts_then_installs(monkeypatch): @@ -784,7 +786,10 @@ def test_connect_claude_no_launch(fake_studio): _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 + command = _launch_command(result.output) + settings = json.loads(command[command.index("--settings") + 1]) + assert settings["env"]["CLAUDE_CODE_SUBAGENT_MODEL"] == "inherit" + assert "--plugin-dir" not in command assert ".claude/settings.json" not in result.output From 387547980305a5ae66c4f86c40cc43d62b9345bb Mon Sep 17 00:00:00 2001 From: oobabooga Date: Fri, 24 Jul 2026 00:48:30 -0300 Subject: [PATCH 014/161] Complete local subagent delegation for Codex, Claude plan mode, and Pi (#7329) Add session-scoped MCP bridges so Codex, Claude plan mode, and Pi subagents run on the loaded local model, with cloud credentials and Codex state isolated per session and a process-wide Pi agent cap. --- unsloth_cli/claude_subagent_mcp.py | 139 +++++-- unsloth_cli/codex_subagent_mcp.py | 171 ++++++++ unsloth_cli/commands/start.py | 373 ++++++++++++++---- unsloth_cli/pi_subagent.ts | 353 ++++++++++++----- unsloth_cli/tests/test_claude_subagent_mcp.py | 67 ++++ unsloth_cli/tests/test_codex_subagent_mcp.py | 228 +++++++++++ unsloth_cli/tests/test_pi_subagent.py | 301 +++++++++++++- unsloth_cli/tests/test_start.py | 274 +++++++++++-- 8 files changed, 1645 insertions(+), 261 deletions(-) create mode 100644 unsloth_cli/codex_subagent_mcp.py create mode 100644 unsloth_cli/tests/test_codex_subagent_mcp.py diff --git a/unsloth_cli/claude_subagent_mcp.py b/unsloth_cli/claude_subagent_mcp.py index b86368515b..e044d78705 100644 --- a/unsloth_cli/claude_subagent_mcp.py +++ b/unsloth_cli/claude_subagent_mcp.py @@ -19,6 +19,8 @@ from unsloth_cli.commands.start import ( _CLAUDE_ENV_UNSET, _SUBAGENT_DESCRIPTION, _SUBAGENT_INSTRUCTIONS, + _SUBAGENT_PLAN_DESCRIPTION, + _SUBAGENT_PLAN_INSTRUCTIONS, _claude_flags, _claude_local_env, _wsl_shim_env, @@ -113,7 +115,11 @@ def _stop_child(process: subprocess.Popen) -> None: pass -def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> str: +def run_local_agent( + task: str, + cancel_event: threading.Event | None = None, + read_only: bool = False, +) -> str: base = _required_env("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL") key = _required_env("UNSLOTH_CLAUDE_SUBAGENT_API_KEY") model = _required_env("UNSLOTH_CLAUDE_SUBAGENT_MODEL") @@ -135,16 +141,20 @@ def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> s *_claude_flags(model), "--permission-mode", ( - "bypassPermissions" - if os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS") == "1" - else "acceptEdits" + "plan" + if read_only + else ( + "bypassPermissions" + if os.environ.get("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS") == "1" + else "acceptEdits" + ) ), "--print", "--output-format", "json", "--no-session-persistence", "--append-system-prompt", - _SUBAGENT_INSTRUCTIONS, + _SUBAGENT_PLAN_INSTRUCTIONS if read_only else _SUBAGENT_INSTRUCTIONS, f"Task: {task}", ] bridged, wsl_names = _wsl_shim_env(command, local_env, _CLAUDE_ENV_UNSET) @@ -196,7 +206,15 @@ def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> s return _result_text(stdout) -def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) -> dict | None: +def _response( + request: dict, + run_agent: Callable[[str], str] = run_local_agent, + tool_name: str = "unsloth_agent", + tool_description: str | None = None, + run_read_only_agent: Callable[[str], str] | None = None, + read_only_tool_name: str | None = None, + instructions: str | None = None, +) -> dict | None: request_id = request.get("id") method = request.get("method") if request_id is None: @@ -208,40 +226,55 @@ def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) "capabilities": {"tools": {"listChanged": False}}, "serverInfo": {"name": "unsloth-local-agent", "version": "1.0.0"}, } + if instructions: + result["instructions"] = instructions elif method == "ping": result = {} elif method == "tools/list": - result = { - "tools": [ - { - "name": "unsloth_agent", - "title": "Unsloth local agent", - "description": _SUBAGENT_DESCRIPTION, - "inputSchema": { - "type": "object", - "properties": { - "task": { - "type": "string", - "description": "The complete task for the local Unsloth agent.", - } - }, - "required": ["task"], - "additionalProperties": False, + + def tool_definition(name: str, description: str, read_only: bool) -> dict: + return { + "name": name, + "title": "Unsloth local plan agent" if read_only else "Unsloth local agent", + "description": description, + "inputSchema": { + "type": "object", + "properties": { + "task": { + "type": "string", + "description": "The complete task for the local Unsloth agent.", + } }, - "annotations": { - "readOnlyHint": False, - "destructiveHint": True, - "idempotentHint": False, - "openWorldHint": True, - }, - "_meta": {"anthropic/maxResultSizeChars": _MAX_RESULT_CHARACTERS}, - } - ] - } + "required": ["task"], + "additionalProperties": False, + }, + "annotations": { + "readOnlyHint": read_only, + "destructiveHint": not read_only, + "idempotentHint": read_only, + "openWorldHint": True, + }, + "_meta": {"anthropic/maxResultSizeChars": _MAX_RESULT_CHARACTERS}, + } + + tools = [tool_definition(tool_name, tool_description or _SUBAGENT_DESCRIPTION, False)] + if read_only_tool_name and run_read_only_agent: + tools.append(tool_definition(read_only_tool_name, _SUBAGENT_PLAN_DESCRIPTION, True)) + result = {"tools": tools} elif method == "tools/call": params = request.get("params") or {} arguments = params.get("arguments") or {} - task = arguments.get("task") if params.get("name") == "unsloth_agent" else None + requested_tool = params.get("name") + selected_agent = ( + run_agent + if requested_tool == tool_name + else ( + run_read_only_agent + if requested_tool == read_only_tool_name and run_read_only_agent + else None + ) + ) + task = arguments.get("task") if selected_agent else None if not isinstance(task, str) or not task.strip(): result = { "content": [{"type": "text", "text": "A non-empty task is required."}], @@ -249,7 +282,7 @@ def _response(request: dict, run_agent: Callable[[str], str] = run_local_agent) } else: try: - text = run_agent(task.strip()) + text = selected_agent(task.strip()) result = {"content": [{"type": "text", "text": text}], "isError": False} except Exception as exc: result = { @@ -269,6 +302,11 @@ def serve( stdin: Any = sys.stdin, stdout: Any = sys.stdout, run_agent: Callable[[str, threading.Event], str] = run_local_agent, + tool_name: str = "unsloth_agent", + tool_description: str | None = None, + run_read_only_agent: Callable[[str, threading.Event], str] | None = None, + read_only_tool_name: str | None = None, + instructions: str | None = None, ) -> None: active: dict[object, threading.Event] = {} workers: list[threading.Thread] = [] @@ -308,6 +346,15 @@ def serve( response = _response( request, run_agent = lambda task: run_agent(task, cancel_event), + tool_name = tool_name, + tool_description = tool_description, + run_read_only_agent = ( + (lambda task: run_read_only_agent(task, cancel_event)) + if run_read_only_agent + else None + ), + read_only_tool_name = read_only_tool_name, + instructions = instructions, ) if not cancel_event.is_set(): send(response) @@ -343,7 +390,18 @@ def serve( worker.start() response = None else: - response = _response(request) + response = _response( + request, + tool_name = tool_name, + tool_description = tool_description, + run_read_only_agent = ( + (lambda task: run_read_only_agent(task, threading.Event())) + if run_read_only_agent + else None + ), + read_only_tool_name = read_only_tool_name, + instructions = instructions, + ) except Exception as exc: response = { "jsonrpc": "2.0", @@ -362,5 +420,14 @@ def serve( signal.signal(signum, handler) +def main() -> None: + serve( + run_read_only_agent = lambda task, cancel_event: run_local_agent( + task, cancel_event, read_only = True + ), + read_only_tool_name = "unsloth_plan_agent", + ) + + if __name__ == "__main__": - serve() + main() diff --git a/unsloth_cli/codex_subagent_mcp.py b/unsloth_cli/codex_subagent_mcp.py new file mode 100644 index 0000000000..f075e66404 --- /dev/null +++ b/unsloth_cli/codex_subagent_mcp.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Small stdio MCP bridge from cloud Codex to an explicit local Codex child.""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import threading +from pathlib import Path +from typing import Any + +from unsloth_cli.claude_subagent_mcp import _bounded, _stop_child, serve +from unsloth_cli.commands.start import ( + _CODEX_ENV_KEY, + _CODEX_ENV_UNSET, + _CODEX_PROFILE, + _CODEX_SUBAGENT_CONFIG_ENV, + _CODEX_SUBAGENT_MCP_TOOL, + _CODEX_SUBAGENT_TOOL_DESCRIPTION, + _CODEX_SUBAGENT_ROUTING_INSTRUCTIONS, + _SUBAGENT_INSTRUCTIONS, + _merge_wslenv, + _wsl_shim_env, +) + +_CANCEL_POLL_SECONDS = 0.1 +_SERVER_INSTRUCTIONS = _CODEX_SUBAGENT_ROUTING_INSTRUCTIONS + + +def _config() -> dict: + path = os.environ.get(_CODEX_SUBAGENT_CONFIG_ENV, "").strip() + if not path: + raise RuntimeError(f"Missing {_CODEX_SUBAGENT_CONFIG_ENV}.") + try: + config = json.loads(Path(path).read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + raise RuntimeError("Could not read the local Codex agent configuration.") from exc + if not isinstance(config, dict): + raise RuntimeError("The local Codex agent configuration must be an object.") + for name in ("api_key", "codex_home"): + if not isinstance(config.get(name), str) or not config[name].strip(): + raise RuntimeError(f"The local Codex agent configuration is missing {name}.") + return config + + +def _result_text(stdout: str) -> str: + messages = [] + errors = [] + for line in stdout.splitlines(): + try: + event = json.loads(line) + except ValueError: + continue + if not isinstance(event, dict): + continue + item = event.get("item") + if ( + event.get("type") == "item.completed" + and isinstance(item, dict) + and item.get("type") == "agent_message" + and isinstance(item.get("text"), str) + and item["text"].strip() + ): + messages.append(item["text"].strip()) + if event.get("type") in ("error", "turn.failed"): + detail = event.get("message") or event.get("error") + if isinstance(detail, dict): + detail = detail.get("message") or json.dumps(detail) + if detail: + errors.append(str(detail)) + if errors: + raise RuntimeError(_bounded(errors[-1])) + if messages: + return _bounded(messages[-1]) + raise RuntimeError("The local Codex agent returned no readable result.") + + +def run_local_agent(task: str, cancel_event: threading.Event | None = None) -> str: + config = _config() + executable = shutil.which("codex") + if executable is None: + raise RuntimeError("`codex` is not installed or is not on PATH.") + cancel_event = cancel_event or threading.Event() + if cancel_event.is_set(): + raise RuntimeError("The local Codex agent was cancelled.") + + permissions = ( + ["--dangerously-bypass-approvals-and-sandbox"] + if config.get("bypass_permissions") is True + else ["--sandbox", "workspace-write", "--ask-for-approval", "never"] + ) + command = [ + "codex", + "--oss", + "--profile", + _CODEX_PROFILE, + *permissions, + "exec", + "--ephemeral", + "--json", + "--skip-git-repo-check", + f"{_SUBAGENT_INSTRUCTIONS}\n\nTask: {task}", + ] + local_env = { + _CODEX_ENV_KEY: config["api_key"], + "CODEX_HOME": config["codex_home"], + "CODEX_SQLITE_HOME": config["codex_home"], + } + bridged, wsl_names = _wsl_shim_env(command, local_env, _CODEX_ENV_UNSET) + child_env = dict(os.environ) + if wsl_names: + bridged = {**bridged, "PWD": os.getcwd()} + child_env["WSLENV"] = _merge_wslenv(child_env.get("WSLENV", ""), wsl_names) + for name in _CODEX_ENV_UNSET: + child_env[name] = "" + else: + for name in _CODEX_ENV_UNSET: + child_env.pop(name, None) + child_env.update(bridged) + popen_kwargs: dict[str, Any] = { + "cwd": os.getcwd(), + "env": child_env, + "stdin": subprocess.DEVNULL, + "stdout": subprocess.PIPE, + "stderr": subprocess.PIPE, + "text": True, + } + if os.name == "nt": + popen_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP + else: + popen_kwargs["start_new_session"] = True + process = subprocess.Popen([executable, *command[1:]], **popen_kwargs) + try: + while True: + try: + stdout, stderr = process.communicate(timeout = _CANCEL_POLL_SECONDS) + break + except subprocess.TimeoutExpired: + if cancel_event.is_set(): + _stop_child(process) + raise RuntimeError("The local Codex agent was cancelled.") + except BaseException: + if process.poll() is None: + _stop_child(process) + raise + if process.returncode != 0: + detail = stderr.strip() or stdout.strip() + raise RuntimeError( + _bounded(detail) or f"Local Codex exited with code {process.returncode}." + ) + return _result_text(stdout) + + +def main() -> None: + if len(sys.argv) > 1: + os.environ[_CODEX_SUBAGENT_CONFIG_ENV] = sys.argv[1] + serve( + run_agent = run_local_agent, + tool_name = _CODEX_SUBAGENT_MCP_TOOL, + tool_description = _CODEX_SUBAGENT_TOOL_DESCRIPTION, + instructions = _SERVER_INSTRUCTIONS, + ) + + +if __name__ == "__main__": + main() diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 3d407576f8..434c4a0ad5 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -84,8 +84,33 @@ _SUBAGENT_INSTRUCTIONS = ( "use the available tools when useful, verify your work, and return a concise result to the " "parent agent." ) +_SUBAGENT_PLAN_DESCRIPTION = ( + "Read-only local coding subagent powered by Unsloth for planning and codebase research. " + "Use this local agent when Claude is in plan mode." +) +_SUBAGENT_PLAN_INSTRUCTIONS = ( + "You are a read-only local coding subagent powered by Unsloth. Investigate the assigned " + "task with read-only tools, produce a concrete plan or answer, and return a concise result " + "to the parent agent. Do not modify files." +) _CLAUDE_SUBAGENT_MCP_MODULE = "unsloth_cli.claude_subagent_mcp" _CLAUDE_SUBAGENT_TOOL = "mcp__plugin_unsloth-local-agent_unsloth__unsloth_agent" +_CLAUDE_SUBAGENT_PLAN_TOOL = "mcp__plugin_unsloth-local-agent_unsloth__unsloth_plan_agent" +_CODEX_SUBAGENT_MCP_MODULE = "unsloth_cli.codex_subagent_mcp" +_CODEX_SUBAGENT_MCP_SERVER = "unsloth_local_agent" +_CODEX_SUBAGENT_MCP_TOOL = "spawn_local_agent" +_CODEX_SUBAGENT_CONFIG_ENV = "UNSLOTH_CODEX_SUBAGENT_CONFIG" +_CODEX_PARENT_OVERLAY_MANIFEST = ".unsloth-parent-overlay.json" +_CODEX_SUBAGENT_TOOL_DESCRIPTION = ( + f"{_SUBAGENT_DESCRIPTION} Use this tool instead of the built-in spawn_agent tool for those " + "requests. Other subagent requests may use the built-in tools normally." +) +_CODEX_SUBAGENT_ROUTING_INSTRUCTIONS = ( + "When the user asks to spawn an Unsloth agent or local agent, you must call the " + "spawn_local_agent MCP tool once with the complete task. Do not answer, simulate the " + "result, call wait, or use a built-in subagent before calling the tool. Use built-in " + "subagents for other delegation requests." +) _PI_SUBAGENT_EXTENSION = Path(__file__).parent.parent / "pi_subagent.ts" # OpenCode selects a model by "/". Use a dedicated id to avoid # colliding with a user's providers; provider filters are set in the launch-time overlay. @@ -113,6 +138,7 @@ class _PassthroughCommand(TyperCommand): _CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") +_CODEX_ENV_UNSET = ("OPENAI_API_KEY", "CODEX_API_KEY", "CODEX_ACCESS_TOKEN") # Shared by every agent command; only the config/env/command differ. # Help is grouped into rich panels so `--help` reads as Model / Server / Session @@ -1093,6 +1119,13 @@ def _write_private_json(path: Path, data: dict) -> None: handle.write(json.dumps(data, indent = 2) + "\n") +def _write_private_text(path: Path, text: str) -> None: + 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", encoding = "utf-8") as handle: + handle.write(text) + + 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). @@ -1599,62 +1632,214 @@ def write_codex_config(base: str, model: dict, home: Path) -> None: typer.echo(f"Updated {profile}") -def write_codex_subagent_config(base: str, key: str, model: dict, home: Path) -> Path: - """Write a session-scoped Codex custom agent without replacing the main model.""" - home.mkdir(parents = True, exist_ok = True) - model_id = model["id"] - window = model.get("context_length") or model.get("max_context_length") - catalog_name = "unsloth-model-catalog.json" - text = ( - f"name = {json.dumps(_SUBAGENT_NAME)}\n" - f"description = {json.dumps(_SUBAGENT_DESCRIPTION)}\n" - f"developer_instructions = {json.dumps(_SUBAGENT_INSTRUCTIONS)}\n" - f"model_provider = {json.dumps(_CODEX_PROFILE)}\n" - f"model = {json.dumps(model_id)}\n" +def write_codex_subagent_bridge( + base: str, key: str, model: dict, home: Path, *, yolo: bool +) -> Path: + """Write private config for an explicit local Codex child launched through MCP.""" + child_home = home / "child" + write_codex_config(base, model, child_home) + path = home / "subagent.json" + _write_private_json( + path, + { + "api_key": key, + "codex_home": str(child_home), + "bypass_permissions": yolo, + }, ) - if _codex_supports_model_catalog() and _CODEX_FALLBACK_PROMPT.is_file(): - catalog = home / catalog_name - catalog_text = json.dumps(_codex_model_catalog(model), indent = 2) + "\n" - if not catalog.exists() or catalog.read_text(encoding = "utf-8") != catalog_text: - catalog.write_text(catalog_text, encoding = "utf-8") - typer.echo(f"Updated {catalog}") - text += f"model_catalog_json = {json.dumps(catalog_name)}\n" - if window: - text += f"model_context_window = {int(window)}\n" - credential = home / "unsloth-auth.json" - _write_private_json(credential, {"token": key}) - auth_command = sys.executable - auth_args = [ - "-c", - "import json,sys; print(json.load(open(sys.argv[1], encoding='utf-8'))['token'])", - str(credential), - ] - if _wsl_windows_executable(["codex"]): - auth_command = "wsl.exe" - auth_args = [ - "-d", - os.environ["WSL_DISTRO_NAME"], - "--", - sys.executable, - *auth_args, - ] - text += ( - f"\n{_PROVIDER_HEADER}\n" - 'name = "Unsloth Studio"\n' - f"base_url = {json.dumps(base + '/v1')}\n" - 'wire_api = "responses"\n' - f"\n{_PROVIDER_HEADER[:-1]}.auth]\n" - f"command = {json.dumps(auth_command)}\n" - f"args = {json.dumps(auth_args)}\n" - "timeout_ms = 5000\n" - ) - path = home / f"{_SUBAGENT_NAME}.toml" - if not path.exists() or path.read_text(encoding = "utf-8") != text: - path.write_text(text, encoding = "utf-8") - typer.echo(f"Updated {path}") return path +def _wsl_windows_user_profile(executable: str) -> Path: + """Return the Windows user profile as a path accessible from WSL.""" + profile = os.environ.get("USERPROFILE", "").strip() + if not profile: + try: + profile = subprocess.check_output( + ["cmd.exe", "/d", "/c", "echo %USERPROFILE%"], + text = True, + stderr = subprocess.DEVNULL, + cwd = str(Path(executable).parent), + ).strip() + except (OSError, subprocess.CalledProcessError) as exc: + _fail(f"Could not find the Windows user profile for Codex: {exc}") + if not profile or profile == "%USERPROFILE%": + _fail("Could not find the Windows user profile for Codex.") + if profile.startswith("/"): + return Path(profile) + try: + translated = subprocess.check_output( + ["wslpath", "-u", profile], + text = True, + stderr = subprocess.DEVNULL, + ).strip() + except (OSError, subprocess.CalledProcessError) as exc: + _fail(f"Could not translate Windows user profile {profile}: {exc}") + if not translated: + _fail(f"Could not translate Windows user profile {profile}.") + return Path(translated) + + +def _codex_source_home(*, ignore_configured: bool = False) -> Path: + configured = None if ignore_configured else os.environ.get("CODEX_HOME") + if configured: + if _wsl_windows_executable(["codex"]) and _looks_like_path(configured): + if not configured.startswith("/"): + try: + configured = subprocess.check_output( + ["wslpath", "-u", configured], + text = True, + stderr = subprocess.DEVNULL, + ).strip() + except (OSError, subprocess.CalledProcessError) as exc: + _fail(f"Could not translate Windows CODEX_HOME {configured}: {exc}") + if not configured: + _fail("Could not translate Windows CODEX_HOME.") + return Path(configured).expanduser() + executable = _wsl_windows_executable(["codex"]) + if executable: + return _wsl_windows_user_profile(executable) / ".codex" + return Path.home() / ".codex" + + +def _remove_overlay_entry(path: Path) -> None: + is_junction = getattr(path, "is_junction", None) + if is_junction and is_junction(): + path.rmdir() + elif path.is_symlink() or path.is_file(): + path.unlink() + elif path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + + +def _create_directory_junction(source: Path, target: Path) -> bool: + if os.name != "nt": + return False + try: + result = subprocess.run( + ["cmd.exe", "/d", "/c", "mklink", "/J", str(target), str(source)], + capture_output = True, + text = True, + timeout = 30, + check = False, + ) + except (OSError, subprocess.SubprocessError): + return False + return result.returncode == 0 + + +def write_codex_parent_overlay(overlay: Path) -> Path: + """Add local-agent routing without replacing the cloud parent's configuration.""" + overlay.mkdir(parents = True, exist_ok = True, mode = 0o700) + + manifest_path = overlay / _CODEX_PARENT_OVERLAY_MANIFEST + try: + manifest = json.loads(manifest_path.read_text(encoding = "utf-8")) + except (FileNotFoundError, OSError, json.JSONDecodeError): + manifest = None + source_home = _codex_source_home() + overlay_key = str(overlay.resolve(strict = False)) + source_key = str(source_home.resolve(strict = False)) + if source_key == overlay_key: + previous_source = manifest.get("source_home") if isinstance(manifest, dict) else None + if isinstance(previous_source, str) and previous_source: + candidate = Path(previous_source).expanduser() + if str(candidate.resolve(strict = False)) != overlay_key: + source_home = candidate + else: + source_home = _codex_source_home(ignore_configured = True) + else: + source_home = _codex_source_home(ignore_configured = True) + source_key = str(source_home.resolve(strict = False)) + same_source = isinstance(manifest, dict) and manifest.get("source_home") == source_key + if same_source: + managed_entries = manifest.get("entries", []) + if not isinstance(managed_entries, list): + managed_entries = [] + for name in managed_entries: + if isinstance(name, str) and name not in {"", ".", ".."} and Path(name).name == name: + _remove_overlay_entry(overlay / name) + else: + # A reused overlay must never mix credentials, config, or plugins from two + # different Codex homes. Legacy overlays have no manifest, so rebuild them once. + for target in list(overlay.iterdir()): + _remove_overlay_entry(target) + + # Keep the user's auth, config, plugins, agents, skills, rules, and session state visible. + # Symlinks make this an overlay rather than a stale copy. If Windows denies them, + # use directory junctions so large runtime state remains shared without a bulk copy. + # Copy the configuration surfaces and sessions only if both link forms are unavailable. + fallback_dirs = {"agents", "skills", "rules", "plugins", "marketplaces", "sessions"} + entries = [] + if source_home.is_dir(): + for source in source_home.iterdir(): + if source.name in { + "AGENTS.md", + "AGENTS.override.md", + _CODEX_PARENT_OVERLAY_MANIFEST, + }: + continue + target = overlay / source.name + _remove_overlay_entry(target) + try: + target.symlink_to(source, target_is_directory = source.is_dir()) + entries.append(source.name) + except OSError: + if source.is_file(): + shutil.copy2(source, target) + entries.append(source.name) + elif source.is_dir(): + if _create_directory_junction(source, target): + entries.append(source.name) + elif source.name in fallback_dirs: + shutil.copytree(source, target) + entries.append(source.name) + + _write_private_json( + manifest_path, + {"source_home": source_key, "entries": sorted(entries)}, + ) + + inherited = "" + instruction_name = "AGENTS.md" + for candidate in (source_home / "AGENTS.override.md", source_home / "AGENTS.md"): + try: + text = candidate.read_text(encoding = "utf-8") + except FileNotFoundError: + continue + except OSError as exc: + _fail(f"Could not preserve Codex instructions from {candidate}: {exc}") + if text.strip(): + inherited = text.rstrip() + instruction_name = candidate.name + break + + other_name = "AGENTS.md" if instruction_name == "AGENTS.override.md" else "AGENTS.override.md" + other = overlay / other_name + if other.is_file() or other.is_symlink(): + other.unlink() + routing = _CODEX_SUBAGENT_ROUTING_INSTRUCTIONS + combined = f"{inherited}\n\n{routing}\n" if inherited else f"{routing}\n" + _write_private_text(overlay / instruction_name, combined) + return overlay + + +@contextlib.contextmanager +def _codex_parent_overlay(session_home: Path, *, launch: bool, persist: bool): + if launch and not persist: + temp_root = _agents_config_root() / ".tmp" + temp_root.mkdir(parents = True, exist_ok = True, mode = 0o700) + overlay = Path(tempfile.mkdtemp(prefix = "codex-parent-", dir = temp_root)) + try: + yield write_codex_parent_overlay(overlay) + finally: + shutil.rmtree(overlay, ignore_errors = True) + else: + yield write_codex_parent_overlay(session_home / "parent") + + def _agent_config_path(path: Path, command: list) -> str: """Translate a generated config path when a Windows agent runs through WSL.""" return _wsl_windows_path(path) if _wsl_windows_executable(command) else str(path) @@ -1778,25 +1963,42 @@ def write_claude_subagent_plugin(path: Path, server_env: dict) -> Path: "description: Delegate a task to the local agent powered by Unsloth. Use when the " "user asks to spawn an Unsloth agent or local agent.\n" "---\n\n" - "Call the Unsloth local agent tool once with the complete task. Return its result " - "to the user without claiming that the cloud parent completed the local work.\n", + "Call the Unsloth local agent tool once with the complete task. In plan mode, call " + "the read-only Unsloth plan agent instead. Return its result to the user without " + "claiming that the cloud parent completed the local work.\n", encoding = "utf-8", ) return plugin def _codex_subagent_flags(path: Path) -> list[str]: - config_path = _agent_config_path(path, ["codex"]) - return [ - "--enable", - "multi_agent", - "-c", - "agents.max_depth=1", - "-c", - f"agents.{_SUBAGENT_NAME}.description={json.dumps(_SUBAGENT_DESCRIPTION)}", - "-c", - f"agents.{_SUBAGENT_NAME}.config_file={json.dumps(config_path)}", - ] + command = sys.executable + package_root = str(Path(__file__).resolve().parents[2]) + bootstrap = ( + f"import sys;sys.path.insert(0,{json.dumps(package_root)});" + f"from {_CODEX_SUBAGENT_MCP_MODULE} import main;main()" + ) + args = ["-c", bootstrap, str(path)] + if _wsl_windows_executable(["codex"]): + command = "wsl.exe" + args = [ + "-d", + os.environ["WSL_DISTRO_NAME"], + "--", + sys.executable, + "-c", + bootstrap, + str(path), + ] + server = ( + "{ " + f"command = {json.dumps(command)}, " + f"args = {json.dumps(args)}, " + f"required = true, enabled_tools = [{json.dumps(_CODEX_SUBAGENT_MCP_TOOL)}], " + 'default_tools_approval_mode = "approve", ' + "startup_timeout_sec = 15, tool_timeout_sec = 3600 }" + ) + return ["-c", f"mcp_servers.{_CODEX_SUBAGENT_MCP_SERVER}={server}"] def _wsl_windows_executable(command: list) -> Optional[str]: @@ -2647,7 +2849,7 @@ def claude( _agent_config_path(plugin, ["claude"]), # Before ctx.args: a forwarded `--` would turn later flags positional. "--allowedTools", - _CLAUDE_SUBAGENT_TOOL, + f"{_CLAUDE_SUBAGENT_TOOL},{_CLAUDE_SUBAGENT_PLAN_TOOL}", *_yolo_command_flags("claude", yolo), *ctx.args, ] @@ -2734,25 +2936,32 @@ def codex( subagent_id = _subagent_model_id(base, key, entry, model, gguf_variant) subagent_model = {**entry, "id": subagent_id} with _session_config("codex-subagent", launch, persist = persist) as home: - agent_config = write_codex_subagent_config(base, key, subagent_model, home) - command = [ - "codex", - *_codex_subagent_flags(agent_config), - *_yolo_command_flags("codex", yolo), - *ctx.args, - ] - typer.echo( - "Unsloth is available as the `unsloth` local agent. " - "Ask Codex to spawn an Unsloth or local agent." - ) - _run( + bridge_config = write_codex_subagent_bridge( base, + key, subagent_model, - {}, - command, - launch = launch, - install_hint = "npm install -g @openai/codex", + home, + yolo = yolo, ) + with _codex_parent_overlay(home, launch = launch, persist = persist) as parent_home: + command = [ + "codex", + *_codex_subagent_flags(bridge_config), + *_yolo_command_flags("codex", yolo), + *ctx.args, + ] + typer.echo( + "Unsloth is available as a local agent. " + "Ask Codex to spawn an Unsloth or local agent." + ) + _run( + base, + subagent_model, + {"CODEX_HOME": str(parent_home)}, + command, + launch = launch, + install_hint = "npm install -g @openai/codex", + ) return command = [ "codex", diff --git a/unsloth_cli/pi_subagent.ts b/unsloth_cli/pi_subagent.ts index d712fc89ae..f4ef0c7d9e 100644 --- a/unsloth_cli/pi_subagent.ts +++ b/unsloth_cli/pi_subagent.ts @@ -7,6 +7,7 @@ import { Type } from "typebox"; const provider = "unsloth"; const maxResultCharacters = 100_000; +const maxParallelAgents = 4; const cancelGraceMilliseconds = 2_000; const configPath = process.env.UNSLOTH_PI_SUBAGENT_CONFIG || ""; delete process.env.UNSLOTH_PI_SUBAGENT_CONFIG; @@ -27,6 +28,8 @@ const baseUrl = typeof config.baseUrl === "string" ? config.baseUrl : ""; const apiKey = typeof config.apiKey === "string" ? config.apiKey : ""; const contextWindow = positiveInt(config.contextWindow, 32768); const maxTokens = positiveInt(config.maxTokens, Math.min(Math.floor(contextWindow / 4), 8192)); +let activeAgents = 0; +const waitingAgents: Array<() => boolean> = []; function positiveInt(value: unknown, fallback: number): number { const parsed = Number.parseInt(typeof value === "string" ? value : String(value || ""), 10); @@ -47,6 +50,45 @@ function boundedResult(text: string): string { return `${text.slice(0, maxResultCharacters)}\n\n[Local agent output truncated]`; } +function agentSlotRelease(): () => void { + let released = false; + return () => { + if (released) return; + released = true; + while (waitingAgents.length) { + if (waitingAgents.shift()!()) return; + } + activeAgents -= 1; + }; +} + +function acquireAgentSlot(signal: AbortSignal | undefined): Promise<() => void> { + if (signal?.aborted) return Promise.reject(new Error("The local Unsloth agent was cancelled.")); + if (activeAgents < maxParallelAgents) { + activeAgents += 1; + return Promise.resolve(agentSlotRelease()); + } + return new Promise((resolve, reject) => { + let waiting = true; + const grant = () => { + if (!waiting) return false; + waiting = false; + signal?.removeEventListener("abort", cancel); + resolve(agentSlotRelease()); + return true; + }; + const cancel = () => { + if (!waiting) return; + waiting = false; + const index = waitingAgents.indexOf(grant); + if (index >= 0) waitingAgents.splice(index, 1); + reject(new Error("The local Unsloth agent was cancelled.")); + }; + waitingAgents.push(grant); + signal?.addEventListener("abort", cancel, { once: true }); + }); +} + function piInvocation(args: string[]): { command: string; args: string[] } { const currentScript = process.argv[1]; const bunVirtualScript = currentScript?.startsWith("/$bunfs/root/"); @@ -107,6 +149,136 @@ async function stopChildTree(child: ChildProcess): Promise { signalProcessGroup(child, "SIGKILL"); } +interface LocalAgentResult { + task: string; + response: string; + transcript: any[]; + error?: string; +} + +async function runLocalAgent( + task: string, + cwd: string, + signal: AbortSignal | undefined, + onProgress: (result: LocalAgentResult) => void, +): Promise { + const extension = fileURLToPath(import.meta.url); + const args = [ + "--mode", + "json", + "--print", + "--no-session", + "--provider", + provider, + "--model", + model, + "--no-extensions", + "--extension", + extension, + `Task: ${task}`, + ]; + const invocation = piInvocation(args); + let output = ""; + let stderr = ""; + let childError = ""; + let aborted = false; + const result: LocalAgentResult = { task, response: "", transcript: [] }; + const transcriptEntries = new Set(); + const appendTranscript = (messages: any[]): boolean => { + let changed = false; + for (const message of messages) { + const entry = JSON.stringify(message); + if (transcriptEntries.has(entry)) continue; + transcriptEntries.add(entry); + result.transcript.push(message); + changed = true; + } + return changed; + }; + const processLine = (line: string) => { + try { + const event = JSON.parse(line); + if (event.type === "message_end" && event.message && appendTranscript([event.message])) { + onProgress(result); + } + if ( + event.type === "turn_end" && + Array.isArray(event.toolResults) && + event.toolResults.length && + appendTranscript(event.toolResults) + ) { + onProgress(result); + } + if (event.type !== "message_end") return; + const message = event.message; + // Pi reports model/API failures as message_end events while still + // exiting 0, so the exit status alone cannot surface them. + if (message?.stopReason === "error" || message?.stopReason === "aborted") { + childError = + (typeof message.errorMessage === "string" && message.errorMessage) || + `The local Unsloth agent stopped: ${message.stopReason}.`; + return; + } + const response = finalText(message); + if (response) { + result.response = boundedResult(response); + childError = ""; + } + } catch { + // Ignore non-JSON diagnostic lines. The exit status still reports failures. + } + }; + + const exitCode = await new Promise((resolve, reject) => { + const child = spawn(invocation.command, invocation.args, { + cwd, + detached: process.platform !== "win32", + shell: false, + stdio: ["ignore", "pipe", "pipe"], + env: { + ...process.env, + UNSLOTH_PI_SUBAGENT_CHILD: "1", + UNSLOTH_PI_SUBAGENT_CONFIG: configPath, + }, + }); + let cleanup: Promise | undefined; + const cancel = () => { + if (aborted) return; + aborted = true; + cleanup = stopChildTree(child); + }; + child.on("error", (error) => { + signal?.removeEventListener("abort", cancel); + reject(error); + }); + child.stdout.on("data", (chunk) => { + output += chunk.toString(); + const lines = output.split("\n"); + output = lines.pop() || ""; + for (const line of lines) processLine(line); + }); + child.stderr.on("data", (chunk) => { + stderr = (stderr + chunk.toString()).slice(-100_000); + }); + child.on("close", async (code) => { + signal?.removeEventListener("abort", cancel); + await cleanup; + if (output.trim()) processLine(output); + resolve(code ?? 1); + }); + signal?.addEventListener("abort", cancel, { once: true }); + if (signal?.aborted) cancel(); + }); + + if (aborted) throw new Error("The local Unsloth agent was cancelled."); + if (exitCode !== 0) { + result.error = stderr.trim() || `The local Unsloth agent exited with code ${exitCode}.`; + } + if (childError) result.error = boundedResult(childError); + if (!result.response && !result.error) result.response = "The local agent returned no text."; + return result; +} + export default function unslothSubagent(pi: ExtensionAPI): void { if (!model || !baseUrl || !apiKey || !configPath) { throw new Error("Unsloth subagent configuration is incomplete."); @@ -137,104 +309,97 @@ export default function unslothSubagent(pi: ExtensionAPI): void { name: "unsloth_agent", label: "Unsloth agent", description: - "Local coding subagent powered by Unsloth for debugging, implementation, and codebase research. Use when the user asks to spawn an Unsloth or local agent.", + "Run local coding agents powered by Unsloth for debugging, implementation, and codebase research. Use task for one agent. To run multiple independent agents, use tasks; up to four run concurrently. The tool returns only after every requested agent finishes.", parameters: Type.Object({ - task: Type.String({ description: "The complete task for the local Unsloth agent." }), + task: Type.Optional( + Type.String({ description: "The complete task for one local Unsloth agent." }), + ), + tasks: Type.Optional( + Type.Array(Type.String({ description: "A complete task for one local Unsloth agent." }), { + description: "Independent tasks to run concurrently, one local agent per task.", + minItems: 2, + maxItems: maxParallelAgents, + }), + ), }), - async execute(_toolCallId, params, signal, _onUpdate, ctx) { - const extension = fileURLToPath(import.meta.url); - const args = [ - "--mode", - "json", - "--print", - "--no-session", - "--provider", - provider, - "--model", - model, - "--no-extensions", - "--extension", - extension, - `Task: ${params.task}`, - ]; - const invocation = piInvocation(args); - let output = ""; - let stderr = ""; - let lastResponse = ""; - let childError = ""; - let aborted = false; - const processLine = (line: string) => { - try { - const event = JSON.parse(line); - if (event.type !== "message_end") return; - const message = event.message; - // Pi reports model/API failures as message_end events while still - // exiting 0, so the exit status alone cannot surface them. - if (message?.stopReason === "error" || message?.stopReason === "aborted") { - childError = - (typeof message.errorMessage === "string" && message.errorMessage) || - `The local Unsloth agent stopped: ${message.stopReason}.`; - return; - } - const response = finalText(message); - if (response) { - lastResponse = boundedResult(response); - childError = ""; - } - } catch { - // Ignore non-JSON diagnostic lines. The exit status still reports failures. - } - }; - - const exitCode = await new Promise((resolve, reject) => { - const child = spawn(invocation.command, invocation.args, { - cwd: ctx.cwd, - detached: process.platform !== "win32", - shell: false, - stdio: ["ignore", "pipe", "pipe"], - env: { - ...process.env, - UNSLOTH_PI_SUBAGENT_CHILD: "1", - UNSLOTH_PI_SUBAGENT_CONFIG: configPath, - }, - }); - let cleanup: Promise | undefined; - const cancel = () => { - if (aborted) return; - aborted = true; - cleanup = stopChildTree(child); - }; - child.on("error", (error) => { - signal?.removeEventListener("abort", cancel); - reject(error); - }); - child.stdout.on("data", (chunk) => { - output += chunk.toString(); - const lines = output.split("\n"); - output = lines.pop() || ""; - for (const line of lines) processLine(line); - }); - child.stderr.on("data", (chunk) => { - stderr = (stderr + chunk.toString()).slice(-100_000); - }); - child.on("close", async (code) => { - signal?.removeEventListener("abort", cancel); - await cleanup; - if (output.trim()) processLine(output); - resolve(code ?? 1); - }); - signal?.addEventListener("abort", cancel, { once: true }); - if (signal?.aborted) cancel(); - }); - - if (aborted) throw new Error("The local Unsloth agent was cancelled."); - if (exitCode !== 0) { - throw new Error(stderr.trim() || `The local Unsloth agent exited with code ${exitCode}.`); + executionMode: "parallel", + async execute(_toolCallId, params, signal, onUpdate, ctx) { + const singleTask = typeof params.task === "string" && params.task.trim() ? params.task.trim() : ""; + const parallelTasks = Array.isArray(params.tasks) + ? params.tasks.map((task) => task.trim()).filter(Boolean) + : []; + if (Boolean(singleTask) === Boolean(parallelTasks.length)) { + throw new Error("Provide exactly one of task or tasks."); } - if (childError) throw new Error(boundedResult(childError)); + if (parallelTasks.length > maxParallelAgents) { + throw new Error(`At most ${maxParallelAgents} local agents can run concurrently.`); + } + if (parallelTasks.length === 1) { + throw new Error("Use task for one local agent, or tasks for two to four agents."); + } + + const tasks = singleTask ? [singleTask] : parallelTasks; + const results: Array = new Array(tasks.length); + let completed = 0; + const details = () => ({ + provider, + model, + mode: tasks.length === 1 ? "single" : "parallel", + results: results.filter((result): result is LocalAgentResult => Boolean(result)), + }); + const emitUpdate = () => { + onUpdate?.({ + content: [ + { + type: "text", + text: `Local agents: ${completed}/${tasks.length} completed`, + }, + ], + details: details(), + }); + }; + await Promise.all( + tasks.map(async (task, index) => { + let releaseAgentSlot: (() => void) | undefined; + try { + releaseAgentSlot = await acquireAgentSlot(signal); + results[index] = await runLocalAgent(task, ctx.cwd, signal, (partial) => { + results[index] = partial; + emitUpdate(); + }); + } catch (error) { + results[index] = { + task, + response: "", + transcript: results[index]?.transcript || [], + error: String(error), + }; + } finally { + releaseAgentSlot?.(); + completed += 1; + emitUpdate(); + } + }), + ); + if (signal?.aborted) throw new Error("The local Unsloth agent was cancelled."); + const completedResults = results.filter( + (result): result is LocalAgentResult => Boolean(result), + ); + const succeeded = completedResults.filter((result) => !result.error).length; + const response = + completedResults.length === 1 + ? completedResults[0].error || completedResults[0].response + : [ + `Parallel: ${succeeded}/${tasks.length} local agents succeeded`, + ...completedResults.map( + (result, index) => + `\n### Agent ${index + 1}${result.error ? " failed" : ""}\n\n${result.error || result.response}`, + ), + ].join("\n"); + if (succeeded !== completedResults.length) throw new Error(response); return { - content: [{ type: "text", text: lastResponse || "The local agent returned no text." }], - details: { provider, model }, + content: [{ type: "text", text: response }], + details: details(), }; }, }); diff --git a/unsloth_cli/tests/test_claude_subagent_mcp.py b/unsloth_cli/tests/test_claude_subagent_mcp.py index 13a9bd6255..568dc76ff5 100644 --- a/unsloth_cli/tests/test_claude_subagent_mcp.py +++ b/unsloth_cli/tests/test_claude_subagent_mcp.py @@ -43,6 +43,41 @@ def test_protocol_lists_and_calls_local_agent(): } +def test_protocol_exposes_read_only_agent_for_claude_plan_mode(): + listed = bridge._response( + {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, + run_read_only_agent = lambda task: task, + read_only_tool_name = "unsloth_plan_agent", + ) + tools = {tool["name"]: tool for tool in listed["result"]["tools"]} + assert tools["unsloth_agent"]["annotations"]["readOnlyHint"] is False + assert tools["unsloth_plan_agent"]["annotations"] == { + "readOnlyHint": True, + "destructiveHint": False, + "idempotentHint": True, + "openWorldHint": True, + } + + called = bridge._response( + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "unsloth_plan_agent", + "arguments": {"task": " inspect this "}, + }, + }, + run_agent = lambda task: f"write: {task}", + run_read_only_agent = lambda task: f"plan: {task}", + read_only_tool_name = "unsloth_plan_agent", + ) + assert called["result"] == { + "content": [{"type": "text", "text": "plan: inspect this"}], + "isError": False, + } + + def test_protocol_returns_tool_errors_to_parent(): response = bridge._response( { @@ -212,6 +247,38 @@ def test_local_child_uses_unsloth_without_overwriting_parent_auth( assert "CLAUDE_CODE_OAUTH_TOKEN" not in child_env +def test_read_only_local_child_uses_plan_mode(monkeypatch, tmp_path): + captured = {} + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_MODEL", "unsloth/model-GGUF:Q4_K_M") + monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BYPASS_PERMISSIONS", "1") + monkeypatch.setenv("CLAUDE_PROJECT_DIR", str(tmp_path)) + monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(bridge, "_claude_flags", lambda model: []) + + class Process: + pid = 1234 + returncode = 0 + + def communicate(self, timeout): + return json.dumps({"is_error": False, "result": "PLAN_OK"}), "" + + def poll(self): + return self.returncode + + def popen(command, **kwargs): + captured["command"] = command + return Process() + + monkeypatch.setattr(bridge.subprocess, "Popen", popen) + assert bridge.run_local_agent("plan this", read_only = True) == "PLAN_OK" + command = captured["command"] + assert command[command.index("--permission-mode") + 1] == "plan" + prompt = command[command.index("--append-system-prompt") + 1] + assert "read-only local coding subagent" in prompt + + def test_local_child_process_is_stopped_on_cancellation(monkeypatch, tmp_path): monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_BASE_URL", "http://127.0.0.1:8888") monkeypatch.setenv("UNSLOTH_CLAUDE_SUBAGENT_API_KEY", "sk-unsloth-test") diff --git a/unsloth_cli/tests/test_codex_subagent_mcp.py b/unsloth_cli/tests/test_codex_subagent_mcp.py new file mode 100644 index 0000000000..c0c97ca123 --- /dev/null +++ b/unsloth_cli/tests/test_codex_subagent_mcp.py @@ -0,0 +1,228 @@ +# 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 io +import json +import os +import subprocess + +import pytest + +import unsloth_cli.codex_subagent_mcp as bridge + + +def _write_config(tmp_path, *, bypass_permissions = False): + path = tmp_path / "subagent.json" + path.write_text( + json.dumps( + { + "api_key": "sk-unsloth-test", + "codex_home": str(tmp_path / "child"), + "bypass_permissions": bypass_permissions, + } + ) + ) + return path + + +def test_protocol_uses_codex_specific_tool_name(): + requests = "\n".join( + [ + json.dumps({"jsonrpc": "2.0", "id": 0, "method": "initialize"}), + json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}), + json.dumps( + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": bridge._CODEX_SUBAGENT_MCP_TOOL, + "arguments": {"task": " inspect this "}, + }, + } + ), + ] + ) + output = io.StringIO() + bridge.serve( + io.StringIO(requests), + output, + run_agent = lambda task, cancel_event: f"completed: {task}", + tool_name = bridge._CODEX_SUBAGENT_MCP_TOOL, + tool_description = bridge._CODEX_SUBAGENT_TOOL_DESCRIPTION, + instructions = bridge._SERVER_INSTRUCTIONS, + ) + responses = { + response["id"]: response for response in map(json.loads, output.getvalue().splitlines()) + } + assert responses[0]["result"]["instructions"] == bridge._SERVER_INSTRUCTIONS + assert len(bridge._SERVER_INSTRUCTIONS) <= 512 + assert responses[1]["result"]["tools"][0]["name"] == "spawn_local_agent" + assert ( + "Use this tool instead of the built-in spawn_agent tool" + in responses[1]["result"]["tools"][0]["description"] + ) + assert responses[1]["result"]["tools"][0]["annotations"]["destructiveHint"] is True + assert responses[2]["result"] == { + "content": [{"type": "text", "text": "completed: inspect this"}], + "isError": False, + } + + +@pytest.mark.parametrize("bypass_permissions", [False, True]) +@pytest.mark.parametrize("wsl_bridge", [False, True]) +def test_local_child_uses_explicit_unsloth_profile( + monkeypatch, tmp_path, bypass_permissions, wsl_bridge +): + config = _write_config(tmp_path, bypass_permissions = bypass_permissions) + monkeypatch.setenv(bridge._CODEX_SUBAGENT_CONFIG_ENV, str(config)) + credential_names = ("OPENAI_API_KEY", "CODEX_API_KEY", "CODEX_ACCESS_TOKEN") + for name in credential_names: + monkeypatch.setenv(name, "cloud-key") + monkeypatch.setenv("CODEX_SQLITE_HOME", str(tmp_path / "parent-sqlite")) + if wsl_bridge: + monkeypatch.setattr( + bridge, + "_wsl_shim_env", + lambda command, env, unset: ( + env, + ( + bridge._CODEX_ENV_KEY, + "CODEX_HOME/p", + "CODEX_SQLITE_HOME/p", + *unset, + "PWD/p", + ), + ), + ) + monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = {} + + class Process: + pid = 1234 + returncode = 0 + + def communicate(self, timeout): + captured["timeout"] = timeout + return ( + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "LOCAL_OK"}, + } + ), + "", + ) + + def poll(self): + return self.returncode + + def popen(command, **kwargs): + captured["command"] = command + captured.update(kwargs) + return Process() + + monkeypatch.setattr(bridge.subprocess, "Popen", popen) + assert bridge.run_local_agent("reply exactly LOCAL_OK") == "LOCAL_OK" + command = captured["command"] + assert command[:4] == ["/usr/local/bin/codex", "--oss", "--profile", "unsloth_api"] + if bypass_permissions: + assert "--dangerously-bypass-approvals-and-sandbox" in command + else: + assert command[4:8] == ["--sandbox", "workspace-write", "--ask-for-approval", "never"] + assert command[command.index("exec") + 1 : command.index("exec") + 4] == [ + "--ephemeral", + "--json", + "--skip-git-repo-check", + ] + assert command[-1].endswith("Task: reply exactly LOCAL_OK") + assert captured["cwd"] == os.getcwd() + assert captured["stdin"] is subprocess.DEVNULL + assert captured["stdout"] is subprocess.PIPE + assert captured["stderr"] is subprocess.PIPE + if os.name == "nt": + assert captured["creationflags"] == subprocess.CREATE_NEW_PROCESS_GROUP + else: + assert captured["start_new_session"] is True + assert captured["env"]["CODEX_HOME"] == str(tmp_path / "child") + assert captured["env"]["CODEX_SQLITE_HOME"] == str(tmp_path / "child") + assert captured["env"][bridge._CODEX_ENV_KEY] == "sk-unsloth-test" + if wsl_bridge: + assert all(captured["env"][name] == "" for name in credential_names) + wslenv = captured["env"]["WSLENV"].split(":") + assert all( + name in {entry.split("/", 1)[0] for entry in wslenv} for name in bridge._CODEX_ENV_UNSET + ) + assert "CODEX_SQLITE_HOME/p" in wslenv + assert "PWD/p" in wslenv + else: + assert all(name not in captured["env"] for name in credential_names) + + +def test_local_child_returns_last_agent_message(): + output = "\n".join( + [ + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "intermediate"}, + } + ), + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "final"}, + } + ), + ] + ) + assert bridge._result_text(output) == "final" + + +def test_local_child_prioritizes_failed_turn_over_progress(): + output = "\n".join( + [ + json.dumps( + { + "type": "item.completed", + "item": {"type": "agent_message", "text": "still working"}, + } + ), + json.dumps({"type": "turn.failed", "error": {"message": "local failure"}}), + ] + ) + with pytest.raises(RuntimeError, match = "local failure"): + bridge._result_text(output) + + +def test_local_child_process_is_stopped_on_cancellation(monkeypatch, tmp_path): + config = _write_config(tmp_path) + monkeypatch.setenv(bridge._CODEX_SUBAGENT_CONFIG_ENV, str(config)) + monkeypatch.setattr(bridge.shutil, "which", lambda _: "/usr/local/bin/codex") + cancel_event = bridge.threading.Event() + stopped = [] + + class Process: + pid = 1234 + returncode = None + + def communicate(self, timeout): + cancel_event.set() + raise subprocess.TimeoutExpired("codex", timeout) + + def poll(self): + return self.returncode + + process = Process() + monkeypatch.setattr(bridge.subprocess, "Popen", lambda *args, **kwargs: process) + + def stop(child): + stopped.append(child) + child.returncode = -15 + + monkeypatch.setattr(bridge, "_stop_child", stop) + with pytest.raises(RuntimeError, match = "cancelled"): + bridge.run_local_agent("wait", cancel_event) + assert stopped == [process] diff --git a/unsloth_cli/tests/test_pi_subagent.py b/unsloth_cli/tests/test_pi_subagent.py index beac6770df..0276366656 100644 --- a/unsloth_cli/tests/test_pi_subagent.py +++ b/unsloth_cli/tests/test_pi_subagent.py @@ -64,7 +64,12 @@ import {{ existsSync }} from "node:fs"; import {{ pathToFileURL }} from "node:url"; mock.module("typebox", () => ({{ - Type: {{ Object: (value) => value, String: (value) => value }}, + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, }})); test("cancellation stops the Pi child process group", async () => {{ @@ -138,10 +143,25 @@ def test_pi_child_error_events_fail_the_tool_call(tmp_path): driver = tmp_path / "pi-driver.js" driver.write_text( """ -const event = { - type: "message_end", - message: { role: "assistant", stopReason: "error", errorMessage: "backend unreachable", content: [] }, -}; +const task = process.argv.at(-1).replace(/^Task: /, ""); +const event = task === "pass" + ? { + type: "message_end", + message: { + role: "assistant", + stopReason: "stop", + content: [{ type: "text", text: "PASS_OK" }], + }, + } + : { + type: "message_end", + message: { + role: "assistant", + stopReason: "error", + errorMessage: "backend unreachable", + content: [], + }, + }; console.log(JSON.stringify(event)); """, encoding = "utf-8", @@ -154,7 +174,12 @@ import {{ expect, mock, test }} from "bun:test"; import {{ pathToFileURL }} from "node:url"; mock.module("typebox", () => ({{ - Type: {{ Object: (value) => value, String: (value) => value }}, + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, }})); test("child error events fail the tool call", async () => {{ @@ -168,14 +193,30 @@ test("child error events fail the tool call", async () => {{ registerTool(value) {{ tool = value; }}, }}); - const execution = tool.execute( + const singleExecution = tool.execute( "call", {{ task: "fail" }}, undefined, undefined, {{ cwd: {str(tmp_path)!r} }}, ); - await expect(execution).rejects.toThrow("backend unreachable"); + await expect(singleExecution).rejects.toThrow("backend unreachable"); + + const parallelExecution = tool.execute( + "call", + {{ tasks: ["pass", "fail"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + const parallelError = await parallelExecution.then( + () => "", + (error) => String(error), + ); + expect(parallelError).toContain("Parallel: 1/2 local agents succeeded"); + expect(parallelError).toContain("PASS_OK"); + expect(parallelError).toContain("Agent 2 failed"); + expect(parallelError).toContain("backend unreachable"); }}, 10_000); """, encoding = "utf-8", @@ -189,3 +230,247 @@ test("child error events fail the tool call", async () => {{ ) assert completed.returncode == 0, completed.stdout + completed.stderr + + +@pytest.mark.skipif(os.name == "nt", reason = "POSIX driver script") +def test_pi_parallel_agents_run_together_and_preserve_transcripts(tmp_path): + bun = shutil.which("bun") + if bun is None: + pytest.skip("Bun is required to execute the bundled Pi extension") + + config = tmp_path / "subagent.json" + config.write_text( + json.dumps( + { + "baseUrl": "http://127.0.0.1:8000/v1", + "apiKey": "private-token", + "model": "local-model", + "contextWindow": 32768, + "maxTokens": 8192, + } + ), + encoding = "utf-8", + ) + starts = tmp_path / "starts" + driver = tmp_path / "pi-driver.js" + driver.write_text( + f""" +import * as fs from "node:fs"; + +const task = process.argv.at(-1).replace(/^Task: /, ""); +fs.appendFileSync({str(starts)!r}, `${{task}}\\n`); +for (let attempt = 0; attempt < 100; attempt++) {{ + const count = fs.readFileSync({str(starts)!r}, "utf8").trim().split("\\n").filter(Boolean).length; + if (count >= 2) break; + await Bun.sleep(20); +}} +const event = {{ + type: "message_end", + message: {{ + role: "assistant", + stopReason: "stop", + content: [{{ type: "text", text: `DONE_${{task}}` }}], + }}, +}}; +console.log(JSON.stringify(event)); +console.log(JSON.stringify({{ + type: "tool_execution_end", + toolCallId: `tool_${{task}}`, + toolName: "read", + result: {{ content: [{{ type: "text", text: `TOOL_${{task}}` }}] }}, + isError: false, +}})); +const toolResult = {{ + role: "toolResult", + toolCallId: `tool_${{task}}`, + toolName: "read", + content: [{{ type: "text", text: `TOOL_${{task}}` }}], + isError: false, +}}; +// Current Pi emits a completed tool result both as message_end and in the +// following turn_end. Preserve it once in the transcript. +console.log(JSON.stringify({{ + type: "message_end", + message: toolResult, +}})); +console.log(JSON.stringify({{ + type: "turn_end", + message: event.message, + toolResults: [toolResult], +}})); +""", + encoding = "utf-8", + ) + extension = Path(__file__).parents[1] / "pi_subagent.ts" + test_file = tmp_path / "pi-parallel.test.ts" + test_file.write_text( + f""" +import {{ expect, mock, test }} from "bun:test"; +import {{ pathToFileURL }} from "node:url"; + +mock.module("typebox", () => ({{ + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, +}})); + +test("parallel tasks launch one child each and retain their transcripts", async () => {{ + process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r}; + process.argv[1] = {str(driver)!r}; + + const loaded = await import(pathToFileURL({str(extension)!r}).href); + let tool; + loaded.default({{ + registerProvider() {{}}, + registerTool(value) {{ tool = value; }}, + }}); + + expect(tool.executionMode).toBe("parallel"); + const result = await tool.execute( + "call", + {{ tasks: ["ALPHA", "BETA"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + expect(result.content[0].text).toContain("Parallel: 2/2 local agents succeeded"); + expect(result.content[0].text).toContain("DONE_ALPHA"); + expect(result.content[0].text).toContain("DONE_BETA"); + expect(result.details.mode).toBe("parallel"); + expect(result.details.results).toHaveLength(2); + expect(result.details.results[0].transcript).toHaveLength(2); + expect(result.details.results[1].transcript).toHaveLength(2); + expect(result.details.results[0].transcript[0].content[0].text).toBe("DONE_ALPHA"); + expect(result.details.results[0].transcript[1].content[0].text).toBe("TOOL_ALPHA"); + expect(result.details.results[1].transcript[0].content[0].text).toBe("DONE_BETA"); + expect(result.details.results[1].transcript[1].content[0].text).toBe("TOOL_BETA"); +}}, 10_000); +""", + encoding = "utf-8", + ) + + completed = subprocess.run( + [bun, "test", str(test_file)], + capture_output = True, + text = True, + timeout = 15, + ) + + assert completed.returncode == 0, completed.stdout + completed.stderr + + +@pytest.mark.skipif(os.name == "nt", reason = "POSIX driver script") +def test_pi_parallel_agent_cap_spans_concurrent_tool_calls(tmp_path): + bun = shutil.which("bun") + if bun is None: + pytest.skip("Bun is required to execute the bundled Pi extension") + + config = tmp_path / "subagent.json" + config.write_text( + json.dumps( + { + "baseUrl": "http://127.0.0.1:8000/v1", + "apiKey": "private-token", + "model": "local-model", + "contextWindow": 32768, + "maxTokens": 8192, + } + ), + encoding = "utf-8", + ) + markers = tmp_path / "active" + markers.mkdir() + peaks = tmp_path / "peaks" + driver = tmp_path / "pi-driver.js" + driver.write_text( + f""" +import * as fs from "node:fs"; + +const task = process.argv.at(-1).replace(/^Task: /, ""); +const marker = `{str(markers)!s}/${{process.pid}}`; +fs.writeFileSync(marker, task); +await Bun.sleep(150); +fs.appendFileSync({str(peaks)!r}, `${{fs.readdirSync({str(markers)!r}).length}}\\n`); +await Bun.sleep(150); +fs.unlinkSync(marker); +console.log(JSON.stringify({{ + type: "message_end", + message: {{ + role: "assistant", + stopReason: "stop", + content: [{{ type: "text", text: `DONE_${{task}}` }}], + }}, +}})); +""", + encoding = "utf-8", + ) + extension = Path(__file__).parents[1] / "pi_subagent.ts" + test_file = tmp_path / "pi-global-cap.test.ts" + test_file.write_text( + f""" +import {{ expect, mock, test }} from "bun:test"; +import {{ pathToFileURL }} from "node:url"; + +mock.module("typebox", () => ({{ + Type: {{ + Object: (value) => value, + String: (value) => value, + Optional: (value) => value, + Array: (value) => value, + }}, +}})); + +test("concurrent tool calls share the four-agent cap", async () => {{ + process.env.UNSLOTH_PI_SUBAGENT_CONFIG = {str(config)!r}; + process.argv[1] = {str(driver)!r}; + + const loaded = await import(pathToFileURL({str(extension)!r}).href); + let tool; + loaded.default({{ + registerProvider() {{}}, + registerTool(value) {{ tool = value; }}, + }}); + + const first = tool.execute( + "call-1", + {{ tasks: ["A1", "A2", "A3", "A4"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + const second = tool.execute( + "call-2", + {{ tasks: ["B1", "B2", "B3", "B4"] }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + const results = await Promise.all([first, second]); + expect(results[0].content[0].text).toContain("4/4 local agents succeeded"); + expect(results[1].content[0].text).toContain("4/4 local agents succeeded"); + const afterQueue = await tool.execute( + "call-3", + {{ task: "C" }}, + undefined, + undefined, + {{ cwd: {str(tmp_path)!r} }}, + ); + expect(afterQueue.content[0].text).toContain("DONE_C"); +}}, 10_000); +""", + encoding = "utf-8", + ) + + completed = subprocess.run( + [bun, "test", str(test_file)], + capture_output = True, + text = True, + timeout = 15, + ) + + assert completed.returncode == 0, completed.stdout + completed.stderr + observed = [int(value) for value in peaks.read_text().splitlines()] + assert max(observed) == 4 diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 8ce0bcc0d4..42745ef9c5 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -621,51 +621,223 @@ def test_write_codex_config_omits_catalog_for_old_codex(tmp_path, monkeypatch): assert not (tmp_path / "model-catalog.json").exists() -def test_write_codex_subagent_config_keeps_parent_model_out(tmp_path, monkeypatch): +def test_write_codex_subagent_bridge_keeps_parent_credentials_out(tmp_path, monkeypatch): monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True) local = {**MODEL, "id": MODEL["id"] + ":UD-Q4_K_XL"} - path = start.write_codex_subagent_config(BASE, "private-token", local, tmp_path) - agent = _parse_toml(path.read_text()) - assert agent["name"] == "unsloth" - assert "local agent" in agent["description"].lower() - assert agent["model_provider"] == start._CODEX_PROFILE - assert agent["model"] == local["id"] - assert agent["model_context_window"] == MODEL["context_length"] - assert agent["model_providers"][start._CODEX_PROFILE] == { - "name": "Unsloth Studio", - "base_url": f"{BASE}/v1", - "wire_api": "responses", - "auth": { - "command": sys.executable, - "args": [ - "-c", - "import json,sys; print(json.load(open(sys.argv[1], encoding='utf-8'))['token'])", - str(tmp_path / "unsloth-auth.json"), - ], - "timeout_ms": 5000, - }, + path = start.write_codex_subagent_bridge( + BASE, + "private-token", + local, + tmp_path, + yolo = False, + ) + assert json.loads(path.read_text()) == { + "api_key": "private-token", + "codex_home": str(tmp_path / "child"), + "bypass_permissions": False, } - assert json.loads((tmp_path / "unsloth-auth.json").read_text()) == {"token": "private-token"} - catalog = json.loads((tmp_path / agent["model_catalog_json"]).read_text()) + assert path.stat().st_mode & 0o077 == 0 + profile = _parse_toml((tmp_path / "child" / "unsloth_api.config.toml").read_text()) + assert profile["model"] == local["id"] + assert profile["model_provider"] == start._CODEX_PROFILE + assert profile["model_context_window"] == MODEL["context_length"] + config = _parse_toml((tmp_path / "child" / "config.toml").read_text()) + assert config["model_providers"][start._CODEX_PROFILE]["base_url"] == f"{BASE}/v1" + catalog = json.loads((tmp_path / "child" / profile["model_catalog_json"]).read_text()) assert catalog["models"][0]["slug"] == local["id"] +def test_write_codex_parent_overlay_preserves_user_state_and_instructions(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + (source / "config.toml").write_text('model = "cloud-model"\n') + (source / "auth.json").write_text('{"auth": "cloud"}\n') + (source / "sessions").mkdir() + (source / "AGENTS.override.md").write_text("Keep my existing instructions.\n") + monkeypatch.setenv("CODEX_HOME", str(source)) + + overlay = start.write_codex_parent_overlay(tmp_path / "managed" / "parent") + + assert (overlay / "config.toml").read_text() == 'model = "cloud-model"\n' + assert (overlay / "auth.json").read_text() == '{"auth": "cloud"}\n' + assert (overlay / "sessions").is_dir() + instructions = (overlay / "AGENTS.override.md").read_text() + assert instructions.startswith("Keep my existing instructions.\n") + assert start._CODEX_SUBAGENT_ROUTING_INSTRUCTIONS in instructions + assert not (overlay / "AGENTS.md").exists() + assert (overlay / "AGENTS.override.md").stat().st_mode & 0o077 == 0 + assert (source / "AGENTS.override.md").read_text() == "Keep my existing instructions.\n" + + +def test_write_codex_parent_overlay_refreshes_reused_entries(tmp_path, monkeypatch): + first = tmp_path / "first-codex" + first.mkdir() + (first / "auth.json").write_text('{"auth": "old"}\n') + (first / "old-only.toml").write_text("old\n") + second = tmp_path / "second-codex" + second.mkdir() + (second / "auth.json").write_text('{"auth": "new"}\n') + overlay_path = tmp_path / "managed" / "parent" + + monkeypatch.setenv("CODEX_HOME", str(first)) + overlay = start.write_codex_parent_overlay(overlay_path) + assert (overlay / "auth.json").read_text() == '{"auth": "old"}\n' + assert (overlay / "old-only.toml").exists() + + monkeypatch.setenv("CODEX_HOME", str(second)) + overlay = start.write_codex_parent_overlay(overlay_path) + assert (overlay / "auth.json").read_text() == '{"auth": "new"}\n' + assert not (overlay / "old-only.toml").exists() + + +def test_write_codex_parent_overlay_does_not_use_itself_as_source(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + (source / "auth.json").write_text('{"auth": "cloud"}\n') + overlay_path = tmp_path / "managed" / "parent" + monkeypatch.setenv("CODEX_HOME", str(source)) + overlay = start.write_codex_parent_overlay(overlay_path) + + monkeypatch.setenv("CODEX_HOME", str(overlay)) + overlay = start.write_codex_parent_overlay(overlay_path) + + assert (overlay / "auth.json").read_text() == '{"auth": "cloud"}\n' + manifest = json.loads((overlay / start._CODEX_PARENT_OVERLAY_MANIFEST).read_text()) + assert manifest["source_home"] == str(source) + + +def test_write_codex_parent_overlay_refreshes_fallback_copies(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + config = source / "config.toml" + config.write_text('model = "first"\n') + sessions = source / "sessions" + sessions.mkdir() + (sessions / "existing.jsonl").write_text("existing session\n") + monkeypatch.setenv("CODEX_HOME", str(source)) + + def deny_symlink(*args, **kwargs): + raise OSError("symlinks unavailable") + + monkeypatch.setattr(Path, "symlink_to", deny_symlink) + monkeypatch.setattr(start, "_create_directory_junction", lambda source, target: False) + overlay = start.write_codex_parent_overlay(tmp_path / "managed" / "parent") + (overlay / "history.jsonl").write_text("session state\n") + config.write_text('model = "second"\n') + + overlay = start.write_codex_parent_overlay(overlay) + + assert (overlay / "config.toml").read_text() == 'model = "second"\n' + assert (overlay / "sessions" / "existing.jsonl").read_text() == "existing session\n" + assert (overlay / "history.jsonl").read_text() == "session state\n" + + config.unlink() + overlay = start.write_codex_parent_overlay(overlay) + assert not (overlay / "config.toml").exists() + assert (overlay / "history.jsonl").read_text() == "session state\n" + + +def test_create_directory_junction_uses_windows_mklink(tmp_path, monkeypatch): + captured = {} + monkeypatch.setattr(start.os, "name", "nt") + + def run(command, **kwargs): + captured["command"] = command + captured["kwargs"] = kwargs + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + source = tmp_path / "source" + target = tmp_path / "target" + + assert start._create_directory_junction(source, target) is True + assert captured["command"] == [ + "cmd.exe", + "/d", + "/c", + "mklink", + "/J", + str(target), + str(source), + ] + assert captured["kwargs"] == { + "capture_output": True, + "text": True, + "timeout": 30, + "check": False, + } + + @pytest.mark.skipif(os.name == "nt", reason = "WSL scenario") -def test_codex_subagent_auth_uses_wsl_for_windows_codex(monkeypatch, tmp_path): +def test_write_codex_parent_overlay_uses_windows_home_for_windows_codex(tmp_path, monkeypatch): + windows_profile = tmp_path / "windows-profile" + source = windows_profile / ".codex" + source.mkdir(parents = True) + (source / "auth.json").write_text('{"auth": "windows"}\n') + executable = "/mnt/c/Users/x/AppData/Roaming/npm/codex" + monkeypatch.delenv("CODEX_HOME", raising = False) + monkeypatch.delenv("USERPROFILE", raising = False) + monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") + monkeypatch.setattr(start.shutil, "which", lambda _: executable) + + def check_output(command, **kwargs): + if command[0] == "cmd.exe": + assert kwargs["cwd"] == str(Path(executable).parent) + return r"C:\Users\x" + "\n" + assert command == ["wslpath", "-u", r"C:\Users\x"] + return str(windows_profile) + "\n" + + monkeypatch.setattr(start.subprocess, "check_output", check_output) + + overlay = start.write_codex_parent_overlay(tmp_path / "managed" / "parent") + + assert (overlay / "auth.json").read_text() == '{"auth": "windows"}\n' + + +def test_codex_parent_overlay_launch_uses_private_temp_root_and_cleans_up(tmp_path, monkeypatch): + source = tmp_path / "user-codex" + source.mkdir() + (source / "auth.json").write_text("{}\n") + monkeypatch.setenv("CODEX_HOME", str(source)) + agents_root = tmp_path / "agents" + monkeypatch.setattr(start, "_agents_config_root", lambda: agents_root) + + with start._codex_parent_overlay(tmp_path / "session", launch = True, persist = False) as overlay: + assert overlay.parent == agents_root / ".tmp" + assert start._CODEX_SUBAGENT_ROUTING_INSTRUCTIONS in (overlay / "AGENTS.md").read_text() + assert overlay.exists() + + assert not overlay.exists() + + +@pytest.mark.skipif(os.name == "nt", reason = "WSL scenario") +def test_codex_subagent_bridge_uses_wsl_for_windows_codex(monkeypatch, tmp_path): monkeypatch.setenv("WSL_DISTRO_NAME", "Ubuntu") - monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: False) monkeypatch.setattr( start.shutil, "which", lambda _: "/mnt/c/Users/x/AppData/Roaming/npm/codex.exe", ) - - path = start.write_codex_subagent_config(BASE, "private-token", MODEL, tmp_path) - auth = _parse_toml(path.read_text())["model_providers"][start._CODEX_PROFILE]["auth"] - - assert auth["command"] == "wsl.exe" - assert auth["args"][:5] == ["-d", "Ubuntu", "--", sys.executable, "-c"] - assert auth["args"][-1] == str(tmp_path / "unsloth-auth.json") + flags = start._codex_subagent_flags(tmp_path / "subagent.json") + prefix = f"mcp_servers.{start._CODEX_SUBAGENT_MCP_SERVER}=" + override = next(value for value in flags if value.startswith(prefix)) + server = _parse_toml("server = " + override.removeprefix(prefix))["server"] + assert server["command"] == "wsl.exe" + assert server["args"] == [ + "-d", + "Ubuntu", + "--", + sys.executable, + "-c", + server["args"][5], + str(tmp_path / "subagent.json"), + ] + assert "sys.path.insert" in server["args"][5] + assert f"from {start._CODEX_SUBAGENT_MCP_MODULE} import main" in server["args"][5] + assert server["required"] is True + assert server["enabled_tools"] == [start._CODEX_SUBAGENT_MCP_TOOL] + assert server["default_tools_approval_mode"] == "approve" + assert not any(value.startswith("developer_instructions=") for value in flags) @pytest.mark.skipif(os.name == "nt", reason = "WSL scenario") @@ -813,7 +985,7 @@ def test_connect_claude_as_subagent_preserves_cloud_parent(fake_studio, tmp_path "--plugin-dir", str(plugin), "--allowedTools", - start._CLAUDE_SUBAGENT_TOOL, + f"{start._CLAUDE_SUBAGENT_TOOL},{start._CLAUDE_SUBAGENT_PLAN_TOOL}", "hello", ] assert "--model" not in command @@ -841,6 +1013,7 @@ def test_connect_claude_as_subagent_preserves_cloud_parent(fake_studio, tmp_path } skill = (plugin / "skills" / "local-agent" / "SKILL.md").read_text() assert "spawn an Unsloth agent or local agent" in skill + assert "In plan mode" in skill assert "Ask Claude to spawn an Unsloth or local agent." in result.output @@ -1016,6 +1189,11 @@ def test_connect_codex_no_launch(fake_studio, tmp_path): def test_connect_codex_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, monkeypatch): monkeypatch.setattr(start, "_codex_supports_model_catalog", lambda: True) + source_home = tmp_path / "user-codex" + source_home.mkdir() + (source_home / "config.toml").write_text('model = "cloud-model"\n') + (source_home / "AGENTS.md").write_text("Keep the user's guidance.\n") + monkeypatch.setenv("CODEX_HOME", str(source_home)) result = CliRunner().invoke( start.start_app, [ @@ -1029,20 +1207,34 @@ def test_connect_codex_as_subagent_preserves_cloud_parent(fake_studio, tmp_path, assert result.exit_code == 0, result.output command = _launch_command(result.output) assert command[0] == "codex" - assert command[1:3] == ["--enable", "multi_agent"] - assert "agents.max_depth=1" in command assert "--oss" not in command assert "--profile" not in command assert "--model" not in command - assert "CODEX_HOME" not in result.output + parent_home = tmp_path / "agents" / "codex-subagent" / "parent" + _assert_env_set(result.output, "CODEX_HOME", str(parent_home)) assert start._CODEX_ENV_KEY not in result.output assert "sk-unsloth-feedfacefeedface" not in result.output home = tmp_path / "agents" / "codex-subagent" - agent_path = home / "unsloth.toml" - agent = _parse_toml(agent_path.read_text()) - assert agent["model"] == MODEL["id"] + ":UD-Q4_K_XL" - assert "env_key" not in agent["model_providers"][start._CODEX_PROFILE] - assert f"agents.unsloth.config_file={json.dumps(str(agent_path))}" in command + bridge_path = home / "subagent.json" + bridge = json.loads(bridge_path.read_text()) + assert bridge["api_key"] == "sk-unsloth-feedfacefeedface" + assert bridge["codex_home"] == str(home / "child") + assert bridge["bypass_permissions"] is False + profile = _parse_toml((home / "child" / "unsloth_api.config.toml").read_text()) + assert profile["model"] == MODEL["id"] + ":UD-Q4_K_XL" + prefix = f"mcp_servers.{start._CODEX_SUBAGENT_MCP_SERVER}=" + override = next(value for value in command if value.startswith(prefix)) + assert override.startswith(prefix) + server = _parse_toml("server = " + override.removeprefix(prefix))["server"] + assert server["command"] == sys.executable + assert server["args"] == ["-c", server["args"][1], str(bridge_path)] + assert "sys.path.insert" in server["args"][1] + assert f"from {start._CODEX_SUBAGENT_MCP_MODULE} import main" in server["args"][1] + assert server["enabled_tools"] == [start._CODEX_SUBAGENT_MCP_TOOL] + assert not any(value.startswith("developer_instructions=") for value in command) + parent_instructions = (parent_home / "AGENTS.md").read_text() + assert parent_instructions.startswith("Keep the user's guidance.\n") + assert start._CODEX_SUBAGENT_ROUTING_INSTRUCTIONS in parent_instructions assert "Ask Codex to spawn an Unsloth or local agent." in result.output From 629cc50f1a7632c250aac90ec65ada8bf5974d58 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 20:49:54 -0700 Subject: [PATCH 015/161] Unsloth run/start: per-model recommended sampling and override flags (#7335) Seed each request with the model's recommended sampling (matching the Chat UI), add per-field override flags, ignore oversized overrides, warn when sampling pins cannot apply to a reused server, and apply pins to the completions endpoint. --- studio/backend/routes/inference.py | 108 ++++++- .../backend/tests/test_audio_sampling_fill.py | 90 ++++++ .../backend/tests/test_sampling_resolution.py | 270 ++++++++++++++++++ .../utils/inference/inference_config.py | 137 +++++++++ unsloth_cli/commands/start.py | 201 ++++++++++++- unsloth_cli/commands/studio.py | 69 ++++- unsloth_cli/tests/test_start.py | 112 ++++++++ .../tests/test_studio_run_parallel_flag.py | 43 ++- 8 files changed, 1014 insertions(+), 16 deletions(-) create mode 100644 studio/backend/tests/test_audio_sampling_fill.py create mode 100644 studio/backend/tests/test_sampling_resolution.py diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b6bbbdfb5f..389a0d245a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6061,6 +6061,7 @@ async def generate_audio( # Advertised repo id after an auto-switch load, else a clean public id, # never the absolute .gguf path. model_name = _llama_public_model_id(llama_backend) + _audio_model_id = getattr(llama_backend, "model_identifier", None) or model_name gen = lambda: llama_backend.generate_audio_response( text = text, audio_type = llama_backend._audio_type, @@ -6079,6 +6080,7 @@ async def generate_audio( if not model_info.get("is_audio"): raise HTTPException(status_code = 400, detail = "Active model is not an audio model.") model_name = public_model_id(backend.active_model_name) + _audio_model_id = getattr(backend, "active_model_name", None) or model_name gen = lambda: backend.generate_audio_response( text = text, temperature = payload.temperature, @@ -6090,6 +6092,13 @@ async def generate_audio( use_adapter = payload.use_adapter, ) + # Apply per-model recommended sampling + any operator UNSLOTH_SAMPLING_* pin before + # generating, so `unsloth run --temperature` (and the other pins) and per-model + # recommendations reach audio (TTS) generation too, not just chat. The gen lambdas read + # payload.* lazily at call time, so filling here takes effect; this covers both the direct + # /audio/generate route and the chat-completions audio branches that delegate here. + _fill_recommended_sampling_openai(payload, _audio_model_id) + try: wav_bytes, sample_rate = await asyncio.to_thread(gen) except Exception as e: @@ -7381,6 +7390,51 @@ async def delete_openai_container( await client.close() +def _fill_recommended_sampling_openai(payload, model_id) -> None: + """Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to a + ChatCompletionRequest in place. + + Only the sampling fields the client did NOT explicitly send (tracked via + ``model_fields_set``) are overwritten, so a client that sets a field stays byte-identical + unless an operator pins it. Fields with neither a recommendation nor a pin keep their + existing (schema-default) value. + """ + from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES + + explicit = { + f: (getattr(payload, f) if f in payload.model_fields_set else None) + for f in SAMPLING_FIELD_NAMES + } + effective = resolve_effective_sampling(model_id, explicit) + for field, value in effective.items(): + setattr(payload, field, value) + + +# /v1/completions is proxied to llama-server verbatim; its repetition knob is "repeat_penalty", +# and every other sampling field keeps its name (mirrors _build_passthrough_payload). +_COMPLETIONS_SAMPLING_BODY_KEY = {"repetition_penalty": "repeat_penalty"} + + +def _fill_recommended_sampling_completions(body: dict, model_id) -> None: + """Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to a raw + ``/v1/completions`` body in place, so the legacy (non-chat) endpoint honors the same pins as + ``/v1/chat/completions``. + + Unlike :func:`_fill_recommended_sampling_openai`, which fills a ChatCompletionRequest whose + schema already carries per-field defaults, this body is proxied to llama-server as-is. A field + with no operator pin, client value, or per-model recommendation is therefore left untouched + (``fill_defaults = False``) so llama-server keeps its own default rather than being forced onto + this schema's value. llama-server names the repetition knob ``repeat_penalty``, so read and + write that alias for the client-sent value and any pin. + """ + from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES + + explicit = {f: body.get(_COMPLETIONS_SAMPLING_BODY_KEY.get(f, f)) for f in SAMPLING_FIELD_NAMES} + effective = resolve_effective_sampling(model_id, explicit, fill_defaults = False) + for field, value in effective.items(): + body[_COMPLETIONS_SAMPLING_BODY_KEY.get(field, field)] = value + + @router.post("/chat/completions") async def openai_chat_completions( payload: ChatCompletionRequest, @@ -7710,6 +7764,13 @@ async def openai_chat_completions( completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) + # Apply recommended sampling + operator pins to the omitted fields before generating, + # so audio-input (non-whisper) generation honors `unsloth run --temperature` and + # per-model recommendations like chat does. Whisper (ASR) ignores these fields. + _fill_recommended_sampling_openai( + payload, getattr(backend, "active_model_name", None) or model_name + ) + def audio_input_generate(): if model_info.get("audio_type") == "whisper": return backend.generate_whisper_response( @@ -7853,6 +7914,18 @@ async def openai_chat_completions( ), ) + # Apply per-model recommended sampling (and any operator UNSLOTH_SAMPLING_* pin) to the + # fields the client omitted, so agents and API clients get the model's tuned defaults + # unless they set the field explicitly. Placed after external-provider routing (which + # returned above) so only local llama-server / transformers requests are touched, and it + # covers both the passthrough and non-passthrough branches below since both read payload.*. + _reco_model_id = ( + getattr(llama_backend, "model_identifier", None) + if using_gguf + else getattr(backend, "active_model_name", None) + ) or model_name + _fill_recommended_sampling_openai(payload, _reco_model_id) + # ── Standard OpenAI function-calling pass-through (GGUF only) ──── # When a client (opencode / Claude Code via OpenAI compat / Cursor / # Continue / ...) sends standard OpenAI `tools` without Unsloth's @@ -10630,6 +10703,10 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if _resolved_max_tokens is not None else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR) ) + # Apply per-model recommended sampling and any operator UNSLOTH_SAMPLING_* pin to the raw + # body so /v1/completions honors the same pins as /v1/chat/completions; it is otherwise a + # verbatim proxy that would keep llama-server's defaults for every omitted sampling field. + _fill_recommended_sampling_completions(body, getattr(llama_backend, "model_identifier", None)) target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) prompt_text = _flatten_monitor_prompt(body.get("prompt", "")) @@ -11588,6 +11665,9 @@ async def _responses_stream( detail = "Image provided but current GGUF model does not support vision.", ) + # Streaming /v1/responses builds the passthrough body directly (bypassing + # openai_chat_completions), so apply recommended sampling here too. + _fill_recommended_sampling_openai(chat_req, getattr(llama_backend, "model_identifier", None)) body = _build_openai_passthrough_body( chat_req, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) @@ -13019,14 +13099,28 @@ async def anthropic_messages( # endpoint matches /v1/chat/completions. _has_image = _normalize_anthropic_openai_images(openai_messages, llama_backend.is_vision) - temperature = payload.temperature if payload.temperature is not None else 0.6 - top_p = payload.top_p if payload.top_p is not None else 0.95 - top_k = payload.top_k if payload.top_k is not None else 20 - min_p = payload.min_p if payload.min_p is not None else 0.01 - repetition_penalty = ( - payload.repetition_penalty if payload.repetition_penalty is not None else 1.0 + # Fill omitted sampling fields with the per-model recommendation (or an operator + # UNSLOTH_SAMPLING_* pin); an explicit client value wins unless the operator pinned it. + # Anthropic sampling fields are Optional, so None already marks "client omitted". + from utils.inference.inference_config import resolve_effective_sampling + + _anthropic_sampling = resolve_effective_sampling( + getattr(llama_backend, "model_identifier", None) or model_name, + { + "temperature": payload.temperature, + "top_p": payload.top_p, + "top_k": payload.top_k, + "min_p": payload.min_p, + "repetition_penalty": payload.repetition_penalty, + "presence_penalty": payload.presence_penalty, + }, ) - presence_penalty = payload.presence_penalty if payload.presence_penalty is not None else 0.0 + temperature = _anthropic_sampling["temperature"] + top_p = _anthropic_sampling["top_p"] + top_k = _anthropic_sampling["top_k"] + min_p = _anthropic_sampling["min_p"] + repetition_penalty = _anthropic_sampling["repetition_penalty"] + presence_penalty = _anthropic_sampling["presence_penalty"] stop = payload.stop_sequences or None # Translate Anthropic tool_choice to OpenAI format for llama-server. Falls diff --git a/studio/backend/tests/test_audio_sampling_fill.py b/studio/backend/tests/test_audio_sampling_fill.py new file mode 100644 index 0000000000..efea18b83e --- /dev/null +++ b/studio/backend/tests/test_audio_sampling_fill.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Audio (TTS) generation applies recommended sampling + operator pins, like chat. + +Regression guard for the fix that moved the sampling fill ahead of the audio generators: a +prior version resolved sampling only after the audio branches returned, so `unsloth run +--temperature` (UNSLOTH_SAMPLING_*) and per-model recommendations never reached audio +generation. These exercise the transformers TTS path of ``generate_audio`` (the direct +``/audio/generate`` route, which the chat-completions audio branches also delegate to). +""" + +import asyncio + +import pytest + +import routes.inference as inference_route +from models.inference import ChatCompletionRequest +from utils.inference import inference_config as ic + + +class _FakeLlama: + # is_loaded False forces the transformers (non-GGUF) TTS branch in generate_audio. + is_loaded = False + _is_audio = False + + +class _FakeTransformersBackend: + def __init__(self): + self.active_model_name = "some/custom-tts" + self.models = {"some/custom-tts": {"is_audio": True}} + self.captured = {} + + def generate_audio_response(self, **kwargs): + self.captured.update(kwargs) + return (b"RIFFfake", 24000) + + +@pytest.fixture(autouse = True) +def _isolate(monkeypatch): + ic._recommended_sampling.cache_clear() + for field in ic.SAMPLING_FIELD_NAMES: + monkeypatch.delenv(ic._SAMPLING_FIELDS[field][0], raising = False) + yield + ic._recommended_sampling.cache_clear() + + +def _run_generate_audio( + monkeypatch, + *, + recommended = None, + temperature = None, +): + backend = _FakeTransformersBackend() + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _FakeLlama()) + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: backend) + + async def _noop_switch(*a, **k): + return None + + monkeypatch.setattr(inference_route, "_maybe_auto_switch_model", _noop_switch) + + # Recommendation source == the Chat UI's .inference block. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(recommended or {})) + ic._recommended_sampling.cache_clear() + + kwargs = {"model": "some/custom-tts", "messages": [{"role": "user", "content": "hi"}]} + if temperature is not None: + kwargs["temperature"] = temperature + payload = ChatCompletionRequest(**kwargs) + + asyncio.run(inference_route.generate_audio(payload, request = None, current_subject = "t")) + return backend.captured + + +def test_audio_uses_recommended_sampling_when_omitted(monkeypatch): + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0, "top_k": 64}) + assert captured["temperature"] == 1.0 + assert captured["top_k"] == 64 + + +def test_audio_operator_pin_overrides_client(monkeypatch): + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0}, temperature = 0.2) + assert captured["temperature"] == 0.9 # operator pin wins even over an explicit client value + + +def test_audio_client_explicit_preserved(monkeypatch): + captured = _run_generate_audio(monkeypatch, recommended = {"temperature": 1.0}, temperature = 0.2) + assert captured["temperature"] == 0.2 # explicit client value preserved over recommendation diff --git a/studio/backend/tests/test_sampling_resolution.py b/studio/backend/tests/test_sampling_resolution.py new file mode 100644 index 0000000000..1ebbae2502 --- /dev/null +++ b/studio/backend/tests/test_sampling_resolution.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Effective sampling resolution: per-model recommendation + operator pins. + +Precedence per field: operator UNSLOTH_SAMPLING_* pin -> client explicit value -> +per-model recommendation (load_inference_config) -> static schema default. +""" + +import pytest + +from utils.inference.inference_config import resolve_effective_sampling, SAMPLING_FIELD_NAMES +from utils.inference import inference_config as ic + +_SCHEMA_DEFAULTS = { + "temperature": 0.6, + "top_p": 0.95, + "top_k": 20, + "min_p": 0.01, + "repetition_penalty": 1.0, + "presence_penalty": 0.0, +} + + +@pytest.fixture(autouse = True) +def _isolate(monkeypatch): + # The recommended lookup is lru-cached; clear it so a patched config takes effect. + ic._recommended_sampling.cache_clear() + for field in SAMPLING_FIELD_NAMES: + monkeypatch.delenv(ic._SAMPLING_FIELDS[field][0], raising = False) + yield + ic._recommended_sampling.cache_clear() + + +def _all_omitted(): + return {f: None for f in SAMPLING_FIELD_NAMES} + + +def _set_recommended(monkeypatch, mapping): + # _recommended_sampling sources from load_inference_config -- the exact block the Chat UI + # seeds from -- so patch that directly. Fields absent from `mapping` fall to schema defaults. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(mapping)) + ic._recommended_sampling.cache_clear() + + +def test_recommended_applies_when_client_omits(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 1.0 + assert eff["top_k"] == 64 + assert eff["min_p"] == 0.0 + # A field with no recommendation keeps the static schema default. + assert eff["top_p"] == 0.95 + + +def test_client_explicit_beats_recommended(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0}) + eff = resolve_effective_sampling("some/model", {**_all_omitted(), "temperature": 0.2}) + assert eff["temperature"] == 0.2 + + +def test_operator_pin_beats_client_and_recommended(monkeypatch): + _set_recommended(monkeypatch, {"temperature": 1.0}) + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + eff = resolve_effective_sampling("some/model", {**_all_omitted(), "temperature": 0.2}) + assert eff["temperature"] == 0.9 + + +def test_unknown_model_matches_ui_inference_block(monkeypatch): + # An unknown model gets the same values the Chat UI would seed (load_inference_config's + # default.yaml fallback: temp 0.7 / top_k -1), NOT the request schema defaults. + ui_block = { + "temperature": 0.7, + "top_p": 0.95, + "top_k": -1, + "min_p": 0.01, + "presence_penalty": 0.0, + "repetition_penalty": 1.0, + } + monkeypatch.setattr(ic, "load_inference_config", lambda mid: dict(ui_block)) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/unknown-model", _all_omitted()) + assert eff["temperature"] == 0.7 + assert eff["top_k"] == -1 + assert eff["min_p"] == 0.01 + + +def test_empty_recommendation_falls_back_to_schema_defaults(monkeypatch): + # If load_inference_config yields nothing usable, the resolver falls back to the request + # schema defaults. + monkeypatch.setattr(ic, "load_inference_config", lambda mid: {}) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff == _SCHEMA_DEFAULTS + + +@pytest.mark.parametrize( + "model", + ["unsloth/gemma-4-E4B", "unsloth/Qwen3-4B", "unsloth/Qwen3.5-9B", "someorg/unknown-xyz"], +) +def test_recommendation_matches_ui_source(model): + # Parity guard: what the server recommends for omitted fields equals the Chat UI's source + # (load_inference_config) for every field the UI adopts (mergeBackendRecommendedInference). + ic._recommended_sampling.cache_clear() + ui = ic.load_inference_config(model) + rec = ic._recommended_sampling(model) + for f in ic._UI_RECOMMENDED_FIELDS: + cleaned = ic._clean_sampling_value(f, ui.get(f)) + if cleaned is not None: + assert rec.get(f) == cleaned, f"{model}:{f} rec={rec.get(f)} ui={ui.get(f)}" + + +def test_repetition_penalty_not_auto_recommended(monkeypatch): + # The Chat UI's mergeBackendRecommendedInference never adopts a backend repetition_penalty + # (e.g. lfm2's family value 1.05), so the server must not auto-apply one either. It stays at + # the schema default unless the client sends it or an operator pins it. + monkeypatch.setattr( + ic, "load_inference_config", lambda mid: {"temperature": 0.7, "repetition_penalty": 1.05} + ) + ic._recommended_sampling.cache_clear() + eff = resolve_effective_sampling("some/lfm2-model", _all_omitted()) + assert eff["temperature"] == 0.7 # a UI-adopted field is recommended + assert eff["repetition_penalty"] == 1.0 # rep is NOT auto-recommended (matches the UI) + # An operator can still pin it explicitly. + monkeypatch.setenv("UNSLOTH_SAMPLING_REPETITION_PENALTY", "1.05") + eff2 = resolve_effective_sampling("some/lfm2-model", _all_omitted()) + assert eff2["repetition_penalty"] == 1.05 + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("0.5", 0.5), + ("abc", None), # unparseable + ("9.0", None), # above temperature max (2.0) + ("-1", None), # below temperature min (0.0) + (" ", None), # blank + ("nan", None), # NaN would pass a naive range check + ("inf", None), # non-finite + ("-inf", None), # non-finite + ], +) +def test_operator_override_parsing(monkeypatch, raw, expected): + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", raw) + assert ic._operator_sampling_override("temperature") == expected + + +def test_out_of_range_recommendation_is_dropped(monkeypatch): + # A malformed model recommendation (out of range) is ignored, so the request keeps the + # schema default rather than forwarding a bad value to llama-server. + _set_recommended(monkeypatch, {"temperature": 5.0, "top_k": 64}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 0.6 # 5.0 is outside [0, 2] -> schema default + assert eff["top_k"] == 64 # a valid recommendation is still applied + + +def test_operator_override_top_k_int_and_range(monkeypatch): + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "40") + assert ic._operator_sampling_override("top_k") == 40 + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "200") # above max 100 + assert ic._operator_sampling_override("top_k") is None + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "-1") # min allowed + assert ic._operator_sampling_override("top_k") == -1 + + +@pytest.mark.parametrize( + "field, val", + [ + ("top_k", 10**400), # oversized int on an int field: int() ok, but math.isfinite raises + ("top_k", float("nan")), # NaN reaching an int field: int(nan) raises ValueError + ("top_k", float("inf")), # inf reaching an int field: int(inf) raises OverflowError + ( + "temperature", + 10**400, + ), # oversized int on a float field: float(huge_int) raises OverflowError + ], +) +def test_clean_sampling_value_rejects_unrepresentable(field, val): + # None of these may raise; each is unusable and must be dropped to None (regression: an + # oversized value used to raise OverflowError before the range check could drop it). + assert ic._clean_sampling_value(field, val) is None + + +def test_oversized_operator_override_ignored(monkeypatch): + # A huge integer string parses via int() but overflows float(); math.isfinite would raise + # OverflowError and 500 the request. It must be ignored like any other bad override and the + # field must fall back to the schema default -- no exception. + monkeypatch.setenv("UNSLOTH_SAMPLING_TOP_K", "9" * 400) + assert ic._operator_sampling_override("top_k") is None + _set_recommended(monkeypatch, {}) # no per-model recommendation -> schema default applies + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["top_k"] == 20 # schema default, resolved without raising + + +def test_oversized_recommendation_ignored(monkeypatch): + # A malformed per-model recommendation carrying an oversized int must not raise while + # resolving either; the field simply falls back to the schema default. + _set_recommended(monkeypatch, {"temperature": 10**400, "top_k": 64}) + eff = resolve_effective_sampling("some/model", _all_omitted()) + assert eff["temperature"] == 0.6 # oversized -> dropped -> schema default + assert eff["top_k"] == 64 # a valid recommendation is still applied + + +def test_fill_recommended_sampling_openai_payload(monkeypatch): + from models.inference import ChatCompletionRequest + from routes.inference import _fill_recommended_sampling_openai + + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + + # Client sent only temperature; top_k / min_p were omitted. + payload = ChatCompletionRequest( + model = "m", messages = [{"role": "user", "content": "hi"}], temperature = 0.2 + ) + _fill_recommended_sampling_openai(payload, "some/model") + assert payload.temperature == 0.2 # explicit client value preserved + assert payload.top_k == 64 # recommended fills the omitted field + assert payload.min_p == 0.0 + assert payload.top_p == 0.95 # no recommendation -> schema default unchanged + + +def test_fill_recommended_sampling_openai_operator_pin_overrides_client(monkeypatch): + from models.inference import ChatCompletionRequest + from routes.inference import _fill_recommended_sampling_openai + + monkeypatch.setattr(ic, "load_model_defaults", lambda mid: {}) + monkeypatch.setattr(ic, "get_family_inference_params", lambda mid: {}) + ic._recommended_sampling.cache_clear() + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + + payload = ChatCompletionRequest( + model = "m", messages = [{"role": "user", "content": "hi"}], temperature = 0.2 + ) + _fill_recommended_sampling_openai(payload, "some/model") + assert payload.temperature == 0.9 # operator pin wins even over an explicit client value + + +def test_fill_recommended_sampling_completions_body(monkeypatch): + # /v1/completions is a raw proxy: recommendations fill omitted fields, but a field with no + # recommendation and no pin is left absent so llama-server keeps its own default (unlike the + # chat schema, which carries per-field defaults). + from routes.inference import _fill_recommended_sampling_completions + + _set_recommended(monkeypatch, {"temperature": 1.0, "top_k": 64, "min_p": 0.0}) + + body = {"prompt": "hi", "temperature": 0.2} + _fill_recommended_sampling_completions(body, "some/model") + assert body["temperature"] == 0.2 # explicit client value preserved + assert body["top_k"] == 64 # recommendation fills the omitted field + assert body["min_p"] == 0.0 + # No recommendation and no pin -> NOT injected (llama-server keeps its default). + assert "top_p" not in body + assert "presence_penalty" not in body + assert "repeat_penalty" not in body + + +def test_fill_recommended_sampling_completions_operator_pin(monkeypatch): + # An operator pin overrides the client's raw-body value, and the repetition pin is written + # under llama-server's "repeat_penalty" key (the schema field is repetition_penalty). + from routes.inference import _fill_recommended_sampling_completions + + monkeypatch.setattr(ic, "load_inference_config", lambda mid: {}) + ic._recommended_sampling.cache_clear() + monkeypatch.setenv("UNSLOTH_SAMPLING_TEMPERATURE", "0.9") + monkeypatch.setenv("UNSLOTH_SAMPLING_REPETITION_PENALTY", "1.2") + + body = {"prompt": "hi", "temperature": 0.2, "repeat_penalty": 1.05} + _fill_recommended_sampling_completions(body, "some/model") + assert body["temperature"] == 0.9 # operator pin wins over the client's explicit value + assert body["repeat_penalty"] == 1.2 # repetition pin lands on llama-server's key + assert "repetition_penalty" not in body # never leak the schema field name into the body diff --git a/studio/backend/utils/inference/inference_config.py b/studio/backend/utils/inference/inference_config.py index 05eb08067c..a264e06c85 100644 --- a/studio/backend/utils/inference/inference_config.py +++ b/studio/backend/utils/inference/inference_config.py @@ -5,7 +5,10 @@ from pathlib import Path from typing import Dict, Any, Optional +from functools import lru_cache import json +import math +import os import yaml import structlog from loggers import get_logger @@ -160,3 +163,137 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]: } return inference_config + + +# ── Effective sampling resolution for `unsloth run` / `unsloth start` ────────── +# +# Per-model recommended sampling is applied to a request only for the fields the +# client omitted; an operator can pin a field from the CLI via UNSLOTH_SAMPLING_* +# (a hard override that wins even over an explicit client value). Precedence per +# field: operator pin -> client explicit -> per-model recommendation -> the static +# schema default (mirroring ChatCompletionRequest, so behavior is unchanged when +# nothing is recommended or pinned). + +# field -> (env var, static default, min, max, is_int) +_SAMPLING_FIELDS = { + "temperature": ("UNSLOTH_SAMPLING_TEMPERATURE", 0.6, 0.0, 2.0, False), + "top_p": ("UNSLOTH_SAMPLING_TOP_P", 0.95, 0.0, 1.0, False), + "top_k": ("UNSLOTH_SAMPLING_TOP_K", 20, -1, 100, True), + "min_p": ("UNSLOTH_SAMPLING_MIN_P", 0.01, 0.0, 1.0, False), + "repetition_penalty": ("UNSLOTH_SAMPLING_REPETITION_PENALTY", 1.0, 1.0, 2.0, False), + "presence_penalty": ("UNSLOTH_SAMPLING_PRESENCE_PENALTY", 0.0, 0.0, 2.0, False), +} + +# Public, ordered tuple of the sampling fields callers resolve. +SAMPLING_FIELD_NAMES = tuple(_SAMPLING_FIELDS) + +# Fields the Studio Chat UI adopts as *per-model recommendations* from the backend +# `.inference` block. Its frontend `mergeBackendRecommendedInference` +# (presets/preset-policy.ts) seeds exactly these five and never reads repetition_penalty, +# so the server auto-recommends the same five for request parity. repetition_penalty stays a +# manual-only knob (client-sent or an UNSLOTH_SAMPLING_REPETITION_PENALTY operator pin), +# matching the UI where it is never auto-filled per model. +_UI_RECOMMENDED_FIELDS = ("temperature", "top_p", "top_k", "min_p", "presence_penalty") + + +def _clean_sampling_value(field: str, val: Any): + """Coerce ``val`` to the field's numeric type when it is a finite, in-range number, else None. + + Rejects bool, non-numeric, NaN/inf, and out-of-range values so neither a bad operator env + var nor a malformed model recommendation can reach llama-server. NaN matters because + ``nan < lo`` and ``nan > hi`` are both False, so a plain range check would let it through. + Coerce before the finiteness check: ``math.isfinite`` and ``float()`` raise ``OverflowError`` + on an int too big for a C double (an oversized UNSLOTH_SAMPLING_TOP_K would otherwise 500 the + request), while an in-range int is range-checked exactly and ``int()`` rejects a NaN/inf that + reached an int field. + """ + if isinstance(val, bool) or not isinstance(val, (int, float)): + return None + _env, _default, lo, hi, is_int = _SAMPLING_FIELDS[field] + try: + val = int(val) if is_int else float(val) + except (ValueError, OverflowError): + # int(nan)/int(inf) and float(oversized_int) raise; treat them as unusable. + return None + # After coercion an int is always finite; only a float can still be NaN/inf. + if isinstance(val, float) and not math.isfinite(val): + return None + if val < lo or val > hi: + return None + return val + + +def _operator_sampling_override(field: str): + """Operator-pinned value for a sampling field from UNSLOTH_SAMPLING_*, or None. + + An unparseable, non-finite, or out-of-range value is ignored so a bad env var can never + reach llama-server; the field then falls back to the client / recommended value. + """ + _env, _default, _lo, _hi, is_int = _SAMPLING_FIELDS[field] + raw = os.environ.get(_env) + if raw is None or raw.strip() == "": + return None + try: + val = int(raw) if is_int else float(raw) + except (TypeError, ValueError): + return None + return _clean_sampling_value(field, val) + + +@lru_cache(maxsize = 128) +def _recommended_sampling(model_id: str) -> Dict[str, Any]: + """Per-model recommended sampling, resolved through the SAME path the Studio Chat UI uses. + + The Chat UI seeds its sampling from the ``.inference`` block of the load/status responses, + which is exactly :func:`load_inference_config` (model-specific YAML -> family defaults + (inference_defaults.json) -> default.yaml). Sourcing recommendations here keeps the values + the server applies to a request identical to what the UI shows for the same model. Only the + fields the UI actually adopts (:data:`_UI_RECOMMENDED_FIELDS`) are recommended; each value + is validated (finite + in range) before use. Cached by model id. + """ + if not model_id: + return {} + try: + cfg = load_inference_config(model_id) or {} + except Exception as e: + logger.debug(f"Could not load recommended sampling for '{model_id}': {e}") + return {} + recommended: Dict[str, Any] = {} + for field in _UI_RECOMMENDED_FIELDS: + cleaned = _clean_sampling_value(field, cfg.get(field)) + if cleaned is not None: + recommended[field] = cleaned + return recommended + + +def resolve_effective_sampling( + model_id: Optional[str], + explicit: Dict[str, Any], + *, + fill_defaults: bool = True, +) -> Dict[str, Any]: + """Resolve the effective sampling params for a request. + + ``explicit`` maps each field in :data:`SAMPLING_FIELD_NAMES` to the client-sent + value, or ``None`` when the client omitted it. Precedence (highest first): an + operator ``UNSLOTH_SAMPLING_*`` pin, then the client's explicit value, then the + per-model recommendation, then the static schema default. + + When ``fill_defaults`` is False a field with no operator pin, client value, or + per-model recommendation is omitted from the result instead of set to the static + schema default, so a raw proxy body (``/v1/completions``) keeps llama-server's own + default for that field rather than being forced onto this schema's value. + """ + recommended = _recommended_sampling(model_id or "") + effective: Dict[str, Any] = {} + for field, (_env, default, _lo, _hi, _int) in _SAMPLING_FIELDS.items(): + override = _operator_sampling_override(field) + if override is not None: + effective[field] = override + elif explicit.get(field) is not None: + effective[field] = explicit[field] + elif field in recommended: + effective[field] = recommended[field] + elif fill_defaults: + effective[field] = default + return effective diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 434c4a0ad5..408ea4bd34 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -145,6 +145,7 @@ _CODEX_ENV_UNSET = ("OPENAI_API_KEY", "CODEX_API_KEY", "CODEX_ACCESS_TOKEN") # instead of one long unaligned list. _PANEL_MODEL = "Model" _PANEL_SERVER = "Server" +_PANEL_SAMPLING = "Sampling" _PANEL_SESSION = "Agent session" _MODEL_OPTION = typer.Option( @@ -212,6 +213,56 @@ _TOOL_CALL_NUDGING_OPTION = typer.Option( help = "Retry once with a nudge when a non-streaming passthrough tool call can't be healed. " "On by default; when the flag is omitted an inherited UNSLOTH_TOOL_CALL_NUDGE is kept.", ) +# Sampling overrides pin a value on the auto-started server (winning over the client and the +# per-model recommendation). Default unset -> the model's recommended sampling is used. +_TEMPERATURE_OPTION = typer.Option( + None, + "--temperature", + min = 0.0, + max = 2.0, + rich_help_panel = _PANEL_SAMPLING, + help = "Pin the sampling temperature. Default: unset (per-model recommendation).", +) +_TOP_P_OPTION = typer.Option( + None, + "--top-p", + min = 0.0, + max = 1.0, + rich_help_panel = _PANEL_SAMPLING, + help = "Pin top-p (nucleus) sampling. Default: unset (per-model recommendation).", +) +_TOP_K_OPTION = typer.Option( + None, + "--top-k", + min = -1, + max = 100, + rich_help_panel = _PANEL_SAMPLING, + help = "Pin top-k sampling. Default: unset (per-model recommendation).", +) +_MIN_P_OPTION = typer.Option( + None, + "--min-p", + min = 0.0, + max = 1.0, + rich_help_panel = _PANEL_SAMPLING, + help = "Pin min-p sampling threshold. Default: unset (per-model recommendation).", +) +_REPETITION_PENALTY_OPTION = typer.Option( + None, + "--repetition-penalty", + min = 1.0, + max = 2.0, + rich_help_panel = _PANEL_SAMPLING, + help = "Pin the repetition penalty. Default: unset (per-model recommendation).", +) +_PRESENCE_PENALTY_OPTION = typer.Option( + None, + "--presence-penalty", + min = 0.0, + max = 2.0, + rich_help_panel = _PANEL_SAMPLING, + help = "Pin the presence penalty. Default: unset (per-model recommendation).", +) # Agent-session knobs. _KEY_OPTION = typer.Option( @@ -415,6 +466,12 @@ class ServerOptions(NamedTuple): enable_tools: bool = False tool_call_healing: Optional[bool] = None tool_call_nudging: Optional[bool] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + top_k: Optional[int] = None + min_p: Optional[float] = None + repetition_penalty: Optional[float] = None + presence_penalty: Optional[float] = None def _split_repo_variant(model: str) -> tuple: @@ -958,6 +1015,18 @@ def _start_studio_server( child_env["UNSLOTH_TOOL_CALL_NUDGE"] = "1" if server.tool_call_nudging else "0" elif "UNSLOTH_TOOL_CALL_NUDGE" not in child_env: child_env["UNSLOTH_TOOL_CALL_NUDGE"] = "1" + # Forward any sampling pin via the env; `unsloth run` reads UNSLOTH_SAMPLING_* and the + # backend resolver applies it as a hard override. Only set fields the operator specified. + for _sampling_env, _sampling_value in ( + ("UNSLOTH_SAMPLING_TEMPERATURE", server.temperature), + ("UNSLOTH_SAMPLING_TOP_P", server.top_p), + ("UNSLOTH_SAMPLING_TOP_K", server.top_k), + ("UNSLOTH_SAMPLING_MIN_P", server.min_p), + ("UNSLOTH_SAMPLING_REPETITION_PENALTY", server.repetition_penalty), + ("UNSLOTH_SAMPLING_PRESENCE_PENALTY", server.presence_penalty), + ): + if _sampling_value is not None: + child_env[_sampling_env] = str(_sampling_value) kwargs: dict = { "stdout": log, "stderr": subprocess.STDOUT, @@ -1045,6 +1114,30 @@ def _require_studio( """Return (base, server). server is a Popen only when WE auto-started it.""" base = find_studio_server() if base is not None: + # Attaching to a server someone else started: UNSLOTH_SAMPLING_* pins only reach the + # server process when WE launch it (via _start_studio_server), so a sampling flag on the + # attach path can't take effect. Warn instead of silently dropping it, so the operator is + # not misled into thinking generation now uses the pinned value. + _pinned = [ + _flag + for _flag, _value in ( + ("--temperature", server_options.temperature), + ("--top-p", server_options.top_p), + ("--top-k", server_options.top_k), + ("--min-p", server_options.min_p), + ("--repetition-penalty", server_options.repetition_penalty), + ("--presence-penalty", server_options.presence_penalty), + ) + if _value is not None + ] + if _pinned: + typer.echo( + f"Warning: an Unsloth server is already running at {base}; sampling pins " + f"({', '.join(_pinned)}) apply only when this command starts the server, so the " + "running server keeps its current sampling. Stop it with `unsloth studio stop` " + "and re-run to apply them.", + err = True, + ) 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 @@ -2807,6 +2900,12 @@ def claude( enable_tools: bool = _ENABLE_TOOLS_OPTION, tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION, tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION, + temperature: Optional[float] = _TEMPERATURE_OPTION, + top_p: Optional[float] = _TOP_P_OPTION, + top_k: Optional[int] = _TOP_K_OPTION, + min_p: Optional[float] = _MIN_P_OPTION, + repetition_penalty: Optional[float] = _REPETITION_PENALTY_OPTION, + presence_penalty: Optional[float] = _PRESENCE_PENALTY_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, persist: bool = _PERSIST_OPTION, @@ -2821,7 +2920,17 @@ def claude( LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), serve = serve, launch = launch, - server_options = ServerOptions(enable_tools, tool_call_healing, tool_call_nudging), + server_options = ServerOptions( + enable_tools = enable_tools, + tool_call_healing = tool_call_healing, + tool_call_nudging = tool_call_nudging, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ), ) model_id = entry["id"] install_hint = ( @@ -2908,6 +3017,12 @@ def codex( enable_tools: bool = _ENABLE_TOOLS_OPTION, tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION, tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION, + temperature: Optional[float] = _TEMPERATURE_OPTION, + top_p: Optional[float] = _TOP_P_OPTION, + top_k: Optional[int] = _TOP_K_OPTION, + min_p: Optional[float] = _MIN_P_OPTION, + repetition_penalty: Optional[float] = _REPETITION_PENALTY_OPTION, + presence_penalty: Optional[float] = _PRESENCE_PENALTY_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, persist: bool = _PERSIST_OPTION, @@ -2922,7 +3037,17 @@ def codex( LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), serve = serve, launch = launch, - server_options = ServerOptions(enable_tools, tool_call_healing, tool_call_nudging), + server_options = ServerOptions( + enable_tools = enable_tools, + tool_call_healing = tool_call_healing, + tool_call_nudging = tool_call_nudging, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ), ) # This preflight runs after _connect may have auto-started a server but before _run # takes over its lifecycle, so tear the server down here if it rejects the model @@ -2990,6 +3115,12 @@ def openclaw( enable_tools: bool = _ENABLE_TOOLS_OPTION, tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION, tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION, + temperature: Optional[float] = _TEMPERATURE_OPTION, + top_p: Optional[float] = _TOP_P_OPTION, + top_k: Optional[int] = _TOP_K_OPTION, + min_p: Optional[float] = _MIN_P_OPTION, + repetition_penalty: Optional[float] = _REPETITION_PENALTY_OPTION, + presence_penalty: Optional[float] = _PRESENCE_PENALTY_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, persist: bool = _PERSIST_OPTION, @@ -3004,7 +3135,17 @@ def openclaw( LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), serve = serve, launch = launch, - server_options = ServerOptions(enable_tools, tool_call_healing, tool_call_nudging), + server_options = ServerOptions( + enable_tools = enable_tools, + tool_call_healing = tool_call_healing, + tool_call_nudging = tool_call_nudging, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ), ) openclaw_args = list(ctx.args) # Default a bare `unsloth start openclaw` to the local TUI. Anything the caller @@ -3054,6 +3195,12 @@ def opencode( enable_tools: bool = _ENABLE_TOOLS_OPTION, tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION, tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION, + temperature: Optional[float] = _TEMPERATURE_OPTION, + top_p: Optional[float] = _TOP_P_OPTION, + top_k: Optional[int] = _TOP_K_OPTION, + min_p: Optional[float] = _MIN_P_OPTION, + repetition_penalty: Optional[float] = _REPETITION_PENALTY_OPTION, + presence_penalty: Optional[float] = _PRESENCE_PENALTY_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, persist: bool = _PERSIST_OPTION, @@ -3068,7 +3215,17 @@ def opencode( LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), serve = serve, launch = launch, - server_options = ServerOptions(enable_tools, tool_call_healing, tool_call_nudging), + server_options = ServerOptions( + enable_tools = enable_tools, + tool_call_healing = tool_call_healing, + tool_call_nudging = tool_call_nudging, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ), ) if as_subagent: subagent_id = _subagent_model_id(base, key, entry, model, gguf_variant) @@ -3198,6 +3355,12 @@ def hermes( enable_tools: bool = _ENABLE_TOOLS_OPTION, tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION, tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION, + temperature: Optional[float] = _TEMPERATURE_OPTION, + top_p: Optional[float] = _TOP_P_OPTION, + top_k: Optional[int] = _TOP_K_OPTION, + min_p: Optional[float] = _MIN_P_OPTION, + repetition_penalty: Optional[float] = _REPETITION_PENALTY_OPTION, + presence_penalty: Optional[float] = _PRESENCE_PENALTY_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, persist: bool = _PERSIST_OPTION, @@ -3214,7 +3377,17 @@ def hermes( LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), serve = serve, launch = launch, - server_options = ServerOptions(enable_tools, tool_call_healing, tool_call_nudging), + server_options = ServerOptions( + enable_tools = enable_tools, + tool_call_healing = tool_call_healing, + tool_call_nudging = tool_call_nudging, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ), ) install_hint = _hermes_install_hint() with _session_config("hermes", launch, persist = persist) as home: @@ -3238,6 +3411,12 @@ def pi( enable_tools: bool = _ENABLE_TOOLS_OPTION, tool_call_healing: Optional[bool] = _TOOL_CALL_HEALING_OPTION, tool_call_nudging: Optional[bool] = _TOOL_CALL_NUDGING_OPTION, + temperature: Optional[float] = _TEMPERATURE_OPTION, + top_p: Optional[float] = _TOP_P_OPTION, + top_k: Optional[int] = _TOP_K_OPTION, + min_p: Optional[float] = _MIN_P_OPTION, + repetition_penalty: Optional[float] = _REPETITION_PENALTY_OPTION, + presence_penalty: Optional[float] = _PRESENCE_PENALTY_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, persist: bool = _PERSIST_OPTION, @@ -3252,7 +3431,17 @@ def pi( LoadOptions(gguf_variant, max_seq_length, load_in_4bit, tensor_parallel), serve = serve, launch = launch, - server_options = ServerOptions(enable_tools, tool_call_healing, tool_call_nudging), + server_options = ServerOptions( + enable_tools = enable_tools, + tool_call_healing = tool_call_healing, + tool_call_nudging = tool_call_nudging, + temperature = temperature, + top_p = top_p, + top_k = top_k, + min_p = min_p, + repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, + ), ) install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent" if as_subagent: diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 9a472d2996..84c22aa4ac 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -1661,6 +1661,7 @@ def _consume_legacy_short_aliases( _RUN_PANEL_MODEL = "Model" _RUN_PANEL_SERVER = "Server & network" _RUN_PANEL_TOOLS = "Tool calls" +_RUN_PANEL_SAMPLING = "Sampling" _RUN_PANEL_ADVANCED = "Advanced" @@ -1758,6 +1759,57 @@ def run( "Default: on. No effect on streaming requests or the server-side agentic loop." ), ), + temperature: Optional[float] = typer.Option( + None, + "--temperature", + min = 0.0, + max = 2.0, + rich_help_panel = _RUN_PANEL_SAMPLING, + help = ( + "Pin the sampling temperature for every request that omits it, overriding the " + "model's recommended value. Default: unset (use the per-model recommendation)." + ), + ), + top_p: Optional[float] = typer.Option( + None, + "--top-p", + min = 0.0, + max = 1.0, + rich_help_panel = _RUN_PANEL_SAMPLING, + help = "Pin top-p (nucleus) sampling. Default: unset (per-model recommendation).", + ), + top_k: Optional[int] = typer.Option( + None, + "--top-k", + min = -1, + max = 100, + rich_help_panel = _RUN_PANEL_SAMPLING, + help = "Pin top-k sampling. Default: unset (per-model recommendation).", + ), + min_p: Optional[float] = typer.Option( + None, + "--min-p", + min = 0.0, + max = 1.0, + rich_help_panel = _RUN_PANEL_SAMPLING, + help = "Pin min-p sampling threshold. Default: unset (per-model recommendation).", + ), + repetition_penalty: Optional[float] = typer.Option( + None, + "--repetition-penalty", + min = 1.0, + max = 2.0, + rich_help_panel = _RUN_PANEL_SAMPLING, + help = "Pin the repetition penalty. Default: unset (per-model recommendation).", + ), + presence_penalty: Optional[float] = typer.Option( + None, + "--presence-penalty", + min = 0.0, + max = 2.0, + rich_help_panel = _RUN_PANEL_SAMPLING, + help = "Pin the presence penalty. Default: unset (per-model recommendation).", + ), yes: bool = typer.Option( False, "--yes", @@ -1841,7 +1893,7 @@ def run( Example: unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL - unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --top-k 20 --seed 42 --parallel 8 + unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --temperature 0.7 --seed 42 --parallel 8 unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja unsloth studio run --model unsloth/Qwen3-27B-GGUF --gguf-variant Q8_0 --tensor-parallel """ @@ -1870,6 +1922,21 @@ def run( elif "UNSLOTH_TOOL_CALL_NUDGE" not in os.environ: os.environ["UNSLOTH_TOOL_CALL_NUDGE"] = "1" + # Sampling overrides: the backend resolver reads UNSLOTH_SAMPLING_* to hard-pin a field + # (winning over both the client and the per-model recommendation). Only write a flag that + # was set explicitly so an omitted flag inherits any value the parent forwarded (e.g. + # `unsloth start`) and, when nothing is set, leaves the per-model recommendation in charge. + for _sampling_env, _sampling_value in ( + ("UNSLOTH_SAMPLING_TEMPERATURE", temperature), + ("UNSLOTH_SAMPLING_TOP_P", top_p), + ("UNSLOTH_SAMPLING_TOP_K", top_k), + ("UNSLOTH_SAMPLING_MIN_P", min_p), + ("UNSLOTH_SAMPLING_REPETITION_PENALTY", repetition_penalty), + ("UNSLOTH_SAMPLING_PRESENCE_PENALTY", presence_penalty), + ): + if _sampling_value is not None: + os.environ[_sampling_env] = str(_sampling_value) + # Set before any re-exec so the in-venv server inherits it via the env. # `run --verbose` used to pass through to llama-server (its own -v); keep # that by forwarding --log-verbose so we add Unsloth logs without dropping it. diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 42745ef9c5..ade82cc06b 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -1984,6 +1984,118 @@ def test_start_studio_server_respects_inherited_tool_call_env(monkeypatch): assert env["UNSLOTH_TOOL_CALL_NUDGE"] == "1" +def test_start_studio_server_forwards_sampling_via_env(monkeypatch): + # Sampling pins ride to the child server through UNSLOTH_SAMPLING_*; unset ones stay absent + # so the backend keeps the per-model recommendation. + captured = {} + + class FakePopen: + def __init__(self, command, **kwargs): + captured["kwargs"] = kwargs + self.pid = 1 + + 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-x") + monkeypatch.setattr(start.time, "sleep", lambda _s: None) + for _v in ("TEMPERATURE", "TOP_P", "TOP_K", "MIN_P", "REPETITION_PENALTY", "PRESENCE_PENALTY"): + monkeypatch.delenv(f"UNSLOTH_SAMPLING_{_v}", raising = False) + + # No sampling flags -> nothing forwarded. + start._start_studio_server("http://127.0.0.1:8888", "unsloth/M-GGUF", start.LoadOptions()) + env = captured["kwargs"]["env"] + assert not any(k.startswith("UNSLOTH_SAMPLING_") for k in env) + + # Pins are forwarded; unset ones stay absent. + start._start_studio_server( + "http://127.0.0.1:8888", + "unsloth/M-GGUF", + start.LoadOptions(), + start.ServerOptions(temperature = 0.3, top_k = 40, min_p = 0.05), + ) + env = captured["kwargs"]["env"] + assert env["UNSLOTH_SAMPLING_TEMPERATURE"] == "0.3" + assert env["UNSLOTH_SAMPLING_TOP_K"] == "40" + assert env["UNSLOTH_SAMPLING_MIN_P"] == "0.05" + assert "UNSLOTH_SAMPLING_TOP_P" not in env + + +def test_require_studio_warns_on_sampling_pin_when_reusing_server(monkeypatch, capsys): + # Attaching to an already-running server can't apply UNSLOTH_SAMPLING_* pins (only + # _start_studio_server forwards them), so a sampling flag on the attach path must warn + # instead of being silently dropped while the command "succeeds". + monkeypatch.setattr(start, "find_studio_server", lambda: BASE) + base, server = start._require_studio( + "unsloth/M-GGUF", + start.LoadOptions(), + serve = True, + launch = True, + server_options = start.ServerOptions(temperature = 0.3, top_k = 40), + ) + assert base == BASE + assert server is None # attach path: we did not start the server + err = capsys.readouterr().err + assert "already running" in err + assert "--temperature" in err and "--top-k" in err + # Only the pinned fields are named; an unset one is not. + assert "--top-p" not in err + + +def test_require_studio_no_sampling_warning_without_pins(monkeypatch, capsys): + # Reusing a server with no sampling pins stays silent (tool flags are out of scope here). + monkeypatch.setattr(start, "find_studio_server", lambda: BASE) + base, server = start._require_studio( + "unsloth/M-GGUF", + start.LoadOptions(), + serve = True, + server_options = start.ServerOptions(enable_tools = True), + ) + assert base == BASE and server is None + assert "sampling" not in capsys.readouterr().err.lower() + + +def test_start_claude_parses_sampling_flags(fake_studio, monkeypatch): + # `unsloth start claude ... --temperature 0.3 --top-k 40` routes the pins into ServerOptions. + monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://127.0.0.1:8888") + monkeypatch.setattr(start, "find_studio_server", lambda: None) + captured = {} + fake = SimpleNamespace(pid = 1, poll = lambda: None) + + def fake_start( + base, + model, + load, + server_options = None, + ): + captured["server_options"] = server_options + 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/gemma-4-E2B-it-GGUF", + "--temperature", + "0.3", + "--top-k", + "40", + ], + ) + assert result.exit_code == 0, result.output + so = captured["server_options"] + assert so.temperature == 0.3 and so.top_k == 40 and so.top_p is None + + 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. diff --git a/unsloth_cli/tests/test_studio_run_parallel_flag.py b/unsloth_cli/tests/test_studio_run_parallel_flag.py index b2fc421359..972abb5d4d 100644 --- a/unsloth_cli/tests/test_studio_run_parallel_flag.py +++ b/unsloth_cli/tests/test_studio_run_parallel_flag.py @@ -305,6 +305,43 @@ def test_run_omitted_flag_respects_inherited_env(monkeypatch, inherited): assert studio_mod.os.environ["UNSLOTH_TOOL_CALL_NUDGE"] == inherited +_SAMPLING_ENV_SUFFIXES = ( + "TEMPERATURE", + "TOP_P", + "TOP_K", + "MIN_P", + "REPETITION_PENALTY", + "PRESENCE_PENALTY", +) + + +def test_run_sampling_flags_set_env(monkeypatch): + """`--temperature`/`--top-k` write UNSLOTH_SAMPLING_* (a hard override the backend applies); + an omitted sampling flag leaves its env unset so the per-model recommendation stays.""" + studio_mod = _load_run_command() + for _v in _SAMPLING_ENV_SUFFIXES: + monkeypatch.delenv(f"UNSLOTH_SAMPLING_{_v}", raising = False) + _invoke_run(monkeypatch, _BASE + ["--temperature", "0.3", "--top-k", "40"]) + assert studio_mod.os.environ["UNSLOTH_SAMPLING_TEMPERATURE"] == "0.3" + assert studio_mod.os.environ["UNSLOTH_SAMPLING_TOP_K"] == "40" + assert "UNSLOTH_SAMPLING_TOP_P" not in studio_mod.os.environ + + +def test_run_no_sampling_flags_leaves_env_unset(monkeypatch): + """Plain `unsloth run` writes no UNSLOTH_SAMPLING_*; the server keeps the recommendation.""" + studio_mod = _load_run_command() + for _v in _SAMPLING_ENV_SUFFIXES: + monkeypatch.delenv(f"UNSLOTH_SAMPLING_{_v}", raising = False) + _invoke_run(monkeypatch, _BASE) + assert not any(k.startswith("UNSLOTH_SAMPLING_") for k in studio_mod.os.environ) + + +def test_run_rejects_out_of_range_sampling_flag(monkeypatch): + """typer enforces the documented ranges before a value can reach the server.""" + result, _captured = _invoke_run(monkeypatch, _BASE + ["--temperature", "9"]) + assert result.exit_code != 0 + + @pytest.mark.parametrize("platform", ["linux", "darwin", "win32"]) def test_reexec_argv_is_consistent_across_platforms(monkeypatch, platform): """Linux/Darwin (execvp) and Windows (Popen) must build the same argv.""" @@ -344,12 +381,14 @@ def test_reexec_mixed_parallel_with_passthrough(monkeypatch): """--parallel + llama-server pass-through flags must all reach the child.""" result, captured = _invoke_run( monkeypatch, - _BASE + ["--parallel", "8", "--top-k", "20", "--temp", "0.7"], + # --top-k is now a first-class sampling flag (routed via UNSLOTH_SAMPLING_*), so use + # --seed / --temp here, which remain genuine llama-server pass-through flags. + _BASE + ["--parallel", "8", "--seed", "42", "--temp", "0.7"], ) assert len(captured) == 1 argv = captured[0]["argv"] assert _value_after(argv, "--parallel") == "8", argv - assert _value_after(argv, "--top-k") == "20", argv + assert _value_after(argv, "--seed") == "42", argv assert _value_after(argv, "--temp") == "0.7", argv From a7761e1740819d522ca098a2205b5e73c93e3929 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 21:02:29 -0700 Subject: [PATCH 016/161] Studio: refine GGUF per-GPU selection (gpu_ids) (#7239) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 246 +++++++++++++++--- studio/backend/models/inference.py | 20 +- studio/backend/routes/inference.py | 209 ++++++++------- studio/backend/routes/training_vram.py | 39 ++- .../tests/test_chat_load_during_training.py | 89 +++---- studio/backend/tests/test_gpu_memory_mode.py | 106 +++++++- studio/backend/tests/test_gpu_selection.py | 203 ++++++++++++++- studio/backend/utils/hardware/hardware.py | 49 ++-- .../lib/apply-inference-status-to-store.ts | 4 +- .../chat/stores/chat-runtime-store.ts | 7 +- .../frontend/src/features/chat/types/api.ts | 6 + tests/studio/test_model_picker_contracts.py | 13 + 12 files changed, 773 insertions(+), 218 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index fc58442a79..7d67339c1e 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2023,6 +2023,10 @@ class LlamaCppBackend: self._tensor_split: Optional[List[float]] = None # User-picked physical GPU indices (None = automatic selection). self._gpu_ids: Optional[List[int]] = None + # RAW requested GPU pin, before the fit narrowed it. self._gpu_ids records the + # EFFECTIVE (fit-narrowed) pin for /status; dedupe compares this raw value so a + # [0, 1] narrowed to [0] and re-sent as [0, 1] still matches (#7239). + self._requested_gpu_ids: Optional[List[int]] = None # Layer load kept multi-GPU only to honor a downgraded tensor request, so a # later explicit tensor-off reloads instead of deduping to it (#6659). self._layer_preserves_tensor_intent: bool = False @@ -2494,6 +2498,46 @@ class LlamaCppBackend: """User-picked physical GPU indices, or None for automatic selection.""" return self._gpu_ids + @property + def requested_gpu_ids(self) -> Optional[List[int]]: + """RAW requested GPU pin (before the fit narrowed it), or None for auto. + gpu_ids echoes the EFFECTIVE pin for /status.""" + return self._requested_gpu_ids + + def matches_gpu_ids(self, gpu_ids: Optional[List[int]]) -> bool: + """Whether a requested pin is already satisfied by the active runner. + + A regular GGUF load may narrow the requested placement pool to the + smallest fitting subset. Accept both the original request and the + effective status-echoed subset so either can round-trip without a + needless reload. Diffusion drives one device and keeps its existing + lowest-device normalization. + """ + if self._is_diffusion: + requested = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + return requested == (self._gpu_ids or None) + + requested = sorted(int(x) for x in gpu_ids) if gpu_ids else None + raw = self._requested_gpu_ids or None + effective = self._gpu_ids or None + return requested == raw or requested == effective + + def _record_matching_gpu_request(self, gpu_ids: Optional[List[int]]) -> None: + """Adopt the caller's explicit pool after a full already-loaded match. + + Matching an effective subset avoids a reload, but the incoming request + is still the user's latest placement intent. Record it so status and a + later reload do not restore GPUs the user just removed. + """ + if self._is_diffusion: + self._requested_gpu_ids = [sorted(int(x) for x in gpu_ids)[0]] if gpu_ids else None + else: + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + if self._last_load_kwargs is not None: + self._last_load_kwargs["gpu_ids"] = ( + list(self._requested_gpu_ids) if self._requested_gpu_ids else None + ) + @property def n_layers(self) -> Optional[int]: """Model layer count (GGUF block_count), or None if unknown.""" @@ -4581,6 +4625,14 @@ class LlamaCppBackend: LlamaCppBackend._gguf_skip_value(f, atype) return None + @classmethod + def _gguf_path_is_diffusion(cls, gguf_path: str, model_identifier: str) -> bool: + """Classify a downloaded GGUF without mutating the active backend.""" + probe = object.__new__(cls) + probe._model_identifier = model_identifier + probe._read_gguf_metadata(gguf_path) + return probe._is_diffusion + def _read_gguf_metadata(self, gguf_path: str) -> None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -5032,11 +5084,14 @@ class LlamaCppBackend: # the unload reset) so /status doesn't misreport TP and an identical # re-Apply doesn't reload against stale tensor-parallel state. self._tensor_parallel = False - # Record only the single device the runner actually uses (the lowest - # selected GPU, chosen above) -- not the whole pick. The diffusion runner - # is single-device, so echoing a multi-GPU list would misreport placement - # in /status and let a re-Apply dedup against GPUs the runner never used. + # The single-device runner records only the lowest selected GPU (chosen + # above), not the whole pick, and clears any explicit pin from a prior + # chat load; a multi-GPU list would misreport placement and mis-dedup. self._gpu_ids = [sorted(gpu_ids)[0]] if gpu_ids else None + # The frontend prefers requested_gpu_ids when hydrating the picker. + # Diffusion uses only one device, so echo the collapsed effective pin, + # not unused members of the original request. + self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None if hf_variant: self._hf_variant = hf_variant elif gguf_path: @@ -6161,6 +6216,8 @@ class LlamaCppBackend: gpu_layers: int = -1, n_cpu_moe: int = 0, tensor_split: Optional[List[float]] = None, + # Explicit GPU placement pool (issue #7164). None/[] = auto-select; + # the fitter may pin the smallest subset of this pool that fits. gpu_ids: Optional[List[int]] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # caller compat, unused @@ -6258,15 +6315,63 @@ class LlamaCppBackend: self._cancel_event.clear() - # ── Phase 1: kill old process (under lock, fast) ────────── - with self._lock: - self._kill_process() - # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() is_vulkan_backend = self._is_vulkan_backend(binary) + # ── Vulkan-ordinal preflight (BEFORE the Phase 1 kill) ──────── + # An explicit Vulkan pin the ggml probe never enumerated cannot be honored. + # Validate it ABOVE the kill so an invalid selection leaves the live model + # untouched: CUDA ids are range-checked at the route, but Vulkan ordinals are + # not, so a stale gpu_ids=[99] used to kill the server then 400, leaving + # nothing running (#7239). _get_gpu_memory needs only the binary (safe pre- + # download) and reuses the later fit's issubset logic. Guarded on a found + # Vulkan build + a pin so a deferred not-found stays deferred for diffusion. + if is_vulkan_backend and gpu_ids and binary: + _pf_wanted = {int(x) for x in gpu_ids} + _pf_probed = {g[0] for g in self._get_gpu_memory(binary)} + if not _pf_wanted.issubset(_pf_probed): + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {sorted(_pf_wanted)} not " + f"present. Available Vulkan devices: {sorted(_pf_probed)}." + ) + + # A remote uncached GGUF may only reveal that it needs the + # single-device diffusion runner after download. On Vulkan, an + # explicit gpu_ids request cannot be mapped from ggml ordinals to + # that runner's CUDA physical index. Download and classify the main + # file before killing the healthy server so this late rejection is + # non-destructive. The Phase 2 call below reuses this cached path. + _preflight_model_path = None + if is_vulkan_backend and gpu_ids and hf_repo: + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo + with _hf_offline_if_dns_dead(): + _preflight_model_path = self._download_gguf( + hf_repo = hf_repo, + hf_variant = hf_variant, + hf_token = hf_token, + ) + if self._gguf_path_is_diffusion(_preflight_model_path, model_identifier): + raise ValueError( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ) + + # ── Phase 1: kill old process (under lock, fast) ────────── + with self._lock: + self._kill_process() + # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected # sibling); for -hf loads it's None here and resolved just below. @@ -6288,7 +6393,7 @@ class LlamaCppBackend: ) hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): - model_path = self._download_gguf( + model_path = _preflight_model_path or self._download_gguf( hf_repo = hf_repo, hf_variant = hf_variant, hf_token = hf_token, @@ -6338,6 +6443,18 @@ class LlamaCppBackend: # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; # serve them with the diffusion runner (same OpenAI-compat interface). if self._is_diffusion: + # The diffusion runner pins its child by CUDA visibility mask, so a + # ggml Vulkan ordinal cannot be honored (wrong GPU / CPU fallback). + # Route and remote-download preflights reject before teardown; keep + # this as a final defense if classification ever disagrees. + if is_vulkan_backend and gpu_ids: + raise ValueError( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ) # Not a tensor/layer GGUF: clear any preserved-fallback flag from a # prior load (this path skips the command builder that clears it). self._layer_preserves_tensor_intent = False @@ -6558,6 +6675,12 @@ class LlamaCppBackend: # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). _layer_min_gpus = 1 + # An explicit Vulkan ordinal absent from the ggml probe cannot be + # honored; flag it in the fit and reject after the try (raising inside + # would be swallowed into the --fit-on fallback). Bound before the try. + _vulkan_explicit_unmatched = False + _vulkan_requested_ids: list[int] = [] + _vulkan_available_ordinals: list[int] = [] try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -6570,6 +6693,28 @@ class LlamaCppBackend: # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] + # Restrict the fit (and thus the layer plan + pin env) to the + # selected GPUs; fail-open if none match so a stale UI choice + # can't strand the load on CPU (issue #7164). + if gpu_ids: + # A Vulkan build indexes by ggml ordinal. An explicit ordinal + # absent from the probe can't be pinned, so reject after the try + # rather than fail-open onto a device the user didn't pick. + _wanted_ids = {int(x) for x in gpu_ids} + # Reject if ANY requested ordinal is absent, not only when none + # match: [0, 99] against {0, 1} silently drops 99. Comparing the + # full requested set (before filter narrows) still lets the fitter + # pick a valid subset later -- that is narrowing, not absence. + _probed_ordinals = {g[0] for g in gpus} + if is_vulkan_backend and not _wanted_ids.issubset(_probed_ordinals): + _vulkan_explicit_unmatched = True + _vulkan_requested_ids = sorted(_wanted_ids) + _vulkan_available_ordinals = sorted(_probed_ordinals) + # Restrict the probed pool to the selection; fail-open (keep the + # full pool) if none match so a stale UI choice can't strand the + # load on CPU (issue #7164). + _sel_gpus = [g for g in gpus if g[0] in _wanted_ids] + gpus = _sel_gpus if _sel_gpus else gpus total_by_idx = {idx: total for idx, _f, total in _gpu_mem} # GPU picker: restrict every mode to the chosen devices, so # auto selection only considers them and manual mask to @@ -7396,6 +7541,17 @@ class LlamaCppBackend: tp_tensor_split = None effective_ctx = requested_ctx # fall back to original + # An unenumerated explicit Vulkan ordinal can't be pinned; fail loudly + # instead of fitting onto an unselected device. Clear the raw selection + # the early state-publish recorded so it never leaks into gpu_ids (#7239). + if _vulkan_explicit_unmatched: + self._gpu_ids = None + self._requested_gpu_ids = None + raise ValueError( + f"Requested Vulkan GPU ordinal(s) {_vulkan_requested_ids} not " + f"present. Available Vulkan devices: {_vulkan_available_ordinals}." + ) + # GPU picker: when no narrower subset was chosen (manual, or # a failed/file-size selection), pin the whole picked set so the # model can't spill onto an unpicked GPU. @@ -7759,11 +7915,45 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) - # Vulkan pins via --device (a cmd arg, unlike the env-based - # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's - # last-wins parsing lets a user --device override Unsloth's pick. - if is_vulkan_backend and gpu_indices is not None: - cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # Vulkan pins via --device (a cmd arg), before user extras so a user + # --device wins. Fall back to raw ids when the fit did not narrow. + _vulkan_pin_ids = gpu_indices if gpu_indices is not None else (gpu_ids or None) + + # Record the pin actually applied (fit-narrowed gpu_indices, else the raw + # request) for the keep-warm loop, dedupe, and /status, so an explicit + # [0, 1] narrowed to [0] records [0] and /status never echoes an ordinal + # the child never saw. Auto selection (no gpu_ids) stays None (#7239). + if is_vulkan_backend: + # Only record an EXPLICIT Vulkan pin: an auto pick still narrows + + # pins below, but recording it would misreport an explicit pin and + # make dedupe miss the loaded server; mirrors the CUDA/ROCm branch. + self._gpu_ids = ( + sorted(int(x) for x in _vulkan_pin_ids) + if (gpu_ids and _vulkan_pin_ids) + else None + ) + elif gpu_ids: + # Physical pin: the fit-selected subset when the fit ran, else the raw + # user selection so an explicit choice is honoured even when the fit + # could not size the model. + _effective_pin_ids = ( + [int(x) for x in gpu_indices] + if gpu_indices is not None + else [int(x) for x in gpu_ids] + ) + self._gpu_ids = ( + sorted(int(x) for x in _effective_pin_ids) if _effective_pin_ids else None + ) + else: + self._gpu_ids = None + + # Also record the RAW requested pin (before the fit narrowed it). Load + # dedupe compares this so a [0, 1] narrowed to [0] and re-sent as [0, 1] + # still matches, while /status keeps echoing the effective pin (#7239). + self._requested_gpu_ids = sorted(int(x) for x in gpu_ids) if gpu_ids else None + + if is_vulkan_backend and _vulkan_pin_ids is not None: + cmd += LlamaCppBackend._vulkan_pin_args(_vulkan_pin_ids) # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Unsloth's auto-set flags. Already @@ -7832,10 +8022,10 @@ class LlamaCppBackend: f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) - # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so - # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device - # (above), not here. + # Pin to selected GPU(s) (issue #7164; resolved above into gpu_indices). + # On ROCm, narrowing only CUDA_VISIBLE_DEVICES leaves the AMD child + # seeing the full set, so set HIP_VISIBLE_DEVICES too. Vulkan is pinned + # via --device (above), not here. # A deliberate zero-offload load with no GPU companions runs # entirely on CPU, yet a visible CUDA device still costs the child # ~0.5 GB (context + compute scratch) that the CPU-only @@ -8756,16 +8946,10 @@ class LlamaCppBackend: ) ): return False - # A changed GPU pick must reload (compare order-insensitively; None/[] - # both mean automatic). The diffusion runner collapses a multi-GPU pick - # to its single lowest device, so self._gpu_ids holds just that device; - # normalize the request the same way, or a multi-GPU pick that resolves - # to the same device needlessly reloads. - if self._is_diffusion: - requested_gpu_pick = [sorted(gpu_ids)[0]] if gpu_ids else None - else: - requested_gpu_pick = sorted(gpu_ids) if gpu_ids else None - if (self._gpu_ids or None) != requested_gpu_pick: + # A changed GPU pick must reload. Regular GGUF accepts either the raw + # requested placement pool or the effective status-echoed subset; + # diffusion compares its normalized single-device pick. + if not self.matches_gpu_ids(gpu_ids): return False # Compare on the canonical requested mode. With --spec-type in @@ -8823,6 +9007,7 @@ class LlamaCppBackend: current = list(self._extra_args) if self._extra_args is not None else [] if list(extra_args) != current: return False + self._record_matching_gpu_request(gpu_ids) return True def _classify_gpu_offload( @@ -8954,12 +9139,15 @@ class LlamaCppBackend: self._supports_preserve_thinking = False self._supports_tools = False self._cache_type_kv = None + # GPU-pin state describes the active runner only; clear it so an explicit + # pin never leaks into the next (or diffusion) runner. + self._gpu_ids = None + self._requested_gpu_ids = None self._tensor_parallel = False self._gpu_memory_mode = "auto" self._gpu_layers = -1 self._n_cpu_moe = 0 self._tensor_split = None - self._gpu_ids = None self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index c7e5ffa36b..099d356a73 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -74,7 +74,7 @@ class LoadRequest(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.", + description = "GPU placement pool, for example [0, 1]. Omit or pass [] to use automatic selection. CUDA/ROCm values are physical GPU indices and are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries; Vulkan values are ggml device ordinals. For GGUF models the fitter may pin the smallest subset of this pool that fits.", ) speculative_type: Optional[str] = Field( None, @@ -485,7 +485,14 @@ class LoadResponse(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.", + ) + requested_gpu_ids: Optional[List[int]] = Field( + None, + description = ( + "GPU placement pool requested by the user before fit-time narrowing, " + "or None for automatic selection." + ), ) @@ -649,7 +656,14 @@ class InferenceStatusResponse(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices the model is pinned to, or None for automatic selection.", + description = "Effective GPU indices the model is using after fit-time narrowing, or None for automatic selection.", + ) + requested_gpu_ids: Optional[List[int]] = Field( + None, + description = ( + "GPU placement pool requested by the user before fit-time narrowing, " + "or None for automatic selection." + ), ) llama_cpp_supports_mtp: bool = Field( True, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 389a0d245a..6dbaa8fcc9 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -3241,15 +3241,10 @@ def _request_matches_loaded_settings( ) ): return False - # A changed GPU pick must reload. The diffusion runner collapses a multi-GPU - # request to its single lowest device (it drives one device only), so the - # backend records just that device; compare the request the same way, or a - # multi-GPU pick that resolves to the same device needlessly reloads. - if llama_backend.is_diffusion: - _req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None - else: - _req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None - if _req_gpu_ids != llama_backend.gpu_ids: + # A regular GGUF may narrow the requested placement pool. Accept either the + # original request or the effective status-echoed subset; diffusion keeps + # its single-device normalization. + if not llama_backend.matches_gpu_ids(request.gpu_ids): return False # Preserved tensor->layer fallback (both report tensor=off, so the check above # matches): if the user now explicitly drops tensor intent, reload so placement @@ -3897,15 +3892,19 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: """Classify a GGUF as diffusion, normal, or unknown before it is loaded. ``None`` is important here: a remote GGUF whose header is not cached can - still be routed to the single-GPU diffusion runner after download. Treating - that case as normal would let Manual mode skip the training guard even - though the runner ignores Manual's llama-server placement controls. + still be routed to the single-GPU diffusion runner after download. Default + placement keeps that unknown case guarded until the header is available. """ identity = " ".join( str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file") ).lower() - if "diffusion" in identity: - return True + # Name-only hint, used ONLY as a pre-download fallback, scoped to the + # DiffusionGemma runner family: a bare "diffusion" substring is common in + # ordinary text-model names/paths (e.g. "stable-diffusion-prompt"), and treating + # those as diffusion falsely rejects a valid Vulkan+gpu_ids GGUF (#7239). Normalize + # non-alphanumerics so "DiffusionGemma"/"diffusion-gemma" collapse to one token. + # The local header below stays authoritative. + name_says_diffusion = "diffusiongemma" in _re.sub(r"[^a-z0-9]+", "", identity) try: main = getattr(config, "gguf_file", None) @@ -3915,23 +3914,86 @@ def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]: if repo and variant: from hub.utils.gguf import resolve_local_gguf_path main = resolve_local_gguf_path(repo, variant) - if not main or not Path(main).is_file(): - return None - - probe = LlamaCppBackend() - probe._read_gguf_metadata(str(main)) - if probe.is_diffusion: - return True - # A successfully decoded architecture proves that this is a normal - # llama-server GGUF. No architecture means the lightweight probe could - # not establish the routing decision, so preserve the unknown state. - if getattr(probe, "_architecture", None): - return False - return None + if main and Path(main).is_file(): + # The local GGUF header is authoritative (same probe the loader uses), so + # it can't be fooled by a "diffusion"-flavored name/path. + probe = LlamaCppBackend() + probe._read_gguf_metadata(str(main)) + if probe.is_diffusion: + return True + # A decoded architecture proves a normal llama-server GGUF; no architecture + # means the probe was inconclusive, so fall through to the name hint below. + if getattr(probe, "_architecture", None): + return False except Exception as e: logger.debug("Could not identify diffusion GGUF for training guard: %s", e) + + # Header unavailable (remote uncached) or inconclusive: True only for the + # DiffusionGemma name family; otherwise None keeps an unknown remote GGUF guarded + # as potentially diffusion until its header proves otherwise. + return True if name_says_diffusion else None + + +async def _resolve_gguf_gpu_ids_for_request( + config: ModelConfig, gpu_ids: Optional[List[int]] +) -> Optional[List[int]]: + """Resolve and fully validate an explicit GGUF GPU placement pool. + + CUDA and ROCm use physical IDs. Vulkan uses ggml ordinals, so its device + existence check comes from the same ggml probe used by the loader. Both + /load and /validate call this before their training guard or any teardown. + """ + if not gpu_ids: return None + from utils.hardware import DeviceType, get_device + from utils.hardware.hardware import resolve_requested_gpu_ids + + is_vulkan = LlamaCppBackend._is_vulkan_backend() + if get_device() == DeviceType.XPU and not is_vulkan: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported on Intel XPU. " + "Omit gpu_ids to use all devices." + ), + ) + + if is_vulkan and _classify_diffusion_gguf(config) is True: + raise HTTPException( + status_code = 400, + detail = ( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." + ), + ) + + try: + resolved = resolve_requested_gpu_ids(gpu_ids, is_vulkan = is_vulkan) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = str(exc)) from exc + + if is_vulkan and resolved: + binary = LlamaCppBackend._find_llama_server_binary() + if binary: + probed = { + gpu[0] for gpu in await asyncio.to_thread(LlamaCppBackend._get_gpu_memory, binary) + } + wanted = {int(gpu_id) for gpu_id in resolved} + if not wanted.issubset(probed): + raise HTTPException( + status_code = 400, + detail = ( + f"Requested Vulkan GPU ordinal(s) {sorted(wanted)} not " + f"present. Available Vulkan devices: {sorted(probed)}." + ), + ) + + return resolved + def _guard_chat_load_against_training( config: ModelConfig, @@ -3971,8 +4033,18 @@ def _guard_chat_load_against_training( if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False: return + # Vulkan GGUF pins are ggml ordinals, not CUDA physical IDs. Detect this + # before deriving a possible diffusion fallback device so an unknown remote + # GGUF never sends its ordinal through the CUDA single-device path. + is_vulkan = False + if is_gguf: + try: + is_vulkan = LlamaCppBackend._is_vulkan_backend() + except Exception as e: + logger.warning("Could not detect Vulkan backend for chat-load guard: %s", e) + diffusion_gpu = None - if is_gguf and diffusion_kind is not False: + if is_gguf and diffusion_kind is not False and not (is_vulkan and requested_gpu_ids): # Use the same token selection as the runner: an explicit pick wins, # followed by DG_GPU, the first parent-visible token, then GPU 0. diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg( @@ -3999,6 +4071,7 @@ def _guard_chat_load_against_training( max_seq_length = max_seq_length, requested_gpu_ids = requested_gpu_ids, is_gguf = is_gguf, + is_vulkan = is_vulkan, required_override_gb = required_override_gb, single_device_gpu = diffusion_gpu, ) @@ -4305,6 +4378,7 @@ async def _load_model_impl( # Skip if a prior audio probe failed -- let load_model retry. and getattr(llama_backend, "_audio_probed", True) ): + llama_backend._record_matching_gpu_request(request.gpu_ids) logger.info( "Model already loaded (GGUF): " f"{model_log_label} variant={request.gguf_variant or llama_backend.hf_variant}, skipping reload" @@ -4351,6 +4425,7 @@ async def _load_model_impl( n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, ) else: if ( @@ -4417,41 +4492,12 @@ async def _load_model_impl( # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # GGUF supports gpu_ids: validate the pick up front (before the training - # guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects - # negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts - # are rejected outright: the picker's indices are torch-xpu ordinals neither - # applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin - # uses ggml's own Vulkan ordinals), so a pick could land on the wrong device. - if config.is_gguf and effective_gpu_ids is not None: - from utils.hardware import DeviceType, get_device - from utils.hardware.hardware import resolve_requested_gpu_ids - - if get_device() == DeviceType.XPU: - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported on Intel XPU. " - "Omit gpu_ids to use all devices." - ), - ) - # Same reasoning for a Vulkan-only build: --device pins ggml's own - # Vulkan ordinals, so a physical pick can land on the wrong card on - # masked or non-contiguous hosts. - if LlamaCppBackend._is_vulkan_backend(): - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported with a Vulkan " - "llama.cpp build: physical GPU ids have no defined " - "mapping to Vulkan device ordinals. Omit gpu_ids to use " - "all devices." - ), - ) - try: - resolve_requested_gpu_ids(effective_gpu_ids) - except ValueError as exc: - raise HTTPException(status_code = 400, detail = str(exc)) from exc + # Validate the full GGUF placement pool before the training guard so an + # invalid physical ID or Vulkan ordinal is a clean 400, not a masked VRAM + # 409. The same helper is used by /validate. + gguf_gpu_ids: Optional[List[int]] = None + if config.is_gguf: + gguf_gpu_ids = await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids) if not config.is_gguf and _mlx_distributed_launch_detected(): raise HTTPException( status_code = 400, @@ -4575,8 +4621,9 @@ async def _load_model_impl( gpu_layers = request.gpu_layers, n_cpu_moe = request.n_cpu_moe, tensor_split = request.tensor_split, - gpu_ids = effective_gpu_ids, n_parallel = _n_parallel, + # Issue #7164: explicit GPU pin resolved to physical ids above. + gpu_ids = gguf_gpu_ids, ) if config.gguf_hf_repo: # HF mode: download via huggingface_hub then start llama-server @@ -4750,6 +4797,7 @@ async def _load_model_impl( n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -5043,36 +5091,8 @@ async def validate_model( # Apply the same training coexistence policy as /load before the frontend # unloads the current model. effective_gpu_ids = request.gpu_ids if request.gpu_ids else None - # Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is - # a clean 400) before the guard sizes the model against training VRAM. - # XPU-host picks are rejected like /load (no defined mapping from the - # picker's torch-xpu ordinals to the launcher's device spaces). - if config.is_gguf and effective_gpu_ids is not None: - from utils.hardware import DeviceType, get_device - from utils.hardware.hardware import resolve_requested_gpu_ids - - if get_device() == DeviceType.XPU: - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported on Intel XPU. " - "Omit gpu_ids to use all devices." - ), - ) - if LlamaCppBackend._is_vulkan_backend(): - raise HTTPException( - status_code = 400, - detail = ( - "GPU selection (gpu_ids) is not supported with a Vulkan " - "llama.cpp build: physical GPU ids have no defined " - "mapping to Vulkan device ordinals. Omit gpu_ids to use " - "all devices." - ), - ) - try: - resolve_requested_gpu_ids(effective_gpu_ids) - except ValueError as exc: - raise HTTPException(status_code = 400, detail = str(exc)) from exc + if config.is_gguf: + await _resolve_gguf_gpu_ids_for_request(config, effective_gpu_ids) effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) # Both checks cover the [adapter, base] set (matching the scan route and workers): @@ -5897,6 +5917,7 @@ async def get_status(current_subject: str = Depends(get_current_subject)): n_layers = llama_backend.n_layers, n_moe_layers = llama_backend.n_moe_layers, gpu_ids = llama_backend.gpu_ids, + requested_gpu_ids = llama_backend.requested_gpu_ids, llama_cpp_supports_mtp = _supports_mtp, spec_fallback_reason = llama_backend.spec_fallback_reason, llama_cpp_prebuilt_stale = _stale, diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index 83bce8b772..ba49528db4 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -225,6 +225,7 @@ def can_load_chat_during_training( max_seq_length: int, requested_gpu_ids: Optional[List[int]], is_gguf: bool = False, + is_vulkan: bool = False, required_override_gb: Optional[float] = None, single_device_gpu: Optional[str] = None, ) -> Tuple[bool, Dict[str, Any]]: @@ -233,11 +234,15 @@ def can_load_chat_during_training( chat model against the free VRAM that remains). Sizes/places it the same way the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an even-share per-GPU floor for device_map="balanced"; GGUF sizes from - required_override_gb over the visible pool. ``single_device_gpu`` is the - exact physical device token selected by a single-device runner. - `load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA - allows the load; default-deny on any CUDA case it can't size, so a load never - OOMs training.""" + required_override_gb over the visible pool. A Vulkan GGUF selection picks by ggml + Vulkan ordinal (separate index space from CUDA ids), so its requested_gpu_ids is + NOT resolved against the CUDA set (which would raise -> invalid_gpu_ids -> bypass + the OOM check); conservatively size an N-device request against the least-free + N visible GPUs instead. + ``single_device_gpu`` is the exact physical device token selected by a + single-device runner. `load_in_4bit` must be effective (LoRA can flip 4-bit + -> 16-bit). Non-CUDA allows the load; default-deny on any CUDA case it can't + size, so a load never OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -258,6 +263,11 @@ def can_load_chat_during_training( max_seq_length = max_seq_length or 2048, ) + # A Vulkan GGUF selection uses ggml Vulkan ordinals, not CUDA physical ids; + # size it against the full visible pool (GGUF self-placement) rather than + # resolving ordinals against the CUDA parent-visible set. + vulkan_gguf = is_gguf and is_vulkan + # HF auto: reuse the loader's selector; fits iff its pick clears the margin. if not requested_gpu_ids and not is_gguf: _selected, meta = auto_select_gpu_ids(model_name, **est_kwargs) @@ -283,7 +293,9 @@ def can_load_chat_during_training( } # Explicit GPUs, or GGUF: size directly and check live free VRAM. - if single_device_gpu is not None: + if requested_gpu_ids and vulkan_gguf: + mode = "gguf_vulkan" + elif single_device_gpu is not None: mode = "single_device" elif is_gguf: mode = "gguf" @@ -296,7 +308,17 @@ def can_load_chat_during_training( return False, {"mode": mode, "reason": "estimate_unavailable"} free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", [])) - if single_device_gpu is not None: + if requested_gpu_ids and vulkan_gguf: + # Vulkan ordinals cannot be mapped to CUDA physical indices. Budget + # the least-free N visible cards for an N-device request. If that + # conservative subset fits, any physical mapping of the ordinals + # fits, without collapsing a multi-GPU request to one card. + visible_free = list(free_by_index.values()) + if not visible_free: + return False, {"mode": "gguf_vulkan", "reason": "no_visible_gpus"} + n_pins = min(len(requested_gpu_ids), len(visible_free)) + free_vals = sorted(visible_free)[:n_pins] + elif single_device_gpu is not None: token = str(single_device_gpu).strip() if not token: # Empty token = a CPU-only single-device runner (e.g. a CPU @@ -324,7 +346,8 @@ def can_load_chat_during_training( return True, {"mode": mode, "reason": "invalid_gpu_ids"} free_vals = [free_by_index.get(i, 0.0) for i in resolved] else: - # GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate. + # GGUF self-placement / auto Vulkan (no requested ids): llama.cpp picks + # the GPU(s), so any visible GPU is a candidate -> size the whole pool. free_vals = list(free_by_index.values()) if not free_vals: diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index a5fd71b6a0..089c5b851e 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -170,6 +170,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): estimate = None, single_device_gpu = None, gpu_ids = None, + is_vulkan = False, ): with ( patch("utils.hardware.get_device", return_value = DeviceType.CUDA), @@ -185,6 +186,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): max_seq_length = 0, requested_gpu_ids = gpu_ids, is_gguf = True, + is_vulkan = is_vulkan, required_override_gb = required_override, single_device_gpu = single_device_gpu, ) @@ -234,6 +236,35 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): self.assertFalse(blocked) self.assertEqual(blocked_info["usable_gb"], 10.0) + def test_vulkan_pin_takes_precedence_over_unknown_diffusion_fallback(self): + # An uncached GGUF can carry a speculative single-device fallback while + # its explicit pin is actually a ggml Vulkan ordinal. Never interpret + # that ordinal as the same-numbered CUDA physical device. + ok, info, _ = self._run( + devices = _devices((0, 80, 0), (1, 80, 78)), + required_override = 20.0, + single_device_gpu = "0", + gpu_ids = [0], + is_vulkan = True, + ) + self.assertFalse(ok) + self.assertEqual(info["mode"], "gguf_vulkan") + self.assertEqual(info["usable_gb"], 2.0) + + def test_vulkan_multi_gpu_guard_counts_requested_devices(self): + # The ordinal mapping is unknown, so use the least-free two visible + # cards for a two-device request. Their aggregate capacity is still + # available instead of collapsing the request to one card. + ok, info, _ = self._run( + devices = _devices((0, 80, 70), (1, 80, 70), (2, 80, 0)), + required_override = 10.0, + gpu_ids = [0, 1], + is_vulkan = True, + ) + self.assertTrue(ok) + self.assertEqual(info["mode"], "gguf_vulkan") + self.assertEqual(info["usable_gb"], 18.5) + def test_single_device_unresolved_token_sizes_against_worst_device(self): # A non-numeric device token (a CUDA UUID / MIG handle) can't map to a # free-VRAM index. The runner still drives ONE device, so size against the @@ -478,58 +509,19 @@ class TestChatLoadGuardRoute(unittest.TestCase): def test_manual_known_normal_gguf_bypasses_training_estimate(self): captured = [] config = SimpleNamespace(is_gguf = True) - with patch.object(self.route, "_classify_diffusion_gguf", return_value = False): + with patch.object(self.route, "_classify_diffusion_gguf", return_value = False) as classify: self._guard( config = config, captured = captured, training_active = True, decision = (False, {"reason": "must not run"}), gpu_memory_mode = "manual", + requested_gpu_ids = [1, 3], ) + classify.assert_called_once_with(config) self.assertEqual(captured, []) - def test_manual_unknown_gguf_keeps_single_device_training_guard(self): - captured = [] - config = SimpleNamespace(is_gguf = True) - with ( - patch.object(self.route, "_classify_diffusion_gguf", return_value = None), - patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), - patch.object( - self.route.LlamaCppBackend, - "_diffusion_gpu_arg", - return_value = "2", - ), - ): - self._guard( - config = config, - captured = captured, - training_active = True, - decision = (True, {"mode": "single_device"}), - gpu_memory_mode = "manual", - ) - self.assertEqual(len(captured), 1) - self.assertEqual(captured[0]["single_device_gpu"], "2") - - def test_manual_diffusion_uses_single_device_guard(self): - captured = [] - config = SimpleNamespace(is_gguf = True) - with ( - patch.object(self.route, "_classify_diffusion_gguf", return_value = True), - patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5), - ): - self._guard( - config = config, - captured = captured, - training_active = True, - decision = (True, {"mode": "gguf"}), - gpu_memory_mode = "manual", - requested_gpu_ids = [3, 1], - ) - self.assertEqual(len(captured), 1) - self.assertEqual(captured[0]["single_device_gpu"], "1") - self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) - - def test_unpinned_diffusion_uses_runner_default_gpu(self): + def test_manual_diffusion_keeps_single_device_training_guard(self): captured = [] config = SimpleNamespace(is_gguf = True) with ( @@ -540,11 +532,6 @@ class TestChatLoadGuardRoute(unittest.TestCase): "_effective_gpu_count", return_value = 2, ), - patch.object( - self.route.LlamaCppBackend, - "_diffusion_gpu_arg", - return_value = "3", - ) as gpu_arg, ): self._guard( config = config, @@ -552,9 +539,11 @@ class TestChatLoadGuardRoute(unittest.TestCase): training_active = True, decision = (True, {"mode": "single_device"}), gpu_memory_mode = "manual", + requested_gpu_ids = [3, 1], ) - gpu_arg.assert_called_once_with(None, cpu_only = False) - self.assertEqual(captured[0]["single_device_gpu"], "3") + self.assertEqual(len(captured), 1) + self.assertEqual(captured[0]["single_device_gpu"], "1") + self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1]) def test_refuses_with_headroom_number(self): info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"} diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index 19ba9e3e05..271a882b11 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -591,10 +591,23 @@ def test_load_request_accepts_gpu_ids(): @pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse]) def test_response_models_emit_gpu_ids(model_cls): if model_cls is LoadResponse: - obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1]) + obj = model_cls( + status = "loaded", + model = "m", + display_name = "m", + inference = {}, + gpu_ids = [1], + requested_gpu_ids = [1, 2], + ) else: - obj = model_cls(gpu_ids = [1]) + obj = model_cls(gpu_ids = [1], requested_gpu_ids = [1, 2]) assert obj.model_dump()["gpu_ids"] == [1] + assert obj.model_dump()["requested_gpu_ids"] == [1, 2] + + +def test_gguf_load_and_status_responses_include_requested_gpu_pool(): + route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") + assert route_src.count("requested_gpu_ids = llama_backend.requested_gpu_ids") == 3 def test_gpu_ids_property_default_and_reset(): @@ -625,6 +638,10 @@ def _target_state_gpu_ids(backend, gpu_ids): def test_gpu_ids_reload_detection_is_order_insensitive(): backend = _loaded_backend("auto") backend._gpu_ids = [0, 1] + # A real non-narrowed load records the raw request too; the non-diffusion + # dedupe now compares that raw pin (#7239). Set it to match the effective pin + # (no narrowing) so this exercises the order-insensitive comparison. + backend._requested_gpu_ids = [0, 1] # Same set, different order -> no reload. assert _target_state_gpu_ids(backend, [1, 0]) is True # Different set -> reload. @@ -633,6 +650,26 @@ def test_gpu_ids_reload_detection_is_order_insensitive(): assert _target_state_gpu_ids(backend, None) is False +def test_gpu_ids_reload_detection_accepts_raw_and_effective_pin(): + backend = _loaded_backend("auto") + backend._requested_gpu_ids = [0, 1] + backend._gpu_ids = [0] + backend._last_load_kwargs = {"gpu_ids": [0, 1], "model_identifier": "owner/repo"} + + # The original request still matches after the fitter narrows it. + assert _target_state_gpu_ids(backend, [1, 0]) is True + assert backend.requested_gpu_ids == [0, 1] + # The status response echoes the effective pin, which must also round-trip. + # Treat the incoming subset as the latest intent so status and a future + # reload do not restore GPU 1 after the user removed it. + assert _target_state_gpu_ids(backend, [0]) is True + assert backend.requested_gpu_ids == [0] + assert backend._last_load_kwargs == {"gpu_ids": [0], "model_identifier": "owner/repo"} + # A genuinely different placement pool still reloads. + assert _target_state_gpu_ids(backend, [1]) is False + assert _target_state_gpu_ids(backend, None) is False + + def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): # The diffusion runner drives only its single lowest device, so the backend # records [lowest]. A later multi-GPU request that still resolves to that @@ -642,6 +679,7 @@ def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): backend._is_diffusion = True backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick assert _target_state_gpu_ids(backend, [3, 1]) is True + assert backend.requested_gpu_ids == [1] assert _target_state_gpu_ids(backend, [1]) is True # Lowest device changes (2, not 1) -> reload. assert _target_state_gpu_ids(backend, [3, 2]) is False @@ -649,6 +687,56 @@ def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device(): assert _target_state_gpu_ids(backend, None) is False +def test_remote_vulkan_diffusion_preflight_runs_before_teardown(monkeypatch): + def _mark_diffusion(probe, path): + assert path == "/cache/model.gguf" + probe._is_diffusion = True + + monkeypatch.setattr(LlamaCppBackend, "_read_gguf_metadata", _mark_diffusion) + assert LlamaCppBackend._gguf_path_is_diffusion("/cache/model.gguf", "owner/model") is True + + src = inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + preflight = src.index("_preflight_model_path = self._download_gguf(") + teardown = src.index("# ── Phase 1: kill old process") + assert preflight < teardown + assert "model_path = _preflight_model_path or self._download_gguf(" in src + + +def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch): + backend = LlamaCppBackend() + killed = [] + monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama") + monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True) + monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)]) + monkeypatch.setattr( + backend, + "_download_gguf", + lambda **_kwargs: "/cache/diffusion.gguf", + ) + monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: True) + monkeypatch.setattr(backend, "_kill_process", lambda: killed.append(True)) + monkeypatch.setattr( + llama_cpp_module, + "_resolve_repo_id_casing", + lambda repo: repo, + ) + monkeypatch.setattr( + llama_cpp_module, + "_hf_offline_if_dns_dead", + lambda: __import__("contextlib").nullcontext(), + ) + + with pytest.raises(ValueError, match = "DiffusionGemma"): + backend.load_model( + hf_repo = "owner/model", + hf_variant = "Q4_K_M", + model_identifier = "owner/model", + gpu_ids = [0], + ) + + assert killed == [] + + def test_start_diffusion_server_resets_tensor_parallel(): # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model # phase 1 only kills the process, it skips the unload reset). Diffusion is never @@ -656,18 +744,16 @@ def test_start_diffusion_server_resets_tensor_parallel(): # diffusion re-Apply reloads against stale tensor-parallel state. src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server) assert "self._tensor_parallel = False" in src + assert "self._requested_gpu_ids = list(self._gpu_ids) if self._gpu_ids else None" in src -def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids(): - # The route-level reload dedupe mirrors the backend: for a loaded diffusion - # model it compares the request against the single recorded device, not the - # full requested list, or a same-device multi-GPU pick reloads needlessly. +def test_route_matches_loaded_settings_uses_shared_gpu_pin_matcher(): + # Route-level and backend race dedupe must share one normalization path so + # raw, effective, and diffusion pins cannot drift apart. route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8") match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :] - guard = match_impl.index("if llama_backend.is_diffusion:") - collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None") - compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:") - assert guard < collapse < compare + assert "if not llama_backend.matches_gpu_ids(request.gpu_ids):" in match_impl + assert "llama_backend._record_matching_gpu_request(request.gpu_ids)" in match_impl # ── Manual tensor split: child enumeration pinned to the picker's order ────── diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index d4f2fbe993..3b44c19e26 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -130,6 +130,26 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): ): self.assertEqual(resolve_requested_gpu_ids([]), [1, 3]) + def test_vulkan_ordinals_bypass_cuda_parent_visible_validation(self): + # Vulkan build on a CPU-only torch host: no CUDA parent-visible set and a + # zero physical count, yet a valid Vulkan ordinal must not be rejected as + # a CUDA physical id (issue #7239). + with ( + patch.dict(os.environ, {}, clear = True), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 0), + ): + # As a CUDA physical id, [0] is outside the empty parent-visible set. + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([0]) + # As Vulkan ordinals, [0] and [0, 1] pass through unchanged. + self.assertEqual(resolve_requested_gpu_ids([0], is_vulkan = True), [0]) + self.assertEqual(resolve_requested_gpu_ids([0, 1], is_vulkan = True), [0, 1]) + # Malformed ordinals are still rejected. + with self.assertRaisesRegex(ValueError, "duplicate GPU IDs"): + resolve_requested_gpu_ids([0, 0], is_vulkan = True) + with self.assertRaisesRegex(ValueError, "non-negative"): + resolve_requested_gpu_ids([-1], is_vulkan = True) + def test_apply_gpu_ids_only_updates_cuda_visible_devices(self): with patch.dict( os.environ, @@ -853,6 +873,171 @@ class TestRouteErrors(unittest.TestCase): self.assertIn("only supported on CUDA devices", str(exc_info.exception)) + def test_inference_route_resolves_gguf_gpu_ids(self): + # GGUF gpu_ids are now supported: /load routes them through the same + # resolution as non-GGUF loads (rejecting only genuinely invalid ids with + # the resolver's actionable message) rather than a blanket "not supported" + # reject, so /validate can stay consistent with /load (#7239). + import utils.hardware.hardware as hardware_mod + + inference_route = _load_route_module( + "inference_route_module_for_gguf_gpu_ids_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1]) + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = None, + gguf_file = "/tmp/test.gguf", + gguf_mmproj_file = None, + gguf_variant = None, + identifier = "unsloth/test.gguf", + display_name = "unsloth/test.gguf", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + def _fake_resolve(ids, is_vulkan = False): + raise ValueError("SENTINEL requested GPUs are outside the parent-visible set") + + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + # Patch both the package re-export and the defining module so the stub + # fires no matter which import path the route uses. + patch("utils.hardware.resolve_requested_gpu_ids", _fake_resolve), + patch.object(hardware_mod, "resolve_requested_gpu_ids", _fake_resolve), + 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_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + + # The selection was routed through resolution (not the old blanket reject). + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("SENTINEL", exc_info.exception.detail) + self.assertNotIn("not supported for GGUF", exc_info.exception.detail) + + def test_load_rejects_unavailable_vulkan_ordinal_before_training_guard(self): + inference_route = _load_route_module( + "inference_route_module_for_vulkan_preflight_test", + "routes/inference.py", + ) + request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [99]) + model_config = SimpleNamespace( + is_gguf = True, + is_lora = False, + gguf_hf_repo = None, + gguf_file = "/tmp/test.gguf", + gguf_mmproj_file = None, + gguf_variant = None, + identifier = "unsloth/test.gguf", + display_name = "unsloth/test.gguf", + is_vision = False, + is_audio = False, + audio_type = None, + has_audio_input = False, + ) + + with ( + patch.object( + inference_route, + "ModelConfig", + SimpleNamespace(from_identifier = lambda **_kwargs: model_config), + ), + patch("utils.hardware.get_device", return_value = DeviceType.CUDA), + patch.object(inference_route, "_classify_diffusion_gguf", return_value = None), + patch.object( + inference_route.LlamaCppBackend, + "_is_vulkan_backend", + return_value = True, + ), + patch.object( + inference_route.LlamaCppBackend, + "_find_llama_server_binary", + return_value = "/tmp/llama-server", + ), + patch.object( + inference_route.LlamaCppBackend, + "_get_gpu_memory", + return_value = [(0, 8 * 1024**3, 16 * 1024**3)], + ), + patch.object( + inference_route, + "_guard_chat_load_against_training", + return_value = None, + ) as training_guard, + 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_impl( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) + ) + + self.assertEqual(exc_info.exception.status_code, 400) + self.assertIn("Vulkan GPU ordinal(s) [99]", exc_info.exception.detail) + training_guard.assert_not_called() + + def test_vulkan_ordinals_are_allowed_on_xpu_hosts(self): + import utils.hardware.hardware as hardware_mod + + inference_route = _load_route_module( + "inference_route_module_for_xpu_vulkan_test", + "routes/inference.py", + ) + config = SimpleNamespace(is_gguf = True) + + with ( + patch("utils.hardware.get_device", return_value = DeviceType.XPU), + patch.object( + inference_route.LlamaCppBackend, + "_is_vulkan_backend", + return_value = True, + ), + patch.object(inference_route, "_classify_diffusion_gguf", return_value = False), + patch.object(hardware_mod, "resolve_requested_gpu_ids", return_value = [0, 1]), + patch.object( + inference_route.LlamaCppBackend, + "_find_llama_server_binary", + return_value = None, + ), + ): + resolved = asyncio.run( + inference_route._resolve_gguf_gpu_ids_for_request(config, [1, 0]) + ) + + self.assertEqual(resolved, [0, 1]) + def test_inference_route_validates_gpu_ids_for_gguf(self): # gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still # validated: a rejected pick surfaces as a clean 400, not the old @@ -861,7 +1046,7 @@ class TestRouteErrors(unittest.TestCase): import utils.hardware.hardware as hardware_mod inference_route = _load_route_module( - "inference_route_module_for_gguf_gpu_ids_test", + "inference_route_module_for_gguf_gpu_ids_test2", "routes/inference.py", ) request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1]) @@ -886,6 +1071,17 @@ class TestRouteErrors(unittest.TestCase): "ModelConfig", SimpleNamespace(from_identifier = lambda **_kwargs: model_config), ), + # Patch both the package re-export and the defining module so the stub + # fires no matter which import path the route uses. + patch( + "utils.hardware.resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), + patch.object( + hardware_mod, + "resolve_requested_gpu_ids", + side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), + ), patch.object( inference_route, "_guard_chat_load_against_training", @@ -893,11 +1089,6 @@ class TestRouteErrors(unittest.TestCase): ), patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread), patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext), - patch.object( - hardware_mod, - "resolve_requested_gpu_ids", - side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"), - ), ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 3d312d4b01..d9a06fb017 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -1639,17 +1639,34 @@ def get_parent_visible_gpu_ids() -> list[int]: return list(parent_visible_ids) if parent_visible_ids is not None else [] -def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: +def resolve_requested_gpu_ids( + gpu_ids: Optional[list[int]], *, is_vulkan: bool = False +) -> list[int]: parent_visible_spec = _get_parent_visible_gpu_spec() parent_visible_ids = get_parent_visible_gpu_ids() physical_gpu_count = get_physical_gpu_count() if gpu_ids is None: - return parent_visible_ids + return [] if is_vulkan else parent_visible_ids requested_ids = list(gpu_ids) if len(requested_ids) == 0: - return parent_visible_ids + return [] if is_vulkan else parent_visible_ids + + if is_vulkan: + # A Vulkan build selects by ggml Vulkan ordinal (--device VulkanN), a separate + # index space from CUDA/ROCm ids that may be empty under CPU-only torch. The + # CUDA parent-visible / physical-count checks below do not apply; only reject + # malformed ordinals (issue #7239). + if len(set(requested_ids)) != len(requested_ids): + raise ValueError(f"Invalid gpu_ids {requested_ids}: duplicate GPU IDs are not allowed.") + negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0] + if negative_ids: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: GPU IDs must be non-negative. " + f"Rejected IDs: {negative_ids}." + ) + return requested_ids if not parent_visible_spec["supports_explicit_gpu_ids"]: raise ValueError( @@ -2193,12 +2210,13 @@ def auto_select_gpu_ids( metadata["selection_mode"] = "auto" metadata["selected_gpu_ids"] = selected logger.debug( - "Selected GPUs automatically", - model_name = model_name, - selected_gpu_ids = selected, - usable_gb = metadata["usable_gb"], - required_gb = metadata.get("required_gb"), - multi_gpu_overhead = multi_gpu_overhead, + "Selected GPUs automatically: model=%s selected=%s usable_gb=%s " + "required_gb=%s multi_gpu_overhead=%s", + model_name, + selected, + metadata["usable_gb"], + metadata.get("required_gb"), + multi_gpu_overhead, ) return selected, metadata @@ -2214,12 +2232,13 @@ def auto_select_gpu_ids( metadata["usable_gb"] = round(fallback_usable, 3) metadata["selected_gpu_ids"] = fallback_all logger.warning( - "Falling back to all visible GPUs -- model may not fit", - model_name = model_name, - selected_gpu_ids = fallback_all, - usable_gb = metadata["usable_gb"], - required_gb = metadata.get("required_gb"), - multi_gpu_overhead = multi_gpu_overhead, + "Falling back to all visible GPUs; model may not fit: model=%s " + "selected=%s usable_gb=%s required_gb=%s multi_gpu_overhead=%s", + model_name, + fallback_all, + metadata["usable_gb"], + metadata.get("required_gb"), + multi_gpu_overhead, ) return fallback_all, metadata diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 547c3f374e..f85ff3246b 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -222,7 +222,9 @@ export function applyActiveModelStatusToStore( incomingGpuMode === "manual" ? (status.n_cpu_moe ?? null) : null; const incomingSplit = incomingGpuMode === "manual" ? (status.tensor_split ?? null) : null; - const incomingGpuIds = status.is_gguf ? (status.gpu_ids ?? null) : null; + const incomingGpuIds = status.is_gguf + ? (status.requested_gpu_ids ?? status.gpu_ids ?? null) + : null; const gpuStatusChanged = prevState.loadedGpuMemoryMode !== incomingGpuMode || prevState.loadedGpuLayers !== incomingGpuLayers || 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 89bc21ee18..15fa962dff 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -605,6 +605,7 @@ export function loadedGpuMemoryFields(resp: { n_layers?: number | null; n_moe_layers?: number; gpu_ids?: number[] | null; + requested_gpu_ids?: number[] | null; }) { // GPU-memory state is meaningful only for a GGUF chat load. A non-GGUF response // still carries gpu_memory_mode (its default "auto" is serialized), so gate on @@ -631,7 +632,9 @@ export function loadedGpuMemoryFields(resp: { }; } const mode = resp.gpu_memory_mode ?? "auto"; - const gpuIds = resp.gpu_ids ?? null; + // Keep the user's placement pool editable across status/load hydration. + // gpu_ids remains the effective fitted subset for diagnostics. + const gpuIds = resp.requested_gpu_ids ?? resp.gpu_ids ?? null; // Layer/MoE/split knobs apply (and are reported) only in manual mode; in auto // the server ignores them, so don't seed the loaded baseline or the editable // knobs with values it never applied. In manual, the server reports gpu_layers @@ -669,7 +672,7 @@ export function loadedGpuMemoryFields(resp: { ggufLayerCount: resp.n_layers ?? null, // MoE expert-layer count: the n_cpu_moe slider max, and 0 hides the slider. moeLayerCount: resp.n_moe_layers ?? null, - // The picker reflects what loaded (the request sent the user's pick). + // The picker reflects the requested placement pool, not a fitted subset. selectedGpuIds: gpuIds, loadedGpuIds: gpuIds, ...manualKnobs, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index c9c06834c1..6c3e919efe 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -188,7 +188,10 @@ export interface LoadModelResponse { n_layers?: number | null; /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ n_moe_layers?: number; + /** Effective GPU placement after fit-time narrowing. */ gpu_ids?: number[] | null; + /** User-requested GPU placement pool before fit-time narrowing. */ + requested_gpu_ids?: number[] | null; } export interface UnloadModelRequest { @@ -240,7 +243,10 @@ export interface InferenceStatusResponse { /** n_ctx the active GGUF load was invoked with (0 = Auto); re-seeds a * Manual + Auto-layers context pin on hydration. Null for non-GGUF. */ requested_context_length?: number | null; + /** Effective GPU placement after fit-time narrowing. */ gpu_ids?: number[] | null; + /** User-requested GPU placement pool before fit-time narrowing. */ + requested_gpu_ids?: number[] | null; n_layers?: number | null; /** Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not MoE. */ n_moe_layers?: number; diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index baf0fdf1bf..62b1a3cf76 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -151,6 +151,19 @@ def test_active_model_config_round_trips_gpu_fields(): assert "export function gpuFieldsSignature" in shared +def test_gpu_picker_round_trips_requested_pool_not_fitted_subset(): + """A GGUF fit may narrow [0, 1] to [0], but load/status hydration must keep + [0, 1] as the editable pool so a later reload can grow back onto GPU 1.""" + types = _read("features/chat/types/api.ts") + assert types.count("requested_gpu_ids?: number[] | null") >= 2 + + store = _read("features/chat/stores/chat-runtime-store.ts") + assert "resp.requested_gpu_ids ?? resp.gpu_ids ?? null" in store + + status = _read("features/chat/lib/apply-inference-status-to-store.ts") + assert "status.requested_gpu_ids ?? status.gpu_ids ?? null" in status + + def test_compare_load_uses_each_models_gpu_config(): src = _read("features/chat/shared-composer.tsx") assert "ownConfig.gpuMemoryMode ?? compareLoadKnobs.gpuMemoryMode" in src From 707b74fac38e518b5fe57a4bc69d50eca7f32bb1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 21:54:41 -0700 Subject: [PATCH 017/161] Studio UI test: recover from voice-picker renderer crash, scoped to macOS runners Downgrades a headless-Chromium renderer crash in the voice model-picker step to a warning plus page recovery on macos-14, where CheckMediaAccessPermission can kill the tab. Linux and Windows strict smoke jobs keep hard crash coverage and any live-page failure stays a hard fail. --- tests/studio/playwright_extra_ui.py | 57 +++++++++++++++++++++++------ 1 file changed, 46 insertions(+), 11 deletions(-) diff --git a/tests/studio/playwright_extra_ui.py b/tests/studio/playwright_extra_ui.py index 0da028a28c..e4de62fd2e 100644 --- a/tests/studio/playwright_extra_ui.py +++ b/tests/studio/playwright_extra_ui.py @@ -36,6 +36,9 @@ ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_extra") ART = Path(ART_DIR) ART.mkdir(parents = True, exist_ok = True) STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1" +# The Voice-picker media-access crash is specific to headless Chromium on macos-14; only there do we +# downgrade a renderer crash to a warning. Linux/Windows strict smoke jobs keep hard crash coverage. +MACOS_RUNNER = os.environ.get("RUNNER_OS", "").lower() == "macos" or sys.platform == "darwin" # Longer turn timeout: gemma-3-270m CPU inference is 3-5x slower on macos-14 runners. TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000")) WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720")) @@ -71,6 +74,18 @@ def runtime_warn(m: str) -> None: info(f"WARN (runtime): {m}") +def page_crashed(pg, exc: Exception) -> bool: + """True when the browser/page/context died (a macos-14 renderer crash) rather than a live-page + assertion failing -- so the caller can downgrade CI-environment flakiness to a runtime warning.""" + try: + if pg.is_closed(): + return True + except Exception: + return True + msg = str(exc).lower() + return "has been closed" in msg or "target closed" in msg or "crash" in msg + + with sync_playwright() as p: _watchdog = install_wall_clock_watchdog( WALL_TIMEOUT_S, @@ -544,13 +559,17 @@ with sync_playwright() as p: if voice_tab.count() == 0: fail("Voice settings tab not found") else: - voice_tab.click() - page.get_by_label("Dictation engine").click() - page.get_by_role("option", name = "Local transcription").click() - page.get_by_label("Speech recognition model").click() - page.get_by_placeholder("Search model").fill("whisper") - results = page.get_by_test_id("stt-model-results") + # The dictation-engine dropdown touches a media-access path that can crash headless + # Chromium on macos-14 (CheckMediaAccessPermission). A resulting TargetClosedError is CI + # flakiness there, not a product bug, so on macOS a crash is a runtime warning + page + # recovery; on Linux/Windows a crash and any live-page failure stay a hard fail. try: + voice_tab.click() + page.get_by_label("Dictation engine").click() + page.get_by_role("option", name = "Local transcription").click() + page.get_by_label("Speech recognition model").click() + page.get_by_placeholder("Search model").fill("whisper") + results = page.get_by_test_id("stt-model-results") page.wait_for_function( """() => { const node = document.querySelector('[data-testid="stt-model-results"]'); @@ -569,10 +588,23 @@ with sync_playwright() as p: ) info("OK Voice model picker mouse wheel changed scrollTop") except Exception as exc: - fail(f"Voice model picker did not wheel-scroll: {exc!r}") - shoot("10-settings-tabs-visited") - page.keyboard.press("Escape") - page.wait_for_timeout(300) + if page_crashed(page, exc) and MACOS_RUNNER: + runtime_warn(f"Voice model picker aborted (browser/page unstable): {exc!r}") + page = recover_or_replace_page( + page, + ctx, + default_timeout_ms = 60_000, + info = lambda m: info(f"recovery: {m}"), + ) + else: + fail(f"Voice model picker did not wheel-scroll: {exc!r}") + # When the crash closed the context/browser (not just the page), recover_or_replace_page + # cannot mint a replacement and hands back the closed page; skip the cosmetic teardown rather + # than re-raise TargetClosedError on it. is_closed() is a local check and never raises. + if not page.is_closed(): + shoot("10-settings-tabs-visited") + page.keyboard.press("Escape") + page.wait_for_timeout(300) info(f"visited Settings tabs: {seen_tabs}") if not seen_tabs: soft_fail("no Settings tabs were visitable") @@ -591,4 +623,7 @@ with sync_playwright() as p: sys.exit(1) info("PASS extra UI flow") _watchdog.cancel() - browser.close() + try: + browser.close() + except Exception: + pass # a crashed browser may already be gone; never fail teardown after PASS From 63d8da34d3476f191a6249ca13a249961885d9e1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 21:55:13 -0700 Subject: [PATCH 018/161] Studio: use text-ui-* tokens instead of raw px in the voice tab Replaces the raw text-[9px] and text-[10px] classes with the text-ui-9 and text-ui-10 scale tokens so the voice tab labels honor the --ui-font-scale typography setting like the rest of the UI. --- studio/frontend/src/features/settings/tabs/voice-tab.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/features/settings/tabs/voice-tab.tsx b/studio/frontend/src/features/settings/tabs/voice-tab.tsx index 59c03bdf14..2dc3d01de8 100644 --- a/studio/frontend/src/features/settings/tabs/voice-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/voice-tab.tsx @@ -278,13 +278,13 @@ function SttModelPicker({ {sttModelName(model)} {twoLines ? ( - + {sttModelSource(model)} ) : null} {sttModelSize(model) ? ( - + {sttModelSize(model)} ) : null} From 47fa4ca6c156c8b9c05663354ee52704213a3cce Mon Sep 17 00:00:00 2001 From: Lei Zhenyuan Date: Fri, 24 Jul 2026 13:22:07 +0800 Subject: [PATCH 019/161] Add Intel XPU support to Unsloth Studio (#4724) --------- Co-authored-by: Daniel Han Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .github/workflows/studio-backend-ci.yml | 12 +- studio/backend/core/inference/inference.py | 25 +- studio/backend/core/inference/orchestrator.py | 4 +- studio/backend/core/inference/worker.py | 2 +- studio/backend/core/training/trainer.py | 16 +- studio/backend/core/training/training.py | 5 +- studio/backend/core/training/worker.py | 2 +- studio/backend/models/inference.py | 11 +- studio/backend/models/training.py | 10 +- studio/backend/routes/training_vram.py | 12 +- .../tests/test_chat_load_during_training.py | 27 +- studio/backend/tests/test_gpu_selection.py | 68 ++- .../tests/test_gpu_selection_sandbox.py | 4 +- .../tests/test_training_vram_coexistence.py | 13 +- studio/backend/utils/hardware/__init__.py | 6 + studio/backend/utils/hardware/hardware.py | 447 +++++++++++++-- studio/backend/utils/utils.py | 28 +- tests/studio/test_xpu_spoof_pipeline.py | 538 ++++++++++++++++++ 18 files changed, 1141 insertions(+), 89 deletions(-) create mode 100644 tests/studio/test_xpu_spoof_pipeline.py diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index d926d5c3e4..3968f2e80a 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -193,6 +193,7 @@ jobs: --ignore=tests/sh \ --ignore=tests/studio/test_hardware_dispatch_matrix.py \ --ignore=tests/studio/test_is_mlx_dispatch_gate.py \ + --ignore=tests/studio/test_xpu_spoof_pipeline.py \ --ignore=tests/vllm_compat \ --ignore=tests/version_compat \ -m 'not server and not e2e' \ @@ -205,14 +206,15 @@ jobs: env: PYTHONPATH: ${{ github.workspace }}/studio UNSLOTH_COMPILE_DISABLE: '1' - # These two files mutate hardware.py module globals at runtime - # via the spoof fixtures, which leaks state into any other test - # that imports hardware. Run them in their own pytest invocation - # so the leak does not cross file boundaries. + # These files mutate hardware.py module globals at runtime via the + # spoof fixtures (CUDA/ROCm/XPU/MLX/CPU), which leaks state into any + # other test that imports hardware. Run them in their own pytest + # invocation so the leak does not cross file boundaries. run: | python -m pytest -q --tb=short \ tests/studio/test_hardware_dispatch_matrix.py \ - tests/studio/test_is_mlx_dispatch_gate.py + tests/studio/test_is_mlx_dispatch_gate.py \ + tests/studio/test_xpu_spoof_pipeline.py - name: Shell installer tests # Subset that does not depend on a writable / pristine install.sh diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 8d262bbb0f..2f46470091 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -8,6 +8,7 @@ from unsloth.chat_templates import get_chat_template from transformers import TextIteratorStreamer, TextStreamer from peft import PeftModel, PeftModelForCausalLM +import contextlib import json import sys import torch @@ -1942,8 +1943,30 @@ class InferenceBackend: + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n" ) + with torch.inference_mode(): - with torch.amp.autocast("cuda", dtype = model.dtype): + # Derive the autocast device from the loaded model, not from the + # global backend: a CPU-fallback DAC on an XPU/CUDA host must not + # open a GPU autocast context around CPU tensors. + device_type = ( + model.device.type + if hasattr(model.device, "type") + else str(model.device).split(":", 1)[0] + ) + # Clamp to autocast-supported backends so exotic devices + # (e.g. "meta" during accelerate offloaded loading) do not raise. + # MPS is autocast-supported since torch 2.3, keep it in the set. + if device_type not in ("cuda", "xpu", "mps", "cpu"): + device_type = "cpu" + # CPU and XPU autocast only accept bfloat16/float16. For a + # float32 model, skip autocast entirely to avoid raising or + # producing a warning on every generate call. + autocast_dtype_supported = model.dtype in (torch.bfloat16, torch.float16) + if device_type in ("cpu", "xpu") and not autocast_dtype_supported: + autocast_ctx = contextlib.nullcontext() + else: + autocast_ctx = torch.amp.autocast(device_type, dtype = model.dtype) + with autocast_ctx: inputs = tokenizer([prompt], return_tensors = "pt").to(model.device) generated = model.generate( **inputs, diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 409132d605..616384386d 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -27,7 +27,7 @@ import uuid from io import BytesIO from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union -from utils.hardware import prepare_gpu_selection +from utils.hardware import get_device, prepare_gpu_selection # Re-exported from the shared helper so GGUF, training, and inference share one # type; kept importable here for backwards compatibility. @@ -1012,6 +1012,8 @@ class InferenceOrchestrator: ) sub_config["resolved_gpu_ids"] = resolved_gpu_ids sub_config["gpu_selection"] = gpu_selection + # Parent-detected backend for the worker's apply_gpu_ids(). + sub_config["device_backend"] = get_device().value # Recheck the sidecar reservation BEFORE tearing the old worker down, # for REPAIRS only: an install holds this same lifecycle gate, so it diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 9f301ba37e..367de196f7 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -794,7 +794,7 @@ def run_inference_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) - apply_gpu_ids(config.get("resolved_gpu_ids")) + apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) model_name = config["model_name"] diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 8e419849cb..53cba522ed 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -1481,6 +1481,9 @@ class UnslothTrainer: SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz" SNAC_SAMPLE_RATE = 24000 + + # SNAC codec unvalidated on Intel XPU; keep the pre-PR CPU + # fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" max_length = self.max_seq_length or 2048 tokenizer = self.tokenizer @@ -1642,7 +1645,8 @@ class UnslothTrainer: del snac_model gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: @@ -1669,6 +1673,8 @@ class UnslothTrainer: import numpy as np import torchaudio.transforms as T + # Spark-TTS BiCodec unvalidated on Intel XPU; keep the pre-PR CPU + # fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" # sparktts lives in the SparkAudio/Spark-TTS GitHub repo, not the HF model @@ -1857,7 +1863,8 @@ class UnslothTrainer: del audio_tokenizer gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: @@ -1894,6 +1901,8 @@ class UnslothTrainer: from datasets import Dataset as HFDataset from utils.paths import ensure_dir, tmp_root + # OuteTTS DAC/Whisper preprocess unvalidated on Intel XPU; keep the + # pre-PR CPU fallback for non-CUDA hosts. device = "cuda" if torch.cuda.is_available() else "cpu" # Clone OuteTTS repo (same as audio_codecs._load_dac) @@ -2065,7 +2074,8 @@ class UnslothTrainer: del prompt_processor gc.collect() - torch.cuda.empty_cache() + + clear_gpu_cache() self._cuda_audio_used = True if not processed_examples: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index bfbd11a427..b87585ffe3 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -30,7 +30,7 @@ from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt -from utils.hardware import prepare_gpu_selection +from utils.hardware import get_device, prepare_gpu_selection from utils.native_path_leases import ( native_path_secret_removed_for_child_start, run_without_native_path_secret, @@ -219,6 +219,9 @@ def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: config[key] = values.get(key) if config["training_type"] == "Full Finetuning": config["load_in_4bit"] = False + # The parent's detected backend: the worker's apply_gpu_ids() targets the + # right visibility env var from this, without probing torch pre-mask. + config["device_backend"] = get_device().value return config diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 5ded18ea45..81010fb5df 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -2373,7 +2373,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> env = os.getenv("ENVIRONMENT_TYPE", "production"), ) - apply_gpu_ids(config.get("resolved_gpu_ids")) + apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) model_name = config["model_name"] diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 099d356a73..2663242187 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -74,7 +74,16 @@ class LoadRequest(BaseModel): ) gpu_ids: Optional[List[int]] = Field( None, - description = "GPU placement pool, for example [0, 1]. Omit or pass [] to use automatic selection. CUDA/ROCm values are physical GPU indices and are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries; Vulkan values are ggml device ordinals. For GGUF models the fitter may pin the smallest subset of this pool that fits.", + description = ( + "GPU placement pool, for example [0, 1]. Omit or pass [] to use " + "automatic selection. CUDA/ROCm and Intel XPU values are physical " + "GPU indices; Vulkan values are ggml device ordinals. Explicit " + "physical IDs are unsupported when the parent visibility mask uses " + "non-numeric or subdevice entries, including CUDA_VISIBLE_DEVICES " + "with UUID/MIG entries and ZE_AFFINITY_MASK with subdevice tokens " + "(for example '0.0,0.1') or FLAT-hierarchy tile handles. For GGUF " + "models the fitter may pin the smallest subset of this pool that fits." + ), ) speculative_type: Optional[str] = Field( None, diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 0f88b78f9f..6416b86069 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -496,7 +496,15 @@ class TrainingStartRequest(BaseModel): # GPU selection gpu_ids: Optional[List[int]] = Field( None, - description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.", + description = ( + "Physical GPU indices to use, for example [0, 1]. Omit or pass " + "[] to use automatic selection. Explicit gpu_ids are unsupported " + "when the parent visibility mask uses non-numeric or subdevice " + "entries -- this includes CUDA_VISIBLE_DEVICES with UUID/MIG " + "entries on NVIDIA, and ZE_AFFINITY_MASK with subdevice tokens " + "(e.g. '0.0,0.1') or FLAT-hierarchy (default) tile handles on " + "Intel XPU." + ), ) # S3 dataset source configuration diff --git a/studio/backend/routes/training_vram.py b/studio/backend/routes/training_vram.py index ba49528db4..8ddda11b1e 100644 --- a/studio/backend/routes/training_vram.py +++ b/studio/backend/routes/training_vram.py @@ -135,8 +135,8 @@ def can_keep_chat_during_training( resolve_requested_gpu_ids, ) - if get_device() != DeviceType.CUDA: - return False, {"mode": "non_cuda", "reason": "non_cuda"} + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + return False, {"mode": "non_accelerator", "reason": "non_accelerator"} # Full finetuning runs in 16-bit, so ignore the 4-bit request or we under-count. effective_4bit = False if training_type == "Full Finetuning" else load_in_4bit @@ -241,8 +241,8 @@ def can_load_chat_during_training( N visible GPUs instead. ``single_device_gpu`` is the exact physical device token selected by a single-device runner. `load_in_4bit` must be effective (LoRA can flip 4-bit - -> 16-bit). Non-CUDA allows the load; default-deny on any CUDA case it can't - size, so a load never OOMs training.""" + -> 16-bit). CPU/MLX allows the load; default-deny on any CUDA/XPU case it + can't size, so a load never OOMs training.""" try: from utils.hardware import ( DeviceType, @@ -253,8 +253,8 @@ def can_load_chat_during_training( resolve_requested_gpu_ids, ) - if get_device() != DeviceType.CUDA: - return True, {"mode": "non_cuda", "reason": "non_cuda"} + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + return True, {"mode": "non_accelerator", "reason": "non_accelerator"} est_kwargs = dict( hf_token = hf_token or None, diff --git a/studio/backend/tests/test_chat_load_during_training.py b/studio/backend/tests/test_chat_load_during_training.py index 089c5b851e..f1d973f004 100644 --- a/studio/backend/tests/test_chat_load_during_training.py +++ b/studio/backend/tests/test_chat_load_during_training.py @@ -326,7 +326,7 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase): class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase): - def test_non_cuda_allows(self): + def test_non_accelerator_allows(self): with patch("utils.hardware.get_device", return_value = DeviceType.MLX): ok, info = tv.can_load_chat_during_training( model_name = "m", @@ -336,7 +336,30 @@ class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase): requested_gpu_ids = None, ) self.assertTrue(ok) - self.assertEqual(info["mode"], "non_cuda") + self.assertEqual(info["mode"], "non_accelerator") + + def test_xpu_overcommit_is_refused(self): + # XPU must NOT get the blanket non-accelerator allow: an oversized + # chat model during resident training is refused, like CUDA. + with ( + patch("utils.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.auto_select_gpu_ids", + return_value = ( + None, + {"selection_mode": "auto", "required_gb": 50.0, "usable_gb": 4.0}, + ), + ), + ): + ok, info = tv.can_load_chat_during_training( + model_name = "m", + hf_token = None, + load_in_4bit = True, + max_seq_length = 0, + requested_gpu_ids = None, + ) + self.assertFalse(ok) + self.assertNotEqual(info.get("mode"), "non_accelerator") def test_no_visible_gpus_refuses(self): # GGUF with an empty device list -> no candidate GPU -> default-deny. diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index 3b44c19e26..3dab7ef368 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -119,7 +119,8 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase): patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), ): with self.assertRaisesRegex( - ValueError, "unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG" + ValueError, + "unsupported when CUDA_VISIBLE_DEVICES uses non-numeric or subdevice", ): resolve_requested_gpu_ids([1]) @@ -866,12 +867,12 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase): class TestRouteErrors(unittest.TestCase): - def test_prepare_gpu_selection_rejects_gpu_ids_on_non_cuda_backend(self): + def test_prepare_gpu_selection_rejects_gpu_ids_on_non_accelerator_backend(self): with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU): with self.assertRaises(ValueError) as exc_info: prepare_gpu_selection([0], model_name = "unsloth/test") - self.assertIn("only supported on CUDA devices", str(exc_info.exception)) + self.assertIn("only supported on CUDA and Intel XPU", str(exc_info.exception)) def test_inference_route_resolves_gguf_gpu_ids(self): # GGUF gpu_ids are now supported: /load routes them through the same @@ -1630,18 +1631,61 @@ class TestAutoSelectWithNoneRequired(_GpuCacheResetMixin, unittest.TestCase): self.assertEqual(metadata["selection_mode"], "fallback_all") -class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase): - def test_auto_select_returns_non_cuda_for_xpu(self): - with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): +class TestXpuSelection(_GpuCacheResetMixin, unittest.TestCase): + def test_auto_select_supports_xpu(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware.estimate_required_model_memory_gb", + return_value = (1.0, {}), + ), + patch( + "utils.hardware.hardware.get_visible_gpu_utilization", + return_value = { + "devices": [ + {"index": 0, "vram_total_gb": 8, "vram_used_gb": 1}, + ] + }, + ), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = { + "raw": None, + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch( + "utils.hardware.hardware.get_parent_visible_gpu_ids", + return_value = [0], + ), + ): selected, metadata = auto_select_gpu_ids("unsloth/test") - self.assertIsNone(selected) - self.assertEqual(metadata["selection_mode"], "non_cuda") + self.assertEqual(selected, [0]) + self.assertEqual(metadata["selection_mode"], "auto") - def test_prepare_gpu_selection_rejects_explicit_ids_on_xpu(self): - with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): - with self.assertRaisesRegex(ValueError, "only supported on CUDA"): - prepare_gpu_selection([0], model_name = "unsloth/test") + def test_prepare_gpu_selection_accepts_explicit_ids_on_xpu(self): + with ( + patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU), + patch( + "utils.hardware.hardware._get_parent_visible_gpu_spec", + return_value = { + "raw": "0", + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch( + "utils.hardware.hardware.get_parent_visible_gpu_ids", + return_value = [0], + ), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 1), + ): + selected, metadata = prepare_gpu_selection([0], model_name = "unsloth/test") + + self.assertEqual(selected, [0]) + self.assertEqual(metadata["selection_mode"], "explicit") class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase): diff --git a/studio/backend/tests/test_gpu_selection_sandbox.py b/studio/backend/tests/test_gpu_selection_sandbox.py index 733933271b..ba6d057123 100644 --- a/studio/backend/tests/test_gpu_selection_sandbox.py +++ b/studio/backend/tests/test_gpu_selection_sandbox.py @@ -294,13 +294,13 @@ class TestAutoSelectGpuIds(unittest.TestCase): # 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB self.assertEqual(len(selected), 2) - def test_non_cuda_returns_none(self): + def test_non_accelerator_returns_none(self): from utils.hardware.hardware import auto_select_gpu_ids import utils.hardware.hardware as hw with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): selected, meta = auto_select_gpu_ids("test/model") self.assertIsNone(selected) - self.assertEqual(meta["selection_mode"], "non_cuda") + self.assertEqual(meta["selection_mode"], "non_accelerator") class TestGetDeviceMap(unittest.TestCase): diff --git a/studio/backend/tests/test_training_vram_coexistence.py b/studio/backend/tests/test_training_vram_coexistence.py index 6683cb9aaa..217caaa4fb 100644 --- a/studio/backend/tests/test_training_vram_coexistence.py +++ b/studio/backend/tests/test_training_vram_coexistence.py @@ -326,12 +326,21 @@ class TestCanKeepAuto(_GpuCacheResetMixin, unittest.TestCase): keep, _, _ = self._run((None, meta)) self.assertFalse(keep) - def test_unload_on_non_cuda(self): + def test_unload_on_non_accelerator(self): keep, info, auto_mock = self._run(([0], {}), device = DeviceType.CPU) self.assertFalse(keep) - self.assertEqual(info["mode"], "non_cuda") + self.assertEqual(info["mode"], "non_accelerator") auto_mock.assert_not_called() + def test_xpu_gets_sized_like_cuda(self): + # XPU is a first-class training backend: the keep-guard must size it, + # not blanket-unload it as a non-accelerator. + meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0} + keep, info, auto_mock = self._run(([0], meta), device = DeviceType.XPU) + self.assertTrue(keep) + self.assertNotEqual(info.get("mode"), "non_accelerator") + auto_mock.assert_called_once() + def test_full_finetuning_forces_16bit_in_estimate(self): meta = {"selection_mode": "auto", "required_gb": 10.0, "usable_gb": 30.0} _keep, _info, auto_mock = self._run( diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index 62b537fbac..138238533f 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -50,6 +50,11 @@ def export_capability() -> dict: return _hardware.export_capability() +def get_torch_device_str() -> str: + """Return the torch device string ("cuda", "xpu", "cpu") for the detected hardware.""" + return _hardware.get_torch_device_str() + + __all__ = [ "DeviceType", "DEVICE", @@ -75,6 +80,7 @@ __all__ = [ "estimate_required_model_memory_gb", "auto_select_gpu_ids", "prepare_gpu_selection", + "get_torch_device_str", "safe_num_proc", "safe_thread_num_proc", "dataset_map_num_proc", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index d9a06fb017..38ebc0b6d4 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -175,18 +175,64 @@ def detect_hardware() -> DeviceType: Call once at FastAPI lifespan startup; idempotent. Detection order: - 1. CUDA (NVIDIA GPU, requires torch) - 2. MLX (Apple Silicon via MLX framework) - 3. CPU (fallback) + 1. XPU-preferred hint: only on an unambiguous "prefer XPU" signal + (CUDA hidden via ``CUDA_VISIBLE_DEVICES="" / "-1"``, + ``UNSLOTH_FORCE_XPU=1``, or CUDA unavailable) AND a non-empty + ``ZE_AFFINITY_MASK`` AND ``torch.xpu`` reports a device. A stray + inherited mask is not enough: CUDA still wins on hybrid hosts. + 2. CUDA (NVIDIA GPU, requires torch) + 3. XPU (Intel GPU, requires torch with XPU support) + 4. MLX (Apple Silicon via MLX framework) + 5. CPU (fallback) """ global DEVICE, CHAT_ONLY, CHAT_ONLY_REASON, IS_ROCM CHAT_ONLY = True # reset -- only CUDA/ROCm/XPU/MLX sets it to False CHAT_ONLY_REASON = None IS_ROCM = False - # --- CUDA / ROCm: try PyTorch --- + # --- CUDA / ROCm / XPU: try PyTorch --- if _has_torch(): import torch + + # --- Explicit-XPU hint --- + # Prefer XPU on UNSLOTH_FORCE_XPU=1, or ZE_AFFINITY_MASK set + CUDA + # hidden/unavailable. A bare mask alone is NOT enough (can leak from + # unrelated Intel tooling); torch.xpu must report a device. + ze_mask = os.environ.get("ZE_AFFINITY_MASK") + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + cuda_hidden = cvd is not None and cvd.strip() in ("", "-1") + force_xpu = os.environ.get("UNSLOTH_FORCE_XPU") == "1" + try: + cuda_unavailable = not torch.cuda.is_available() + except Exception: + cuda_unavailable = True + + prefer_xpu = force_xpu or (bool(ze_mask) and (cuda_hidden or cuda_unavailable)) + if prefer_xpu: + try: + xpu_ok = hasattr(torch, "xpu") and torch.xpu.is_available() + except Exception: + xpu_ok = False + if xpu_ok: + # Forced XPU on a hybrid host: unsloth's device_type picks + # CUDA before XPU and ignores this Studio-only env var, so + # hide CUDA or spawned workers would silently train on CUDA. + if force_xpu and not cuda_hidden and not cuda_unavailable: + os.environ["CUDA_VISIBLE_DEVICES"] = "" + DEVICE = DeviceType.XPU + CHAT_ONLY = False + CHAT_ONLY_REASON = None + device_name = torch.xpu.get_device_name(0) + if force_xpu and not ze_mask: + reason = "UNSLOTH_FORCE_XPU=1" + elif force_xpu: + reason = "UNSLOTH_FORCE_XPU=1 + ZE_AFFINITY_MASK" + else: + reason = "ZE_AFFINITY_MASK hint honoured" + print(f"Hardware detected: XPU -- {device_name} ({reason})") + return DEVICE + + # --- CUDA: NVIDIA GPU --- if torch.cuda.is_available(): DEVICE = DeviceType.CUDA CHAT_ONLY = False @@ -327,9 +373,18 @@ def clear_gpu_cache(): torch.cuda.empty_cache() torch.cuda.ipc_collect() elif device == DeviceType.XPU: - import torch - torch.xpu.synchronize() - torch.xpu.empty_cache() + # Guard synchronize/empty_cache: older torch-xpu builds may lack + # them, and an unguarded AttributeError would propagate to callers. + # torch.xpu has no ipc_collect(), so do not call it here. + try: + import torch + if hasattr(torch, "xpu"): + if hasattr(torch.xpu, "synchronize"): + torch.xpu.synchronize() + if hasattr(torch.xpu, "empty_cache"): + torch.xpu.empty_cache() + except Exception as e: + logger.debug("Failed to clear XPU cache: %s", e) elif device == DeviceType.MLX: # MLX manages memory automatically; gc.collect() above is enough. pass @@ -500,14 +555,27 @@ def get_package_versions() -> Dict[str, Optional[str]]: except PackageNotFoundError: versions[name] = None - # GPU runtime version bundled with torch + # GPU runtime versions bundled with torch (CUDA, ROCm/HIP, Intel XPU) try: import torch + versions["cuda"] = getattr(torch.version, "cuda", None) versions["rocm"] = getattr(torch.version, "hip", None) + # Isolated probe: a broken Intel runtime raising in is_available() + # must not blank the already-read cuda/rocm versions. + try: + if hasattr(torch, "xpu") and torch.xpu.is_available(): + # torch.version.xpu may be None on modern builds; fall back to + # "available" so the UI distinguishes present-but-unknown from + # "package not found". + xpu_ver = getattr(torch.version, "xpu", None) + versions["xpu"] = xpu_ver if xpu_ver is not None else "available" + except Exception: + versions["xpu"] = None except Exception: versions["cuda"] = None versions["rocm"] = None + versions["xpu"] = None return versions @@ -547,6 +615,7 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] if mod is None: return [] + device = get_device() # free==total is a Windows-ROCm-only quirk. _win_rocm = sys.platform == "win32" and IS_ROCM devices = [] @@ -558,11 +627,30 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] used_bytes: Optional[int] # Prefer mem_get_info (system-wide) so auto-select sees other consumers. if hasattr(mod, "mem_get_info"): - free_bytes, total_bytes = mod.mem_get_info(ordinal) - used_bytes = total_bytes - free_bytes - # free==total is the broken-API sentinel, not an idle GPU. - if _win_rocm and free_bytes == total_bytes: + try: + free_bytes, total_bytes = mod.mem_get_info(ordinal) + used_bytes = total_bytes - free_bytes + except Exception as e: + if device != DeviceType.XPU: + raise + # Arc B580 and Lunar Lake can report properties while + # rejecting free-memory queries. Preserve the usable + # device and its total memory with unknown utilization. + logger.debug( + "XPU free-memory query failed for ordinal %d: %s", + ordinal, + e, + ) used_bytes = None + else: + # free==total is the broken-API sentinel, not an idle GPU. + if _win_rocm and free_bytes == total_bytes: + used_bytes = None + elif device == DeviceType.XPU: + # XPU without mem_get_info: memory_allocated() is process-local + # and misleading for placement, so return None for the + # selector's no-telemetry fallback. + used_bytes = None else: used_bytes = mod.memory_allocated(ordinal) devices.append( @@ -571,7 +659,9 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] "visible_ordinal": ordinal, "name": props.name, "total_gb": round(total_bytes / (1024**3), 2), - "used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None, + "used_gb": ( + round(used_bytes / (1024**3), 2) if used_bytes is not None else None + ), } ) except Exception as e: @@ -582,6 +672,43 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any] # ========== Live GPU Utilization ========== +def _xpu_hierarchy_is_composite() -> bool: + """Return True iff Level Zero is running in COMPOSITE device hierarchy. + + COMPOSITE: numeric ``ZE_AFFINITY_MASK`` entries address root GPU IDs + (tiles use ``N.M``). FLAT (the oneAPI default; also assumed when + ``ZE_FLAT_DEVICE_HIERARCHY`` is unset): entries address tile/device + handles, so mapping them back to root GPU IDs is unsafe. Only COMPOSITE + gives stable root-ID semantics. + """ + hierarchy = (os.environ.get("ZE_FLAT_DEVICE_HIERARCHY") or "FLAT").strip().upper() + return hierarchy == "COMPOSITE" + + +def _parse_ze_mask_roots(mask: str) -> list[int]: + """Parse a ``ZE_AFFINITY_MASK`` value into an ordered list of root device IDs. + + One root ID per mask token, preserving order and duplicates so logical + ordinals map 1-to-1 to physical root IDs (e.g. ``"0.0,0.1"`` -> ``[0, 0]``, + ``"2.0,0.1,0.2"`` -> ``[2, 0, 0]``); empty list if no parseable digits. + Only meaningful in COMPOSITE hierarchy -- callers needing a stable + root-ID mapping must gate on ``_xpu_hierarchy_is_composite()``. + """ + roots: list[int] = [] + if not mask: + return roots + for token in mask.split(","): + token = token.strip() + if not token: + continue + root = token.split(".", 1)[0] + # isdecimal() (not isdigit()) rejects Unicode superscripts like + # "²"/"³", which pass isdigit() but crash int() with ValueError. + if root.isdecimal(): + roots.append(int(root)) + return roots + + def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]: """Query the appropriate SMI backend (amd-smi or nvidia-smi). @@ -1504,6 +1631,13 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: for td in torch_devices: total = td["total_gb"] used = td["used_gb"] + # used=None is a deliberate "telemetry unavailable" signal + # from _torch_get_per_device_info (e.g. XPU without + # mem_get_info); propagate None instead of dividing by it. On + # CUDA/ROCm used is always an int, so this stays byte-identical. + vram_pct = ( + round((used / total) * 100, 1) if used is not None and total > 0 else None + ) devices.append( { "index": td["index"], @@ -1513,9 +1647,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]: "temperature_c": None, "vram_used_gb": used, "vram_total_gb": total, - "vram_utilization_pct": round((used / total) * 100, 1) - if total > 0 and used is not None - else None, + "vram_utilization_pct": vram_pct, "power_draw_w": None, "power_limit_w": None, "power_utilization_pct": None, @@ -1583,6 +1715,82 @@ _visible_gpu_count: Optional[int] = None def _get_parent_visible_gpu_spec() -> Dict[str, Any]: + # On Intel XPU, visibility is controlled by ZE_AFFINITY_MASK (Level Zero), + # not CUDA_VISIBLE_DEVICES. + if get_device() == DeviceType.XPU: + xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK") + composite = _xpu_hierarchy_is_composite() + + if xpu_mask_raw is None: + # COMPOSITE: root GPU IDs are stable physical IDs. + if composite: + return { + "raw": None, + "numeric_ids": list(range(get_physical_gpu_count())), + "supports_explicit_gpu_ids": True, + } + # FLAT (oneAPI default): ordinals are tile/device handles, not + # physical GPU IDs. numeric_ids=None so telemetry uses relative + # ordinals; explicit selection needs ZE_FLAT_DEVICE_HIERARCHY=COMPOSITE. + return { + "raw": None, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + xpu_mask = xpu_mask_raw.strip() + if xpu_mask == "": + return { + "raw": xpu_mask, + "numeric_ids": [], + "supports_explicit_gpu_ids": True, + } + + # Subdevice syntax ("N.M") expands one root into multiple + # logical devices -- not addressable by explicit root-ID selection. + has_subdevice = any("." in token.strip() for token in xpu_mask.split(",") if token.strip()) + if has_subdevice: + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + # FLAT numeric entries are tile handles, not physical GPU IDs. Keep + # numeric_ids unresolved so every telemetry and picker consumer uses + # relative torch ordinals and cannot advertise them as pinnable roots. + if not composite: + tokens = [token.strip() for token in xpu_mask.split(",") if token.strip()] + if tokens and all(token.isdecimal() for token in tokens): + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + # COMPOSITE + pure numeric (subdevice handled above). _parse_ze_mask_roots + # maps to root GPU IDs, dropping non-decimal tokens so "*"/"GPU-uuid" -> []. + roots_with_dupes = _parse_ze_mask_roots(xpu_mask) + if not roots_with_dupes: + # Unparseable mask (e.g. "*", "GPU-uuid") -- cannot map to + # physical root IDs. + return { + "raw": xpu_mask, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + return { + "raw": xpu_mask, + "numeric_ids": roots_with_dupes, + "supports_explicit_gpu_ids": True, + } + # ROCm uses HIP/ROCR_VISIBLE_DEVICES on top of CUDA_VISIBLE_DEVICES; check # them first. Explicit None checks (not `or`) so "" reads as "no visible GPUs". cuda_visible = None @@ -1669,11 +1877,14 @@ def resolve_requested_gpu_ids( return requested_ids if not parent_visible_spec["supports_explicit_gpu_ids"]: + env_var_name = ( + "ZE_AFFINITY_MASK" if get_device() == DeviceType.XPU else "CUDA_VISIBLE_DEVICES" + ) raise ValueError( f"Invalid gpu_ids {requested_ids}: explicit physical GPU IDs are " - f"unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries " - f"({parent_visible_spec['raw']!r}). Omit gpu_ids to use the " - "parent-visible devices." + f"unsupported when {env_var_name} uses non-numeric or subdevice " + f"entries ({parent_visible_spec['raw']!r}). Omit gpu_ids to use " + "the parent-visible devices." ) if len(set(requested_ids)) != len(requested_ids): @@ -2112,8 +2323,11 @@ def auto_select_gpu_ids( ) -> tuple[Optional[list[int]], Dict[str, Any]]: metadata: Dict[str, Any] = {"selection_mode": "auto"} - if get_device() != DeviceType.CUDA: - metadata["selection_mode"] = "non_cuda" + # Auto-selection needs per-device free-VRAM telemetry, available on CUDA + # (nvidia-smi) and XPU (torch.xpu) but not MLX/CPU, which fall + # through to inheriting parent visibility. + if get_device() not in (DeviceType.CUDA, DeviceType.XPU): + metadata["selection_mode"] = "non_accelerator" return None, metadata required_gb, estimate_metadata = estimate_required_model_memory_gb( @@ -2272,10 +2486,10 @@ def prepare_gpu_selection( to a Hugging Face ``device_map`` string) and to ``apply_gpu_ids()`` in the worker subprocess (narrows ``CUDA_VISIBLE_DEVICES`` before torch/CUDA init). """ - if gpu_ids and get_device() != DeviceType.CUDA: + if gpu_ids and get_device() not in (DeviceType.CUDA, DeviceType.XPU): raise ValueError( - f"gpu_ids {list(gpu_ids)} is only supported on CUDA devices, " - f"but the current backend is '{get_device().value}'." + f"gpu_ids {list(gpu_ids)} is only supported on CUDA and Intel XPU " + f"devices, but the current backend is '{get_device().value}'." ) if gpu_ids: @@ -2348,11 +2562,14 @@ def get_physical_gpu_count() -> int: def _backend_visible_devices_env() -> Optional[str]: """Return the raw visibility env string that applies to this backend. - On ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over - CUDA_VISIBLE_DEVICES; this mirrors ``_get_parent_visible_gpu_spec`` so + On XPU the control is ``ZE_AFFINITY_MASK`` (not ``CUDA_VISIBLE_DEVICES``); + on ROCm, HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES take precedence over + CUDA_VISIBLE_DEVICES. Mirrors ``_get_parent_visible_gpu_spec`` so ``backend_cuda_visible_devices`` reports the value actually narrowing the - visible device set. + visible device set on the current backend. """ + if get_device() == DeviceType.XPU: + return os.environ.get("ZE_AFFINITY_MASK") if IS_ROCM: return _get_parent_visible_gpu_spec().get("raw") return os.environ.get("CUDA_VISIBLE_DEVICES") @@ -2467,6 +2684,43 @@ def get_visible_gpu_count() -> int: if _visible_gpu_count is not None: return _visible_gpu_count + # Prefer torch.xpu.device_count() on Intel XPU: the Level Zero runtime + # correctly interprets ZE_AFFINITY_MASK semantics (e.g. subdevice syntax + # "0.0,0.1" collapses onto one root GPU). Supersedes the torch fallback below. + if get_device() == DeviceType.XPU: + xpu_mask_raw = os.environ.get("ZE_AFFINITY_MASK") + xpu_mask_set = xpu_mask_raw is not None + xpu_visible = (xpu_mask_raw or "").strip() + if xpu_mask_set and xpu_visible == "": + _visible_gpu_count = 0 + return _visible_gpu_count + + try: + import torch + _visible_gpu_count = torch.xpu.device_count() + except Exception as e: + logger.debug( + "torch.xpu.device_count() failed, falling back to mask parsing: %s", + e, + ) + if xpu_visible: + # Fallback: count unique root device IDs from the mask. + # "device.subdevice" notation means "0.0,0.1" is 1 root, not 2. + # Without torch the hierarchy mode is unknown, so root-device + # counting is the conservative choice. + if xpu_visible == "*": + # Documented wildcard: all physical XPUs visible. + _visible_gpu_count = get_physical_gpu_count() + else: + roots = _parse_ze_mask_roots(xpu_visible) + # Non-parseable masks (",,,", "GPU-abc") yield an empty + # roots list, treated as 0 visible devices, not "all + # visible" -- no evidence the whole fleet was intended. + _visible_gpu_count = len(set(roots)) + else: + _visible_gpu_count = get_physical_gpu_count() + return _visible_gpu_count + # _get_parent_visible_gpu_spec() already handles HIP_VISIBLE_DEVICES / # ROCR_VISIBLE_DEVICES on ROCm. visible_spec = _get_parent_visible_gpu_spec() @@ -2480,20 +2734,18 @@ def get_visible_gpu_count() -> int: _visible_gpu_count = len([x for x in raw.split(",") if x.strip()]) return _visible_gpu_count - # No visibility env var set -- try torch, else physical count + # No visibility env var set -- try torch, else physical count. XPU is + # handled by the early return above, so only torch.cuda is needed here. try: import torch - if get_device() == DeviceType.XPU and hasattr(torch, "xpu"): - _visible_gpu_count = torch.xpu.device_count() - else: - _visible_gpu_count = torch.cuda.device_count() + _visible_gpu_count = torch.cuda.device_count() except Exception: _visible_gpu_count = get_physical_gpu_count() return _visible_gpu_count -def apply_gpu_ids(gpu_ids) -> None: +def apply_gpu_ids(gpu_ids, backend: Optional[str] = None) -> None: if gpu_ids is None: return @@ -2509,6 +2761,62 @@ def apply_gpu_ids(gpu_ids) -> None: else: value = str(gpu_ids) + # Intel XPU honors ZE_AFFINITY_MASK, not CUDA_VISIBLE_DEVICES; route XPU + # pinning through it so worker subprocesses are restricted to the intended GPU. + # Decide WITHOUT get_device(): workers call this before detect_hardware(), + # and a lazy detect would probe torch.cuda against the unmasked parent env, + # latching device enumeration before the mask below is written. Pre-detect, + # use env + torch BUILD attributes only (no runtime init, like the ROCm + # mirror below). + _is_xpu = DEVICE == DeviceType.XPU + if backend is not None: + # The spawning parent's detected backend (config["device_backend"]): + # exact and probe-free, so the mask target always matches what + # detect_hardware() decided in the parent, including its XPU + # availability check and CUDA fallback. + _is_xpu = backend == DeviceType.XPU.value + elif DEVICE is None: + # No parent backend passed (direct caller). version.xpu can be None + # on a working XPU build, so also accept torch.xpu._is_compiled() + # (a pure symbol-presence check, no runtime init). UNSLOTH_FORCE_XPU + # counts only on an XPU-capable build: detect_hardware() falls back + # to CUDA when XPU is missing, and the mask target must follow. + try: + import torch as _torch + + _ver = _torch.version + _is_comp = getattr(getattr(_torch, "xpu", None), "_is_compiled", None) + _xpu_build = (callable(_is_comp) and bool(_is_comp())) or ( + getattr(_ver, "xpu", None) is not None + ) + if os.environ.get("UNSLOTH_FORCE_XPU") == "1": + _is_xpu = _xpu_build + else: + # Mirror detect_hardware: hidden CUDA prefers XPU on an + # XPU-capable build (with or without a ZE mask -- detection + # falls through to XPU either way), where writing these ids + # to CUDA_VISIBLE_DEVICES would re-expose the deliberately + # hidden CUDA. + _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + _cuda_hidden = _cvd is not None and _cvd.strip() in ("", "-1") + _is_xpu = _xpu_build and ( + _cuda_hidden + or (getattr(_ver, "cuda", None) is None and getattr(_ver, "hip", None) is None) + ) + except Exception as e: + logger.debug( + "apply_gpu_ids: torch XPU probe skipped (%s: %s)", + type(e).__name__, + e, + ) + if _is_xpu: + os.environ["ZE_AFFINITY_MASK"] = value + # Leave inherited CUDA_VISIBLE_DEVICES alone -- clearing it could let + # the worker flip back to CUDA on hybrid hosts. + _visible_gpu_count = None + logger.info("Applied gpu_ids: ZE_AFFINITY_MASK='%s'", value) + return + os.environ["CUDA_VISIBLE_DEVICES"] = value # Keep ROCm visibility env vars in sync. Workers may call apply_gpu_ids() # before detect_hardware() (IS_ROCM still False), so also mirror when the @@ -2553,26 +2861,41 @@ def get_device_map(gpu_ids: Optional[list[int]] = None) -> str: Returns ``"balanced"`` (shard evenly across GPUs) when: - ``gpu_ids`` explicitly lists >1 GPU, **or** - - ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and - >1 GPU is visible (fallback: numeric IDs unresolvable, so assume - multi-GPU is intended). + - ``CUDA_VISIBLE_DEVICES``/``ZE_AFFINITY_MASK`` uses non-numeric + identifiers (UUID/MIG/wildcard) and >1 GPU is visible (fallback: + numeric IDs unresolvable, so assume multi-GPU is intended). - Returns ``"sequential"`` (single device) otherwise, including non-CUDA - backends (CPU, MLX). + Returns ``"sequential"`` (single device) otherwise, including CPU/MLX + backends. Use ``prepare_gpu_selection()`` upstream to determine ``gpu_ids`` -- it handles auto-selecting the minimum GPUs needed for a model. """ device = get_device() - if device == DeviceType.CUDA: + if device in (DeviceType.CUDA, DeviceType.XPU): multi_gpu = gpu_ids is not None and len(gpu_ids) > 1 if not multi_gpu: - # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU - # means multi-GPU sharding is intended. parent_visible_spec = _get_parent_visible_gpu_spec() - if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: - multi_gpu = True + if device == DeviceType.CUDA: + # UUID/MIG masks can't be split into numeric IDs; >1 visible GPU + # means multi-GPU sharding is intended. + if parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1: + multi_gpu = True + elif device == DeviceType.XPU and gpu_ids is None: + # Shard across visible XPU ordinals via HF (no mask rewrite), + # only when no gpu_ids were passed -- an explicit gpu_ids=[0] + # means "use exactly device 0" and must stay sequential. + supports_physical = parent_visible_spec["supports_explicit_gpu_ids"] + has_multiple_numeric = ( + parent_visible_spec["numeric_ids"] is not None + and len(parent_visible_spec["numeric_ids"]) > 1 + ) + has_multiple_unresolved = ( + parent_visible_spec["numeric_ids"] is None and get_visible_gpu_count() > 1 + ) + if has_multiple_unresolved or (not supports_physical and has_multiple_numeric): + multi_gpu = True if multi_gpu: return "balanced" @@ -2607,6 +2930,19 @@ def raise_if_offloaded( ) +def get_torch_device_str() -> str: + """ + Return the torch device string for the detected hardware. + E.g. "cuda", "xpu", or "cpu". + """ + device = get_device() + if device == DeviceType.CUDA: + return "cuda" + elif device == DeviceType.XPU: + return "xpu" + return "cpu" + + def safe_num_proc(desired: Optional[int] = None) -> int: """ Return a safe ``num_proc`` for ``dataset.map()`` calls. @@ -2674,7 +3010,32 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]: Returns ``None`` on spawn platforms (Windows, macOS) because ``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``); only ``num_proc=None`` guarantees in-process execution. + + Also returns ``None`` on XPU once its runtime is initialized in this + process: ``os.fork()`` corrupts the Level-Zero context, making Triton + kernels fail with "Pointer argument doesn't reference XPU device memory". + Pre-init XPU hosts can still parallelize CPU-side preprocessing. """ if sys.platform in ("win32", "darwin"): return None + + if get_device() == DeviceType.XPU: + try: + import torch + except Exception: + # No torch means no active XPU runtime, so CPU-side dataset + # parallelism is still safe. + return safe_num_proc(desired) + + xpu = getattr(torch, "xpu", None) + is_initialized = getattr(xpu, "is_initialized", None) + if callable(is_initialized): + try: + if is_initialized(): + return None + except Exception as e: + # Treat a failing probe as "runtime not touched yet" so + # pre-init CPU preprocessing can still parallelize. + logger.debug("torch.xpu.is_initialized() probe failed: %s", e) + return safe_num_proc(desired) diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index bf8348dd82..ce3d6704b7 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -294,17 +294,31 @@ def format_error_message(error: Exception, model_name: str) -> str: return "Invalid HF token. Please check your token and try again." if ( - "memory" in error_str - or "cuda" in error_str - or "mlx" in error_str - or "out of memory" in error_str + "out of memory" in error_str + or "out of device memory" in error_str + or "out_of_device_memory" in error_str # ZE_RESULT_ERROR_OUT_OF_DEVICE_MEMORY + or "out_of_host_memory" in error_str # ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY + or "not enough memory" in error_str + or "cannot allocate memory" in error_str + or "memory allocation failed" in error_str + or "cublas_status_alloc_failed" in error_str # cuBLAS workspace OOM + or ("cuda error" in error_str and "alloc" in error_str) + or ("xpu" in error_str and ("alloc" in error_str or "memory" in error_str)) + or isinstance(error, MemoryError) + or ("mlx" in error_str and ("memory" in error_str or "allocate" in error_str)) ): + # Resolve get_device() at call time (not import time) so tests that + # monkey-patch utils.hardware.get_device after this module is loaded + # still see the patched backend. from utils.hardware import get_device device = get_device() - device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get( - device.value, "GPU" - ) + device_label = { + "cuda": "GPU", + "xpu": "Intel GPU", + "mlx": "Apple Silicon GPU", + "cpu": "system", + }.get(device.value, "GPU") return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory." return str(error) diff --git a/tests/studio/test_xpu_spoof_pipeline.py b/tests/studio/test_xpu_spoof_pipeline.py new file mode 100644 index 0000000000..4a458a545d --- /dev/null +++ b/tests/studio/test_xpu_spoof_pipeline.py @@ -0,0 +1,538 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""Full Intel XPU spoof pipeline: fake torch.xpu on a GPU-less/NVIDIA runner so +Studio's hardware selection + training-device path (detect -> select -> apply -> +device_map -> cache clear) runs exactly as the CUDA path does, with no real +Intel hardware. The XPU sibling of tests/_zoo_aggressive_cuda_spoof.py. + +State-sensitive: it fresh-imports the Studio hardware module under the spoof and +mutates its module globals, so studio-backend-ci.yml runs it in the isolated +"Hardware-spoof tests" step (never alongside tests that import hardware). + +torch.xpu surface faked here mirrors the PyTorch 2.6+ API hardware.py calls: +is_available, device_count, current_device, get_device_name, +get_device_properties(idx).total_memory, memory_allocated/reserved, mem_get_info +(incl. the Arc B580 / Lunar Lake RuntimeError), is_initialized, synchronize, +empty_cache, plus torch.version.xpu. +""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_BACKEND = REPO_ROOT / "studio" / "backend" + + +def _make_fake_xpu( + *, + available: bool = True, + device_count: int = 2, + total_gb: float = 16.0, + used_gb: float = 1.0, + device_name: str = "Intel(R) Arc(TM) B580 Graphics (spoofed)", + mem_get_info: str = "ok", # "ok" | "raise" | "absent" + is_initialized: bool = False, +): + """Build a fake torch.xpu namespace + a call counter for synchronize/empty_cache. + + mem_get_info: "ok" returns (free, total); "raise" models the Arc B580 / Lunar + Lake "device doesn't support querying free memory" RuntimeError; "absent" + omits the attribute so the memory_allocated fallback path is exercised. + """ + total_bytes = int(total_gb * 1024**3) + used_bytes = int(used_gb * 1024**3) + calls = {"synchronize": 0, "empty_cache": 0} + props = types.SimpleNamespace(name = device_name, total_memory = total_bytes) + + def _mem_get_info(idx = 0): + if mem_get_info == "raise": + raise RuntimeError( + "The device (Intel(R) Arc(TM) B580 Graphics) doesn't support " + "querying the available free memory." + ) + return (total_bytes - used_bytes, total_bytes) + + def _sync(*a, **k): + calls["synchronize"] += 1 + + def _empty(*a, **k): + calls["empty_cache"] += 1 + + xpu = types.SimpleNamespace( + is_available = lambda: available, + device_count = lambda: device_count, + current_device = lambda: 0, + get_device_name = lambda idx = 0: device_name, + get_device_properties = lambda idx = 0: props, + memory_allocated = lambda idx = 0: used_bytes, + memory_reserved = lambda idx = 0: used_bytes, + is_initialized = lambda: is_initialized, + synchronize = _sync, + empty_cache = _empty, + ) + if mem_get_info != "absent": + xpu.mem_get_info = _mem_get_info + return xpu, calls + + +def _import_studio_hardware_module(): + """Fresh-import Studio's hardware module so detect_hardware re-runs under the + current spoofs (mirrors test_hardware_dispatch_matrix.py).""" + if str(STUDIO_BACKEND) not in sys.path: + sys.path.insert(0, str(STUDIO_BACKEND)) + sys.modules.pop("utils.hardware.hardware", None) + sys.modules.pop("utils.hardware", None) + from utils.hardware import hardware as hw # type: ignore + + return hw + + +@pytest.fixture +def spoof_xpu(monkeypatch): + """Apply a full torch.xpu spoof and return (hardware_module, xpu_call_counter). + + Defaults present an unambiguous "prefer XPU" host: CUDA hidden, a numeric + ZE_AFFINITY_MASK, and torch.xpu reporting devices. Override cuda_available / + cuda_visible / force_xpu / ze_mask to model hybrid or canary hosts. + """ + + def _apply( + *, + cuda_available: bool = False, + cuda_visible: str = "", # "" hides CUDA; None unsets; else passthrough + ze_mask: str = "0,1", # None unsets the mask + force_xpu: bool = False, + xpu_version = "2.7", + **xpu_kwargs, + ): + import torch + + monkeypatch.setattr(torch.cuda, "is_available", lambda: cuda_available) + if cuda_available: + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda i = 0: types.SimpleNamespace(name = "Stub NVIDIA GPU"), + raising = False, + ) + fake_xpu, calls = _make_fake_xpu(**xpu_kwargs) + monkeypatch.setattr(torch, "xpu", fake_xpu, raising = False) + monkeypatch.setattr(torch.version, "xpu", xpu_version, raising = False) + + if ze_mask is None: + monkeypatch.delenv("ZE_AFFINITY_MASK", raising = False) + else: + monkeypatch.setenv("ZE_AFFINITY_MASK", ze_mask) + if cuda_visible is None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + else: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda_visible) + if force_xpu: + monkeypatch.setenv("UNSLOTH_FORCE_XPU", "1") + else: + monkeypatch.delenv("UNSLOTH_FORCE_XPU", raising = False) + # FLAT is the oneAPI default; pin it so the test is host-independent. + monkeypatch.delenv("ZE_FLAT_DEVICE_HIERARCHY", raising = False) + + hw = _import_studio_hardware_module() + hw._visible_gpu_count = None + return hw, calls + + return _apply + + +# ---------- detection ---------- + + +def test_detect_hardware_routes_to_xpu(spoof_xpu): + hw, _ = spoof_xpu() + assert hw.detect_hardware() == hw.DeviceType.XPU + assert hw.CHAT_ONLY is False + assert hw.IS_ROCM is False + + +def test_force_xpu_env_routes_to_xpu_even_without_mask(spoof_xpu): + hw, _ = spoof_xpu(force_xpu = True, ze_mask = None, cuda_visible = None) + assert hw.detect_hardware() == hw.DeviceType.XPU + + +def test_bare_mask_with_cuda_present_stays_cuda(spoof_xpu): + # Canary: a stray inherited ZE_AFFINITY_MASK must NOT steal a CUDA host. + hw, _ = spoof_xpu(cuda_available = True, cuda_visible = None, ze_mask = "0,1") + assert hw.detect_hardware() == hw.DeviceType.CUDA + + +def test_force_xpu_on_hybrid_hides_cuda_for_workers(spoof_xpu): + # Forced XPU with CUDA still visible must hide CUDA: unsloth's + # device_type picks CUDA before XPU and ignores UNSLOTH_FORCE_XPU, + # so workers would otherwise silently train on CUDA. + hw, _ = spoof_xpu(force_xpu = True, cuda_available = True, cuda_visible = None, ze_mask = None) + assert hw.detect_hardware() == hw.DeviceType.XPU + import os + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "" + + +def test_force_xpu_without_working_xpu_leaves_cuda_untouched(spoof_xpu): + # Canary: FORCE_XPU on a CUDA host with no working XPU must fall + # through to CUDA and must NOT hide it. + hw, _ = spoof_xpu( + force_xpu = True, + cuda_available = True, + cuda_visible = None, + ze_mask = None, + available = False, + ) + assert hw.detect_hardware() == hw.DeviceType.CUDA + import os + + assert "CUDA_VISIBLE_DEVICES" not in os.environ + + +def test_apply_gpu_ids_predetect_never_probes_torch(spoof_xpu, monkeypatch): + # Workers call apply_gpu_ids() BEFORE detect_hardware(); a lazy detect + # would probe torch.cuda against the unmasked parent env, latching device + # enumeration before the mask is written. Pre-detect it must decide from + # env/build attributes only. + import torch + + hw, _ = spoof_xpu(ze_mask = None, cuda_visible = None) + assert hw.DEVICE is None # fresh import, pre-detect + + def _poisoned_detect(): + raise AssertionError("apply_gpu_ids triggered detect_hardware pre-mask") + + monkeypatch.setattr(hw, "detect_hardware", _poisoned_detect) + monkeypatch.setattr(torch.cuda, "is_available", _poisoned_detect, raising = False) + # CUDA-build torch (torch.version.cuda set on this box or spoofed): + monkeypatch.setattr(torch.version, "cuda", "12.8", raising = False) + monkeypatch.setattr(torch.version, "xpu", None, raising = False) + hw.apply_gpu_ids([1]) + import os + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "1" + assert "ZE_AFFINITY_MASK" not in os.environ + + +def test_apply_gpu_ids_predetect_xpu_build_writes_ze_mask(spoof_xpu, monkeypatch): + # Pre-detect on an XPU-build torch (version.xpu set, no cuda/hip): + # the mask must go to ZE_AFFINITY_MASK without any runtime probe. + import torch + + hw, _ = spoof_xpu(ze_mask = None, cuda_visible = None) + assert hw.DEVICE is None + + def _poisoned_detect(): + raise AssertionError("apply_gpu_ids triggered detect_hardware pre-mask") + + monkeypatch.setattr(hw, "detect_hardware", _poisoned_detect) + monkeypatch.setattr(torch.version, "cuda", None, raising = False) + monkeypatch.setattr(torch.version, "hip", None, raising = False) + monkeypatch.setattr(torch.version, "xpu", "2.7", raising = False) + hw.apply_gpu_ids([0]) + import os + + assert os.environ["ZE_AFFINITY_MASK"] == "0" + assert "CUDA_VISIBLE_DEVICES" not in os.environ + + +def test_apply_gpu_ids_predetect_xpu_compiled_with_null_version(spoof_xpu, monkeypatch): + # version.xpu can be None on a working XPU build; torch.xpu._is_compiled() + # must be accepted as the build signal so the mask still goes to + # ZE_AFFINITY_MASK. + import torch + + hw, _ = spoof_xpu(ze_mask = None, cuda_visible = None) + assert hw.DEVICE is None + monkeypatch.setattr( + hw, "detect_hardware", lambda: (_ for _ in ()).throw(AssertionError("detect ran")) + ) + monkeypatch.setattr(torch.version, "cuda", None, raising = False) + monkeypatch.setattr(torch.version, "hip", None, raising = False) + monkeypatch.setattr(torch.version, "xpu", None, raising = False) + monkeypatch.setattr(torch.xpu, "_is_compiled", lambda: True, raising = False) + hw.apply_gpu_ids([0]) + import os + + assert os.environ["ZE_AFFINITY_MASK"] == "0" + assert "CUDA_VISIBLE_DEVICES" not in os.environ + + +def test_apply_gpu_ids_predetect_force_on_cuda_build_writes_cvd(spoof_xpu, monkeypatch): + # UNSLOTH_FORCE_XPU=1 on a CUDA build (no XPU compiled in): detect falls + # back to CUDA, so the pre-detect mask must go to CUDA_VISIBLE_DEVICES, + # not ZE_AFFINITY_MASK. + import torch + + hw, _ = spoof_xpu(force_xpu = True, ze_mask = None, cuda_visible = None) + assert hw.DEVICE is None + monkeypatch.setattr( + hw, "detect_hardware", lambda: (_ for _ in ()).throw(AssertionError("detect ran")) + ) + monkeypatch.setattr(torch.version, "cuda", "12.8", raising = False) + monkeypatch.setattr(torch.version, "xpu", None, raising = False) + monkeypatch.setattr(torch.xpu, "_is_compiled", lambda: False, raising = False) + hw.apply_gpu_ids([1]) + import os + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "1" + assert "ZE_AFFINITY_MASK" not in os.environ + + +def test_apply_gpu_ids_predetect_dual_build_honors_xpu_hint(spoof_xpu, monkeypatch): + # Dual CUDA+XPU build launched the documented XPU way (CUDA hidden + ZE + # mask): the mask must narrow ZE_AFFINITY_MASK, not re-expose the hidden + # CUDA via CUDA_VISIBLE_DEVICES. Mirrors detect_hardware's hint. + import torch + + hw, _ = spoof_xpu(ze_mask = "0,1", cuda_visible = "") + assert hw.DEVICE is None + monkeypatch.setattr( + hw, "detect_hardware", lambda: (_ for _ in ()).throw(AssertionError("detect ran")) + ) + monkeypatch.setattr(torch.version, "cuda", "12.8", raising = False) + monkeypatch.setattr(torch.version, "xpu", "2.7", raising = False) + hw.apply_gpu_ids([0]) + import os + + assert os.environ["ZE_AFFINITY_MASK"] == "0" + assert os.environ["CUDA_VISIBLE_DEVICES"] == "" # stays hidden + + +def test_apply_gpu_ids_predetect_dual_build_cuda_active_writes_cvd(spoof_xpu, monkeypatch): + # Canary: dual build with CUDA active (no hint) keeps CUDA masking, same + # as detect_hardware picking CUDA on a hybrid host. + import torch + + hw, _ = spoof_xpu(ze_mask = "0,1", cuda_visible = None) + assert hw.DEVICE is None + monkeypatch.setattr( + hw, "detect_hardware", lambda: (_ for _ in ()).throw(AssertionError("detect ran")) + ) + monkeypatch.setattr(torch.version, "cuda", "12.8", raising = False) + monkeypatch.setattr(torch.version, "xpu", "2.7", raising = False) + hw.apply_gpu_ids([1]) + import os + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "1" + assert os.environ["ZE_AFFINITY_MASK"] == "0,1" # untouched + + +def test_apply_gpu_ids_trusts_parent_backend_param(spoof_xpu, monkeypatch): + # Workers pass the parent's detected backend (config["device_backend"]): + # it must win over build heuristics in both directions, mirroring + # detect_hardware's availability check and CUDA fallback exactly. + import torch + + hw, _ = spoof_xpu(ze_mask = None, cuda_visible = None, force_xpu = True) + assert hw.DEVICE is None + monkeypatch.setattr( + hw, "detect_hardware", lambda: (_ for _ in ()).throw(AssertionError("detect ran")) + ) + # Forced XPU + XPU build, but the parent detected CUDA (xpu had no + # device): backend="cuda" must route to CUDA_VISIBLE_DEVICES. + monkeypatch.setattr(torch.version, "cuda", "12.8", raising = False) + monkeypatch.setattr(torch.version, "xpu", "2.7", raising = False) + hw.apply_gpu_ids([1], backend = "cuda") + import os + + assert os.environ["CUDA_VISIBLE_DEVICES"] == "1" + assert "ZE_AFFINITY_MASK" not in os.environ + + # And backend="xpu" routes to ZE_AFFINITY_MASK even on a CUDA build. + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False) + monkeypatch.setattr(torch.version, "xpu", None, raising = False) + hw.apply_gpu_ids([0], backend = "xpu") + assert os.environ["ZE_AFFINITY_MASK"] == "0" + assert "CUDA_VISIBLE_DEVICES" not in os.environ + + +def test_apply_gpu_ids_predetect_hidden_cuda_without_mask_prefers_xpu(spoof_xpu, monkeypatch): + # Hidden CUDA on an XPU-capable build prefers XPU even with NO ZE mask + # set (detection falls through to XPU in that state); writing the ids to + # CUDA_VISIBLE_DEVICES would re-expose the hidden CUDA. + import torch + + hw, _ = spoof_xpu(ze_mask = None, cuda_visible = "") + assert hw.DEVICE is None + monkeypatch.setattr( + hw, "detect_hardware", lambda: (_ for _ in ()).throw(AssertionError("detect ran")) + ) + monkeypatch.setattr(torch.version, "cuda", "12.8", raising = False) + monkeypatch.setattr(torch.version, "xpu", "2.7", raising = False) + hw.apply_gpu_ids([0]) + import os + + assert os.environ["ZE_AFFINITY_MASK"] == "0" + assert os.environ["CUDA_VISIBLE_DEVICES"] == "" # stays hidden + + +# ---------- visibility / selection ---------- + + +def test_apply_gpu_ids_writes_ze_affinity_mask(spoof_xpu, monkeypatch): + hw, _ = spoof_xpu() + hw.detect_hardware() + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "sentinel") + hw.apply_gpu_ids([0, 1]) + import os + + assert os.environ["ZE_AFFINITY_MASK"] == "0,1" + # XPU pinning must not touch CUDA_VISIBLE_DEVICES (hybrid-host safety). + assert os.environ["CUDA_VISIBLE_DEVICES"] == "sentinel" + + +def test_get_visible_gpu_count_uses_device_count(spoof_xpu): + hw, _ = spoof_xpu(ze_mask = "0,1", device_count = 2) + hw.detect_hardware() + assert hw.get_visible_gpu_count() == 2 + + +def test_get_visible_gpu_count_empty_mask_is_zero(spoof_xpu): + hw, _ = spoof_xpu(ze_mask = "") + hw.detect_hardware() + assert hw.get_visible_gpu_count() == 0 + + +def test_flat_numeric_mask_reports_relative_ordinals(spoof_xpu): + hw, _ = spoof_xpu(ze_mask = "4,7", device_count = 2) + hw.detect_hardware() + + spec = hw._get_parent_visible_gpu_spec() + assert spec["numeric_ids"] is None + assert spec["supports_explicit_gpu_ids"] is False + + for result in (hw.get_visible_gpu_utilization(), hw.get_backend_visible_gpu_info()): + assert result["available"] is True + assert result["index_kind"] == "relative" + assert result["parent_visible_gpu_ids"] == [] + assert [device["index"] for device in result["devices"]] == [0, 1] + + +def test_composite_numeric_mask_reports_physical_ids(spoof_xpu, monkeypatch): + hw, _ = spoof_xpu(ze_mask = "4,7", device_count = 2) + monkeypatch.setenv("ZE_FLAT_DEVICE_HIERARCHY", "COMPOSITE") + hw.detect_hardware() + + spec = hw._get_parent_visible_gpu_spec() + assert spec["numeric_ids"] == [4, 7] + assert spec["supports_explicit_gpu_ids"] is True + + for result in (hw.get_visible_gpu_utilization(), hw.get_backend_visible_gpu_info()): + assert result["available"] is True + assert result["index_kind"] == "physical" + assert result["parent_visible_gpu_ids"] == [4, 7] + assert [device["index"] for device in result["devices"]] == [4, 7] + + +def test_get_device_map_multi_is_balanced(spoof_xpu): + hw, _ = spoof_xpu(ze_mask = "0,1", device_count = 2) + hw.detect_hardware() + assert hw.get_device_map([0, 1]) == "balanced" + + +def test_get_device_map_explicit_single_is_sequential(spoof_xpu): + hw, _ = spoof_xpu(ze_mask = "0,1", device_count = 2) + hw.detect_hardware() + # Explicit gpu_ids=[0] is a deliberate single-device request. + assert hw.get_device_map([0]) == "sequential" + + +# ---------- cache / telemetry / versions ---------- + + +def test_clear_gpu_cache_calls_xpu(spoof_xpu): + hw, calls = spoof_xpu() + hw.detect_hardware() + hw.clear_gpu_cache() + assert calls["synchronize"] >= 1 + assert calls["empty_cache"] >= 1 + + +def test_package_versions_survive_broken_xpu_runtime(spoof_xpu, monkeypatch): + # A broken Intel runtime raising in is_available() must not blank the + # CUDA/ROCm versions on NVIDIA/AMD hosts. + import torch + + hw, _ = spoof_xpu(cuda_available = True, cuda_visible = None, ze_mask = None) + monkeypatch.setattr(torch.version, "cuda", "12.8", raising = False) + + def _broken(): + raise RuntimeError("Level Zero init failed") + + monkeypatch.setattr(torch.xpu, "is_available", _broken) + versions = hw.get_package_versions() + assert versions["cuda"] == "12.8" + assert versions.get("xpu") is None + + +def test_package_versions_reports_xpu(spoof_xpu): + hw, _ = spoof_xpu(xpu_version = "2.7") + hw.detect_hardware() + assert hw.get_package_versions().get("xpu") == "2.7" + + +def test_package_versions_xpu_available_fallback(spoof_xpu): + hw, _ = spoof_xpu(xpu_version = None) + hw.detect_hardware() + assert hw.get_package_versions().get("xpu") == "available" + + +def test_per_device_info_mem_get_info_ok(spoof_xpu): + hw, _ = spoof_xpu(total_gb = 16.0, used_gb = 1.0) + hw.detect_hardware() + info = hw._torch_get_per_device_info([0]) + assert len(info) == 1 + assert info[0]["total_gb"] == pytest.approx(16.0, abs = 0.1) + assert info[0]["used_gb"] == pytest.approx(1.0, abs = 0.1) + + +def test_mem_get_info_runtimeerror_keeps_device_with_unknown_usage(spoof_xpu): + # Arc B580 and Lunar Lake can reject mem_get_info while remaining usable. + hw, _ = spoof_xpu(mem_get_info = "raise", total_gb = 16.0, device_count = 2) + hw.detect_hardware() + + info = hw._torch_get_per_device_info([0]) + assert len(info) == 1 + assert info[0]["total_gb"] == pytest.approx(16.0, abs = 0.1) + assert info[0]["used_gb"] is None + + utilization = hw.get_visible_gpu_utilization() + assert utilization["available"] is True + assert len(utilization["devices"]) == 2 + assert all(device["vram_total_gb"] == pytest.approx(16.0) for device in utilization["devices"]) + assert all(device["vram_used_gb"] is None for device in utilization["devices"]) + + visibility = hw.get_backend_visible_gpu_info() + assert visibility["available"] is True + assert len(visibility["devices"]) == 2 + assert all(device["memory_total_gb"] == pytest.approx(16.0) for device in visibility["devices"]) + + +def test_per_device_info_no_mem_get_info_uses_none(spoof_xpu): + hw, _ = spoof_xpu(mem_get_info = "absent") + hw.detect_hardware() + info = hw._torch_get_per_device_info([0]) + assert len(info) == 1 + assert info[0]["used_gb"] is None + + +# ---------- training-device wiring ---------- + + +def test_get_torch_device_str_is_xpu(spoof_xpu): + hw, _ = spoof_xpu() + hw.detect_hardware() + assert hw.get_torch_device_str() == "xpu" + + +def test_dataset_map_num_proc_none_after_xpu_init(spoof_xpu): + # os.fork() after Level-Zero init corrupts the XPU context -> force in-process. + hw, _ = spoof_xpu(is_initialized = True) + hw.detect_hardware() + assert hw.dataset_map_num_proc(4) is None From 140b3fbe057398fa22b075b20366c62d10e936e5 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Fri, 24 Jul 2026 00:48:54 -0700 Subject: [PATCH 020/161] Studio: register text-ui tokens with tailwind-merge so cn() keeps them (#7396) * Studio: register text-ui tokens with tailwind-merge so cn keeps them Stock tailwind-merge classifies text-ui-* as a text color, so cn() dropped the size class whenever a color utility followed it in the same call. The element then fell back to the unscaled 16px root font, which made hub tabs and capability pills look oversized at small UI font sizes. Extend the merge config so text-ui-* and leading-ui-* resolve as font-size and line-height groups, and cover the failure in the contract and Playwright regression tests. * Studio: rename the Models page to Model hub Page heading, sidebar navigation label in all locales, and the chat download toasts that point at the tab. --- .../frontend/src/features/chat/chat-page.tsx | 4 ++-- .../features/hub/catalog/models-header.tsx | 2 +- studio/frontend/src/i18n/locales/ar.ts | 2 +- studio/frontend/src/i18n/locales/de.ts | 2 +- studio/frontend/src/i18n/locales/en.ts | 2 +- studio/frontend/src/i18n/locales/es.ts | 2 +- studio/frontend/src/i18n/locales/fr.ts | 2 +- studio/frontend/src/i18n/locales/hi.ts | 2 +- studio/frontend/src/i18n/locales/ja.ts | 2 +- studio/frontend/src/i18n/locales/ko.ts | 2 +- studio/frontend/src/i18n/locales/pt-br.ts | 2 +- studio/frontend/src/i18n/locales/ru.ts | 2 +- studio/frontend/src/i18n/locales/zh-CN.ts | 2 +- studio/frontend/src/lib/utils.ts | 16 +++++++++++++++- tests/studio/playwright_ui_font_scale.py | 19 +++++++++++++++++++ tests/studio/test_ui_font_scale_contract.py | 11 +++++++++++ 16 files changed, 59 insertions(+), 15 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index c241607e28..7452cf3447 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -2289,7 +2289,7 @@ export function ChatPage({ } else if (outcome === "conflict") { toast.info("Resume this download from Models", { description: - "An earlier partial download used a different transport. Open the Models tab to resume or restart it.", + "An earlier partial download used a different transport. Open the Model hub tab to resume or restart it.", }); } else if (outcome === "busy") { toast.info("Download already in progress", { @@ -2410,7 +2410,7 @@ export function ChatPage({ // surface's onComplete auto-loads, mirroring the "started" branch. toast.info("Resume this download from Models", { description: - "An earlier partial download used a different transport. Open the Models tab to resume or restart it.", + "An earlier partial download used a different transport. Open the Model hub tab to resume or restart it.", }); return; } diff --git a/studio/frontend/src/features/hub/catalog/models-header.tsx b/studio/frontend/src/features/hub/catalog/models-header.tsx index f0e0950871..f844336d4a 100644 --- a/studio/frontend/src/features/hub/catalog/models-header.tsx +++ b/studio/frontend/src/features/hub/catalog/models-header.tsx @@ -64,7 +64,7 @@ export function ModelsHeader({ return (
/^ui-\d+(p5)?$/.test(value); + +const twMerge = extendTailwindMerge({ + extend: { + classGroups: { + "font-size": [{ text: [isUiToken] }], + leading: [{ leading: [isUiToken] }], + }, + }, +}); export function cn(...inputs: ClassValue[]): string { return twMerge(clsx(inputs)); diff --git a/tests/studio/playwright_ui_font_scale.py b/tests/studio/playwright_ui_font_scale.py index 89c7894929..0f14c42422 100644 --- a/tests/studio/playwright_ui_font_scale.py +++ b/tests/studio/playwright_ui_font_scale.py @@ -193,6 +193,25 @@ def main(): page.set_viewport_size({"width": 1440, "height": 900}) page.wait_for_timeout(400) + step("cn keeps text-ui-* next to color classes (hub tabs)") + page.keyboard.press("Escape") + page.wait_for_timeout(400) + page.goto(f"{BASE}/hub", wait_until = "domcontentloaded") + page.wait_for_timeout(2000) + open_appearance(page) + set_input(page, "UI font size", 12) + page.keyboard.press("Escape") + page.wait_for_timeout(400) + tab = page.get_by_role("radio").filter(has_text = "Discover").first + tab.wait_for(state = "visible", timeout = 15000) + tab_font = tab.evaluate("el => parseFloat(getComputedStyle(el).fontSize)") + # text-ui-12p5 at scale 0.75; 16px means twMerge dropped the token. + if not near(tab_font, 12.5 * 12 / 16): + fail(f"hub tab font did not scale (twMerge drop?): {tab_font}") + page.goto(BASE, wait_until = "domcontentloaded") + page.wait_for_timeout(1500) + open_appearance(page) + step("default restores exactly") page.get_by_role("dialog").get_by_role("button").filter(has_text = "Appearance").first.click() page.wait_for_timeout(500) diff --git a/tests/studio/test_ui_font_scale_contract.py b/tests/studio/test_ui_font_scale_contract.py index 393c60dd6b..1153eea643 100644 --- a/tests/studio/test_ui_font_scale_contract.py +++ b/tests/studio/test_ui_font_scale_contract.py @@ -17,6 +17,7 @@ SRC = REPO / "studio/frontend/src" INDEX_CSS = (SRC / "index.css").read_text(encoding = "utf-8") STORE = (SRC / "features/settings/stores/appearance-custom-store.ts").read_text(encoding = "utf-8") SELECT = (SRC / "components/ui/select.tsx").read_text(encoding = "utf-8") +UTILS = (SRC / "lib/utils.ts").read_text(encoding = "utf-8") # Raw numeric fontSize props are only allowed where a scaled stylesheet rule # (.recharts-text) overrides the presentation attribute at render time. @@ -87,6 +88,16 @@ def test_radix_select_viewport_owns_the_scroll_state(): assert "overflow-y-auto" not in content_cls.group(1) +def test_cn_knows_the_ui_typography_tokens(): + """Stock tailwind-merge classifies text-ui-* as a text color and deletes + it whenever a real color class follows in the same cn() call, so the + element falls back to the unscaled inherited font size.""" + assert "extendTailwindMerge" in UTILS + assert '"font-size": [{ text: [isUiToken] }]' in UTILS + assert "leading: [{ leading: [isUiToken] }]" in UTILS + assert "/^ui-\\d+(p5)?$/.test(value)" in UTILS + + def test_no_raw_pixel_text_utilities(): offenders = [] for path in _frontend_sources(): From 6e91d1dff8c8a8720e2be5afb8797f7819b63a55 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Jul 2026 02:12:00 -0700 Subject: [PATCH 021/161] Studio: scan HF cache snapshot loads by their repo id (#7398) * Studio: scan HF cache snapshot loads by their repo id Inactive Hugging Face caches (legacy, default, and previously selected download locations) are loaded by their resolved snapshot path so they keep using the selected cache instead of re-downloading. That path is a local filesystem path, so evaluate_file_security exempted it with "local path; no Hub scan" and skipped Hugging Face's pickle/malware scan. Active caches load by repo id and are still scanned, so the same model could dodge the gate simply by being in an inactive cache. An HF cache snapshot keeps the canonical models--org--repo/snapshots/ layout, so recover the repo id from that path and scan it instead of exempting it. Non-cache local paths (models directory, custom folders) still skip the scan, and a remote ref is still scanned by repo id. Adds a regression test that a flagged pickle in an inactive-cache snapshot path blocks the load. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: scan the exact cached commit for inactive HF caches An HF cache snapshot path encodes the commit, not just the repo id (models--org--repo/snapshots/). Recover the revision alongside the repo id and pass it to model_info and the shard-index lookup so the scan covers the exact files that will be deserialized, rather than the repo's default branch. Without this, a pickle in an older cached commit that was later removed from the branch would scan clean and still load. Extends the regression test to assert the recovered revision is forwarded to the Hub scan. --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/tests/test_file_security.py | 17 +++++++ .../backend/utils/security/file_security.py | 51 ++++++++++++++++--- 2 files changed, 61 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_file_security.py b/studio/backend/tests/test_file_security.py index b4c8f5d242..e02c33a0f1 100644 --- a/studio/backend/tests/test_file_security.py +++ b/studio/backend/tests/test_file_security.py @@ -165,6 +165,23 @@ def test_skips_local_path(): assert "local" in d.reason +def test_scans_inactive_hf_cache_snapshot_path(tmp_path): + # An inactive HF cache loads by snapshot path; the gate must recover the repo id + + # commit from models--org--repo/snapshots/ and scan that exact commit, not exempt + # it and not fall back to the default branch (an older commit may hold a dropped pickle). + snapshot = tmp_path / "models--evil--repo" / "snapshots" / "deadbeef" + snapshot.mkdir(parents = True) + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with _patch_status(status) as model_info: + d = evaluate_file_security(str(snapshot)) + assert d.blocked is True + assert model_info.call_args.args[0] == "evil/repo" + assert model_info.call_args.kwargs["revision"] == "deadbeef" + + def test_remote_gguf_named_repo_is_still_scanned(): # Only LOCAL paths skip the Hub scan, so a remote .gguf repo is still scanned and a # poisoned pickle smuggled into it is blocked. diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 892f7862a9..3e12c15096 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -114,6 +114,28 @@ def _file_suffix(path: str) -> str: return "." + base.rsplit(".", 1)[1].lower() if "." in base else "" +def _hf_cache_snapshot_ref(local_path: str) -> Optional[tuple]: + """``(repo_id, revision)`` for an HF-cache snapshot path, else None. An inactive Studio + cache loads by its snapshot path but keeps the ``models--org--repo/snapshots/`` + layout, so the gate recovers its provenance and scans that exact commit instead of + exempting it (an older cached commit can hold a pickle since dropped from the branch).""" + try: + path = Path(local_path).resolve(strict = False) + except (OSError, ValueError): + return None + for parent in path.parents: + if parent.name != "snapshots": + continue + encoded = parent.parent.name + if not encoded.startswith("models--"): + return None + repo_id = encoded.removeprefix("models--").replace("--", "/") + if not repo_id: + return None + return repo_id, path.relative_to(parent).parts[0] # dir under snapshots/ + return None + + def _load_relative_path(norm: str, load_subdirs) -> str: """``norm`` relative to a ``from_pretrained`` load root. Some loads read from a snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``/LLM``), where a file @@ -141,13 +163,14 @@ def _indexed_shard_paths( model_name: str, hf_token: Optional[str], load_subdirs = (), + revision: Optional[str] = None, ): """Repo-relative weight paths a load could fetch via weight-index files. Returns a set (empty when the repo ships no index files -- a definitive "nothing sharded"), or None when the lookup was inconclusive (transient error) so the caller treats a flagged subdir pickle conservatively. Reads only small JSON indexes, never weights. Indexes are looked up at the root and each ``load_subdirs`` root, with ``weight_map`` - entries re-prefixed to repo-relative paths. + entries re-prefixed to repo-relative paths. ``revision`` scopes to a cached commit. """ import json @@ -166,6 +189,7 @@ def _indexed_shard_paths( index_path = hf_hub_download( model_name, prefix + filename, + revision = revision, token = hf_token or None, cache_dir = active_hf_hub_cache(), ) @@ -260,9 +284,14 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: return model_name, load_subdirs -def _fetch_security_status(model_name: str, hf_token: Optional[str]): +def _fetch_security_status( + model_name: str, + hf_token: Optional[str], + revision: Optional[str] = None, +): """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; retries once on a transient error, then returns None so the caller fails open. + ``revision`` scopes the scan to a specific cached commit (else the default branch). """ from huggingface_hub import model_info as hf_model_info @@ -272,6 +301,7 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): try: info = hf_model_info( model_name, + revision = revision, token = token_arg, securityStatus = True, timeout = timeout, @@ -485,12 +515,17 @@ def evaluate_file_security( # fails open): the Spark-TTS "/LLM" alias is really unsloth/ from LLM/. model_name, load_subdirs = _load_scan_target(model_name, tuple(load_subdirs)) - # Local paths (including a local .gguf) have no Hub scan. A remote ref is scanned - # even if named "*.gguf", so a repo cannot dodge the scan via its name. + # Local paths have no Hub scan, EXCEPT an HF-cache snapshot whose canonical path + # encodes a repo id + commit: scan that exact commit so an inactive-cache load can't + # dodge the gate. A remote ref is scanned even if named "*.gguf" (name can't dodge it). + snapshot_revision = None try: from utils.paths import is_local_path if is_local_path(model_name): - return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan") + cache_ref = _hf_cache_snapshot_ref(model_name) + if cache_ref is None: + return FileSecurityDecision(model_name, False, reason = "local path; no Hub scan") + model_name, snapshot_revision = cache_ref except Exception: # Cannot classify the path -> do not block on that account. return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked") @@ -499,7 +534,7 @@ def evaluate_file_security( if local_only_load: return _evaluate_local_only(model_name) - status = _fetch_security_status(model_name, hf_token) + status = _fetch_security_status(model_name, hf_token, revision = snapshot_revision) if not isinstance(status, dict): return FileSecurityDecision( model_name, False, reason = "scan unavailable; allowed (fail-open)" @@ -536,7 +571,9 @@ def evaluate_file_security( maybe_shard.append({"path": path, "level": level, "norm": norm}) if maybe_shard: - indexed = _indexed_shard_paths(model_name, hf_token, load_subdirs) + indexed = _indexed_shard_paths( + model_name, hf_token, load_subdirs, revision = snapshot_revision + ) for m in maybe_shard: # Block if a root index lists this shard, or if the lookup was inconclusive # (transient error -> stay conservative). A definitive "no index / not listed" From 418ae14388318dd9f9d152afb3d36b2afc086f2f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Jul 2026 02:12:16 -0700 Subject: [PATCH 022/161] Fix ROCm wheel-index unit test: extract the gfx-arch probe helpers get_torch_index_url now calls (#7399) * Fix ROCm wheel-index test: extract the gfx-arch probe helpers get_torch_index_url now calls get_torch_index_url gained a gfx-arch probe on the ROCm path (Strix reroute work) and now calls _ensure_rocm_probe_env, _probe_amd_gfx_arch, _infer_linux_amd_gfx_arch and friends. The unit test in tests/sh/test_get_torch_index_url.sh sources a curated subset of install.sh functions, and that list was never updated, so those helpers were undefined in the harness. On the ROCm path the gfx probe hit an undefined function, the branch silently fell through to the CPU wheel index, and every ROCm assertion failed (9 failures: all ROCm versions resolved to /whl/cpu). Extract the six missing helpers so the ROCm branch runs end to end. All 49 assertions pass. Adds a comment noting these must stay in sync with install.sh. * Keep the ROCm wheel-index test hermetic: redirect the /opt/rocm prefix Extracting _ensure_rocm_probe_env pulled its absolute-path host probe into the harness: it appends /opt/rocm/bin to PATH and runs the real host rocminfo, and version detection reads /opt/rocm/.info/version. On a host with ROCm installed that leaks the host GPU into the minimal-PATH test, so the no-GPU and CUDA-visible-device assertions could select a host ROCm wheel index instead of their expected CPU result, making the test host-dependent. Redirect the whole /opt/rocm prefix to an empty temp dir in the same sed pass that stubs /usr/bin/nvidia-smi, so the probes stay hermetic. All 49 assertions pass and the generated harness contains no real /opt/rocm path. --------- Co-authored-by: danielhanchen --- tests/sh/test_get_torch_index_url.sh | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 23902097ef..633d8ca17b 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -14,6 +14,13 @@ FAIL=0 # controllable path so we can test the "no GPU" scenario on GPU machines. _FUNC_FILE=$(mktemp) _FAKE_SMI_DIR=$(mktemp -d) +# The ROCm probe helpers read the real /opt/rocm prefix by ABSOLUTE path +# (_ensure_rocm_probe_env appends /opt/rocm/bin to PATH and runs the host +# rocminfo; version detection reads /opt/rocm/.info/version). On a real ROCm +# host that leaks the host GPU into the minimal-PATH harness and makes the +# no-GPU / CPU assertions host-dependent. Redirect the whole prefix to an empty +# temp dir so the probes stay hermetic (same idea as the nvidia-smi rewrite). +_FAKE_ROCM_DIR=$(mktemp -d) { sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH" echo "" @@ -23,10 +30,28 @@ _FAKE_SMI_DIR=$(mktemp -d) echo "" sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH" echo "" + # ROCm gfx-arch probe helpers that get_torch_index_url / _has_amd_rocm_gpu + # now call. These MUST stay in sync with install.sh: if get_torch_index_url + # references a helper that is not extracted here, the ROCm branch hits an + # undefined function, silently falls through to the CPU wheel index, and the + # ROCm assertions below fail. + sed -n '/^_ensure_rocm_probe_env()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_probe_amd_gfx_arch()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_amd_gpu_present_via_pci()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_infer_amd_gfx_arch_from_gpu_name()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_infer_linux_amd_gfx_arch()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_amd_arch_index_family_for_gfx()/,/^}/p' "$INSTALL_SH" + echo "" sed -n '/^_trim_index_path_slashes()/,/^}/p' "$INSTALL_SH" echo "" sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH" -} | sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \ +} | sed -e "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \ + -e "s|/opt/rocm|$_FAKE_ROCM_DIR|g" \ > "$_FUNC_FILE" # Save system PATH so we always have basic tools (uname, grep, head, etc.) @@ -438,6 +463,7 @@ assert_eq "url override preserves fragment slash" "https://mirror.example.com/wh rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" +rm -rf "$_FAKE_ROCM_DIR" rm -rf "$_TOOLS_DIR" echo "" From 0e800d213aacd48ae8c6dcf6e0745182fb8991af Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:43:52 +0530 Subject: [PATCH 023/161] fix(studio): stop false MTP/vision capability reports (#7332) * fix(studio): stop false MTP/vision capability reports (#7302) MTP probing only inspected the first physical --spec-type help line and treated empty/crash --help output as "lacks MTP", which false-warned on otherwise capable builds. Parse the full --spec-type help block, fail open when the probe is inconclusive, and stop blaming bare mmproj crashes on a projector-format mismatch when the text-only retry also fails. Fixes #7302 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): tighten MTP probe semantics per Codex review (#7302) Treat nonempty --help without --spec-type as definitive no-MTP, keep only empty/crash probes inconclusive, skip binary_no_mtp UI hint on inconclusive loads, and stop reporting supports_mtp=True in /status for unknown probes. * Treat failed llama-server --help probes as inconclusive (#7302) Gate definitive no-MTP results on a zero exit code so crash diagnostics with nonempty stderr do not re-enable the false lacks-MTP warning path. * Add returncode to probe test mock so probe_ok gating passes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fail open in /status when the MTP probe is inconclusive * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Report missing llama-server as lacking MTP in /status * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten comments in MTP/mmproj probe changes --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 155 +++++++++++++----- studio/backend/main.py | 6 +- studio/backend/routes/inference.py | 9 +- .../tests/test_llama_cpp_mmproj_fallback.py | 21 +++ .../tests/test_llama_cpp_mtp_detection.py | 114 ++++++++++++- 5 files changed, 262 insertions(+), 43 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 7d67339c1e..1fe134c3f9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2781,6 +2781,7 @@ class LlamaCppBackend: "found": False, "mtp_token": None, "supports_mtp": False, + "mtp_probe_inconclusive": True, "ngram_mod_flavor": None, "supports_ngram_mod": False, "spec_draft_n_max_flag": None, @@ -2813,6 +2814,9 @@ class LlamaCppBackend: supports_no_cache_prompt = False supports_metrics = False supports_slot_save = False + saw_spec_type = False + probe_ok = False + help_text = "" try: probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( @@ -2824,6 +2828,7 @@ class LlamaCppBackend: check = False, env = probe_env, ) + probe_ok = result.returncode == 0 help_text = (result.stdout or "") + "\n" + (result.stderr or "") # Split into per-flag blocks (each --flag line + its indented # continuation), so the "argument has been removed" description @@ -2868,17 +2873,19 @@ class LlamaCppBackend: return False return "argument has been removed" not in desc - # MTP token from the --spec-type line. - spec_line = "" - for line in help_text.splitlines(): - if "--spec-type" in line: - spec_line = line - break - # PR #22673 used draft-mtp; later renamed to mtp. - if "draft-mtp" in spec_line: - mtp_token = "draft-mtp" - elif re.search(r"[|,\[]mtp[|,\]]", spec_line): - mtp_token = "mtp" + # MTP token from the full --spec-type help block (decl + indented + # continuation). First-line-only probing missed builds putting the + # enum on the next line (#7302). Prefer draft-mtp (PR #22673) over mtp. + spec_help = blocks.get("--spec-type") or "" + if not spec_help: + # Fallback: join --spec-type lines, avoiding incidental "mtp" in --help. + spec_help = "\n".join( + line for line in help_text.splitlines() if "--spec-type" in line + ) + mtp_token = cls._mtp_token_from_spec_help(spec_help) + # Only a resolved --spec-type block confirms missing MTP; empty/crash + # leaves saw_spec_type False so supports_mtp fails open. + saw_spec_type = bool(spec_help.strip()) and "--spec-type" in spec_help # ngram-mod flag flavor. Post-rename builds advertise both new # args (real) and legacy ones (stubs); pre-rename builds only @@ -2914,11 +2921,29 @@ class LlamaCppBackend: supports_slot_save = _is_real("--slot-save-path") except (OSError, subprocess.SubprocessError) as exc: logger.debug(f"llama-server --help probe failed: {exc}") + saw_spec_type = False + probe_ok = False + help_text = "" + + help_nonempty = bool(help_text.strip()) + # Confirmed only when a successful --help lists a --spec-type block with + # mtp/draft-mtp; nonempty --help without it is a definitive pre-spec + # binary; failed/empty probes stay inconclusive (#7302). + if saw_spec_type and probe_ok: + supports_mtp = mtp_token is not None + mtp_probe_inconclusive = False + elif help_nonempty and probe_ok: + supports_mtp = False + mtp_probe_inconclusive = False + else: + supports_mtp = False + mtp_probe_inconclusive = True info = { "found": True, "mtp_token": mtp_token, - "supports_mtp": mtp_token is not None, + "supports_mtp": supports_mtp, + "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": ngram_mod_flavor, "supports_ngram_mod": ngram_mod_flavor is not None, "spec_draft_n_max_flag": spec_draft_n_max_flag, @@ -2934,6 +2959,21 @@ class LlamaCppBackend: cls._capability_cache[cache_key] = info return info + @staticmethod + def _mtp_token_from_spec_help(spec_help: str) -> Optional[str]: + """Extract ``draft-mtp`` / ``mtp`` from a ``--spec-type`` help snippet. + + Prefers ``draft-mtp`` (llama.cpp PR #22673) over the later bare ``mtp`` + rename. Returns ``None`` when neither token appears as an enum value. + """ + text = spec_help or "" + if "draft-mtp" in text: + return "draft-mtp" + # Bare `mtp` enum token (`|mtp|`, `,mtp,`, ...), not a substring. + if re.search(r"(? str: + """User-facing error when the text-only --mmproj strip retry also fails. + + Confirmed projector-format mismatches keep the historical wording. + Bare signal crashes (common on some ROCm/driver paths) must not be + reported as "Vision projector incompatible" — that misled #7302. + """ + if projector_confirmed: + return ( + "Vision projector incompatible with this llama.cpp " + "build, and the text-only retry also failed: " + detail + ) + return ( + "Vision model failed to start (llama-server crashed with " + "--mmproj), and the text-only retry also failed: " + detail + ) + @staticmethod def _output_has_nonprojector_diagnostic(output: str) -> bool: """True when the output already names a concrete non-projector cause (out @@ -8388,23 +8446,29 @@ class LlamaCppBackend: self._kill_process() # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). + _projector_msg = self._is_projector_incompatibility(out) + _signal_mmproj_guess = self._is_signal_crash( + _crash_rc + ) and not self._output_has_nonprojector_diagnostic(out) if ( launched_with_mmproj and not self._cancel_event.is_set() - and ( - self._is_projector_incompatibility(out) - or ( - self._is_signal_crash(_crash_rc) - and not self._output_has_nonprojector_diagnostic(out) - ) - ) + and (_projector_msg or _signal_mmproj_guess) ): - logger.warning( - "llama-server could not load this model's vision " - "projector (--mmproj). The installed llama.cpp build is " - "likely too old for it. Loading text-only for this " - "session; run 'unsloth studio update' to enable vision." - ) + if _projector_msg: + logger.warning( + "llama-server could not load this model's vision " + "projector (--mmproj). The installed llama.cpp build is " + "likely too old for it. Loading text-only for this " + "session; run 'unsloth studio update' to enable vision." + ) + else: + logger.warning( + "llama-server crashed while loading this model's vision " + "projector (--mmproj). Retrying text-only for this " + "session; if this persists, run 'unsloth studio update' " + "or check GPU/driver logs." + ) cmd = self._strip_mmproj_args(_last_spawn_cmd) # This retry bypasses _spawn_and_wait, so refresh the # launched-argv snapshot itself -- the zero-offload @@ -8432,14 +8496,16 @@ class LlamaCppBackend: "(e.g. ROCR_VISIBLE_DEVICES=0 exposes only the first " "GPU) before launching Unsloth Studio." ) + _retry_detail = self._classify_llama_start_failure( + "\n".join(self._stdout_lines[-50:]), + gguf_path, + self._model_identifier, + _retry_rc, + ) raise RuntimeError( - "Vision projector incompatible with this llama.cpp " - "build, and the text-only retry also failed: " - + self._classify_llama_start_failure( - "\n".join(self._stdout_lines[-50:]), - gguf_path, - self._model_identifier, - _retry_rc, + self._mmproj_retry_failure_message( + projector_confirmed = _projector_msg, + detail = _retry_detail, ) ) else: @@ -8669,18 +8735,29 @@ class LlamaCppBackend: caps = self.probe_server_capabilities(binary) mtp_token = caps.get("mtp_token") if caps else None if not mtp_token: - logger.warning( - "Requested MTP speculative decoding but " - "llama-server lacks --spec-type mtp/draft-mtp; " - "run `unsloth studio update`. Loading without " - "speculative decoding." - ) + inconclusive = bool(caps.get("mtp_probe_inconclusive")) if caps else True + if inconclusive: + logger.info( + "Requested MTP speculative decoding but llama-server MTP " + "capability probe was inconclusive; loading without " + "speculative decoding." + ) + else: + logger.warning( + "Requested MTP speculative decoding but " + "llama-server lacks --spec-type mtp/draft-mtp; " + "run `unsloth studio update`. Loading without " + "speculative decoding." + ) # Override an inherited LLAMA_ARG_SPEC_TYPE=draft-mtp (CLI wins # over env) so the child matches the binary-capability gate and # the no-MTP budget, like the sibling no-head/non-MTP fallbacks. flags.append("--spec-default") self._speculative_type = "default" - self._spec_fallback_reason = "binary_no_mtp" + if inconclusive: + self._spec_fallback_reason = None + else: + self._spec_fallback_reason = "binary_no_mtp" return False draft_n_max = _resolved_draft_n_max() n_max_flag = caps.get("spec_draft_n_max_flag") or "--spec-draft-n-max" diff --git a/studio/backend/main.py b/studio/backend/main.py index a538f935ff..5af25efa74 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -438,7 +438,11 @@ def _run_llama_cpp_startup_probes(app: FastAPI) -> None: import structlog as _structlog _log = _structlog.get_logger(__name__) - if _caps.get("found") and not _caps.get("supports_mtp"): + if ( + _caps.get("found") + and not _caps.get("supports_mtp") + and not _caps.get("mtp_probe_inconclusive") + ): _msg = ( "llama.cpp prebuilt lacks MTP support " "(--spec-type mtp/draft-mtp). Run `unsloth studio update`. " diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 6dbaa8fcc9..445a26f04d 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5835,10 +5835,15 @@ async def get_status(current_subject: str = Depends(get_current_subject)): try: _bin = type(llama_backend)._find_llama_server_binary() _caps = type(llama_backend).probe_server_capabilities(_bin) - _supports_mtp = bool(_caps.get("supports_mtp", False)) + # Fail open on inconclusive probes: False means a definitive + # "binary lacks MTP" to API consumers. + _supports_mtp = bool( + _caps.get("supports_mtp", False) + or (_caps.get("found", False) and _caps.get("mtp_probe_inconclusive", False)) + ) except Exception: _bin = None - _supports_mtp = True # fail open + _supports_mtp = False # no usable binary: MTP genuinely unavailable try: from utils.llama_cpp_freshness import check_prebuilt_freshness _freshness = check_prebuilt_freshness(_bin) diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 049058e511..4332a440a5 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -335,3 +335,24 @@ class TestRetryContract: def test_external_kill_skips_flash_attn_retry(self): # SIGKILL (-9, OOM killer) is not a program fault: no FA-off retry. assert _signal_crash(-9) is False + + +class TestMmprojRetryFailureMessage: + """#7302: bare mmproj crashes must not be reported as projector-format.""" + + def test_confirmed_projector_keeps_historical_wording(self): + msg = LlamaCppBackend._mmproj_retry_failure_message( + projector_confirmed = True, + detail = "llama-server failed to start", + ) + assert msg.startswith("Vision projector incompatible with this llama.cpp") + assert "llama-server failed to start" in msg + + def test_bare_crash_does_not_claim_projector_incompatibility(self): + msg = LlamaCppBackend._mmproj_retry_failure_message( + projector_confirmed = False, + detail = "llama-server failed to start. Check that the GGUF file is valid", + ) + assert "Vision projector incompatible" not in msg + assert "crashed with --mmproj" in msg + assert "GGUF file is valid" in msg diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 1d15647967..27c1b17a85 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -637,7 +637,9 @@ def test_probe_server_capabilities_uses_binary_library_env(tmp_path, monkeypatch def fake_run(cmd, **kwargs): captured["cmd"] = cmd captured["env"] = kwargs.get("env") - return _types.SimpleNamespace(stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "") + return _types.SimpleNamespace( + stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "", returncode = 0 + ) monkeypatch.setattr("core.inference.llama_cpp.subprocess.run", fake_run) @@ -678,6 +680,95 @@ def test_probe_server_capabilities_reports_outdated_binary(tmp_path): assert caps["found"] is True assert caps["mtp_token"] is None assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_reads_mtp_from_multiline_help(tmp_path): + # Enum on the indented line: first-line-only probing falsely reported + # "lacks MTP" (#7302). + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--spec-type TYPE\n" + " speculative decoding type\n" + " (none,draft-simple,draft-mtp,ngram-mod)\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["mtp_token"] == "draft-mtp" + assert caps["supports_mtp"] is True + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_empty_help_fails_open(tmp_path): + # --help prints nothing: must not claim the prebuilt lacks MTP (#7302). + fake = tmp_path / "llama-server" + fake.write_text("#!/usr/bin/env bash\nexit 0\n") + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +@_NEEDS_BASH +def test_probe_server_capabilities_no_spec_type_is_definitive(tmp_path): + # Nonempty --help without --spec-type: pre-spec binary, not inconclusive. + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--gpu-layers N\n GPU layers to offload\n", + ) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is False + + +@_NEEDS_BASH +def test_probe_server_capabilities_failed_help_with_output_is_inconclusive(tmp_path): + fake = tmp_path / "llama-server" + fake.write_text( + "#!/usr/bin/env bash\n" + 'if [ "$1" = "--help" ]; then\n' + " echo 'illegal instruction'\n" + " exit 1\n" + "fi\n" + ) + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +@_NEEDS_BASH +def test_probe_server_capabilities_crash_on_help_fails_open(tmp_path): + fake = tmp_path / "llama-server" + fake.write_text("#!/usr/bin/env bash\nkill -SEGV $$\n") + fake.chmod(0o755) + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + assert caps["found"] is True + assert caps["mtp_token"] is None + assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True + + +def test_mtp_token_from_spec_help_prefers_draft_mtp(): + assert ( + LlamaCppBackend._mtp_token_from_spec_help("--spec-type none,draft-mtp,mtp,ngram-mod") + == "draft-mtp" + ) + assert LlamaCppBackend._mtp_token_from_spec_help("--spec-type [none|mtp|ngram-cache]") == "mtp" + assert LlamaCppBackend._mtp_token_from_spec_help("--spec-type none,ngram-mod") is None + # No incidental substring matches. + assert LlamaCppBackend._mtp_token_from_spec_help("prompt cache") is None def test_probe_server_capabilities_handles_missing_binary(): @@ -685,6 +776,7 @@ def test_probe_server_capabilities_handles_missing_binary(): caps = LlamaCppBackend.probe_server_capabilities("/no/such/llama-server") assert caps["found"] is False assert caps["supports_mtp"] is False + assert caps["mtp_probe_inconclusive"] is True assert caps["supports_cache_ram"] is False assert caps["supports_ctx_checkpoints"] is False assert caps["supports_no_cache_prompt"] is False @@ -1176,12 +1268,14 @@ def _resolver_backend( *, ngram_supported = True, mtp_token = "draft-mtp", + mtp_probe_inconclusive = False, ): """Backend with a deterministic probe so the resolver is hermetic.""" fake = { "found": True, "mtp_token": mtp_token, "supports_mtp": bool(mtp_token), + "mtp_probe_inconclusive": mtp_probe_inconclusive, "ngram_mod_flavor": "new" if ngram_supported else None, "supports_ngram_mod": bool(ngram_supported), "spec_draft_n_max_flag": "--spec-draft-n-max", @@ -1879,6 +1973,24 @@ def test_spec_fallback_reason_set_when_binary_lacks_mtp(monkeypatch): assert backend.spec_fallback_reason == "binary_no_mtp" +def test_spec_fallback_reason_none_when_mtp_probe_inconclusive(monkeypatch): + backend = _resolver_backend( + monkeypatch, + mtp_token = None, + mtp_probe_inconclusive = True, + ) + backend._build_speculative_flags( + speculative_type = "mtp", + spec_draft_n_max = None, + extra_args = None, + model_identifier = _MTP_MODEL, + model_path = None, + gpus = True, + binary = "/fake/llama-server", + ) + assert backend.spec_fallback_reason is None + + def test_spec_fallback_reason_none_when_mtp_engages(monkeypatch): backend = _resolver_backend(monkeypatch) backend._build_speculative_flags( From 330586de7c9e849c400b8bbc5112fcde8e964081 Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:52:03 +0530 Subject: [PATCH 024/161] feat(studio): expose full KV cache dtype list in model config UI (#7348) Fixes #7244 The Studio per-model config dropdown only surfaced bf16, q8_0, q5_1, and q4_1 even though llama.cpp already accepts q4_0, q5_0, iq4_nl, and f32. Add the missing options to KV_CACHE_DTYPES and align API field descriptions with the backend _valid_cache_types set. Co-authored-by: Daniel Han --- studio/backend/models/inference.py | 16 +++++++++++++--- studio/backend/routes/models.py | 5 ++++- .../components/model-config-page.tsx | 3 ++- .../model-config/per-model-config.ts | 12 +++++++++++- 4 files changed, 30 insertions(+), 6 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 2663242187..1758efe515 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -70,7 +70,10 @@ class LoadRequest(BaseModel): cache_type_kv: Optional[str] = Field( None, - description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')", + description = ( + "KV cache data type for both K and V " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32')" + ), ) gpu_ids: Optional[List[int]] = Field( None, @@ -442,7 +445,10 @@ class LoadResponse(BaseModel): ) cache_type_kv: Optional[str] = Field( None, - description = "KV cache data type for K and V (e.g. 'f16', 'bf16', 'q8_0')", + description = ( + "KV cache data type for K and V " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32')" + ), ) chat_template: Optional[str] = Field( None, @@ -602,7 +608,11 @@ class InferenceStatusResponse(BaseModel): ) cache_type_kv: Optional[str] = Field( None, - description = "KV cache quantization dtype (e.g. 'q8_0'), or None for default", + description = ( + "KV cache quantization dtype " + "(e.g. 'f16', 'bf16', 'q8_0', 'q4_0', 'q4_1', 'q5_0', 'q5_1', 'iq4_nl', 'f32'), " + "or None for default" + ), ) chat_template: Optional[str] = Field( None, description = "Model's default chat template (Jinja2 source), if any" diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 3c0d6ff4ba..ed83a12f48 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2692,7 +2692,10 @@ async def get_kv_cache_estimate( repo_id: str = Query(..., description = "HF repo ID or local path"), quant: str = Query(..., description = "Quantization label (e.g. Q4_K_M)"), n_ctx: int = Query(..., ge = 1, description = "Context length to size the KV cache for"), - cache_type_kv: Optional[str] = Query(None, description = "KV cache dtype (e.g. q8_0)"), + cache_type_kv: Optional[str] = Query( + None, + description = "KV cache dtype (e.g. q8_0, q4_0, q5_0, iq4_nl, f32)", + ), current_subject: str = Depends(get_current_subject), ): """Estimate KV cache + weight bytes for a downloaded GGUF at n_ctx. diff --git a/studio/frontend/src/features/model-picker/components/model-config-page.tsx b/studio/frontend/src/features/model-picker/components/model-config-page.tsx index ee208dc0bc..afbd33af7c 100644 --- a/studio/frontend/src/features/model-picker/components/model-config-page.tsx +++ b/studio/frontend/src/features/model-picker/components/model-config-page.tsx @@ -415,7 +415,8 @@ function GgufAdvancedSettings({ KV Cache Dtype Lower KV cache precision to save VRAM at the cost of some quality. - f16/bf16 are full precision; q8_0/q5_1/q4_1 are quantized. + f16 is the default; bf16 and f32 are full precision; q8_0 through + iq4_nl are quantized. setProjectNameDraft(event.target.value)} - onKeyDown={(event) => { - if (event.key === "Enter") { - event.preventDefault(); - void commitCreateProject(); - } - }} - autoFocus - maxLength={120} - placeholder="Project name" - aria-label="Project name" - className="focus-visible:border-input focus-visible:ring-0" - /> - - - - - - + title={ + projectCreateMoveTarget ? "Move to new project" : "Create project" + } + submitLabel={projectCreateMoveTarget ? "Create and move" : "Create project"} + onCreated={afterCreateProject} + /> ); } diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index a439a91239..7e0544e2d1 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -185,6 +185,10 @@ import { listStoredChatThreads, } from "./utils/chat-history-storage"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; +import { + consumeProjectSourcesPending, + hasProjectSourcesPending, +} from "@/features/rag/components/project-source-dropzone"; const ProjectSourcesPanel = lazy(() => @@ -998,7 +1002,14 @@ function ProjectLanding({ const active = useChatActive(); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const initialActiveThreadRef = useRef(null); - const [projectTab, setProjectTab] = useState<"chats" | "sources">("chats"); + // Land on Sources when the project was just created with dropped files. + const [projectTab, setProjectTab] = useState<"chats" | "sources">(() => + hasProjectSourcesPending(projectId) ? "sources" : "chats", + ); + // Drop the marker once committed: React may replay the initializer above. + useEffect(() => { + consumeProjectSourcesPending(projectId); + }, [projectId]); const [pendingNewThreadId, setPendingNewThreadId] = useState( null, ); diff --git a/studio/frontend/src/features/chat/components/new-project-dialog.tsx b/studio/frontend/src/features/chat/components/new-project-dialog.tsx index 880129ac6c..6aca3d36c3 100644 --- a/studio/frontend/src/features/chat/components/new-project-dialog.tsx +++ b/studio/frontend/src/features/chat/components/new-project-dialog.tsx @@ -2,7 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useNavigate } from "@tanstack/react-router"; -import { useState } from "react"; +import { useEffect, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; import { @@ -12,31 +12,92 @@ import { DialogHeader, DialogTitle, } from "@/components/ui/dialog"; -import { Input } from "@/components/ui/input"; +import { + ProjectSourceDropzone, + type StagedSource, + uploadStagedSources, +} from "@/features/rag/components/project-source-dropzone"; import { toast } from "@/lib/toast"; +import { Folder02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { createChatProject } from "../hooks/use-chat-projects"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import type { ProjectRecord } from "../types"; -// Create-project dialog usable from the composer + menu. Creating opens the new -// project straight away rather than dropping the user on the projects list. +function currentRoute(): string { + if (typeof window === "undefined") return ""; + return window.location.pathname + window.location.search; +} + +// Create-project dialog for the composer, sidebar, and projects page. Creating +// opens the new project; `onCreated` overrides that for callers with their own +// follow-up (the sidebar's "move this chat to a new project"). export function NewProjectDialog({ open, onOpenChange, + title = "Create project", + submitLabel = "Create project", + onCreated, }: { open: boolean; onOpenChange: (open: boolean) => void; + title?: string; + submitLabel?: string; + onCreated?: ( + project: ProjectRecord, + context: { stayedOnRoute: boolean }, + ) => void | Promise; }) { const navigate = useNavigate(); const [name, setName] = useState(""); + const [staged, setStaged] = useState([]); + const [busy, setBusy] = useState(false); + // Uploads outlive this component, so a slow one must not yank the user to the + // new project after they have navigated away. + const mounted = useRef(true); + useEffect(() => { + // Set on setup, not just cleared on cleanup: StrictMode replays + // setup/cleanup/setup, which would otherwise leave this false forever. + mounted.current = true; + return () => { + mounted.current = false; + }; + }, []); + + function reset() { + setName(""); + setStaged([]); + } + + // Every close path routes through here: callers keep this mounted, so a draft + // left behind would resurface (and upload) on the next project. + function close() { + if (busy) return; + reset(); + onOpenChange(false); + } async function commitCreate() { const trimmed = name.trim(); - if (!trimmed) return; + if (!trimmed || busy) return; + setBusy(true); + // Sidebar callers keep this mounted across routes, so unmounting alone + // cannot tell whether the user has moved on during a slow upload. + const origin = currentRoute(); try { const project = await createChatProject(trimmed); + // Upload before closing so the Sources panel lists them on first fetch. + await uploadStagedSources(project.id, staged); + if (!mounted.current) return; + const stayedOnRoute = currentRoute() === origin; onOpenChange(false); - setName(""); + reset(); + if (onCreated) { + await onCreated(project, { stayedOnRoute }); + return; + } + if (!stayedOnRoute) return; const runtime = useChatRuntimeStore.getState(); runtime.setActiveThreadId(null); runtime.setActiveProjectId(project.id); @@ -45,6 +106,8 @@ export function NewProjectDialog({ toast.error("Failed to create project", { description: err instanceof Error ? err.message : undefined, }); + } finally { + setBusy(false); } } @@ -52,43 +115,59 @@ export function NewProjectDialog({ { - if (!next) setName(""); - onOpenChange(next); + if (next) { + onOpenChange(true); + return; + } + close(); }} > - + - New project + {title} - setName(e.target.value)} - onKeyDown={(e) => { - if (e.key === "Enter") { - e.preventDefault(); - void commitCreate(); - } - }} - autoFocus={true} - maxLength={120} - placeholder="Project name" - aria-label="Project name" - className="focus-visible:border-input focus-visible:ring-0" + {/* Name field: folder glyph in its own cell, divided from the input. */} +
+ + + +
+ -
diff --git a/studio/frontend/src/features/chat/projects-page.tsx b/studio/frontend/src/features/chat/projects-page.tsx index e20e517787..192c4e2331 100644 --- a/studio/frontend/src/features/chat/projects-page.tsx +++ b/studio/frontend/src/features/chat/projects-page.tsx @@ -34,7 +34,6 @@ import { isTauri } from "@/lib/api-base"; import { isDownloadCancelled, pickNativeChatImport } from "@/lib/native-files"; import { toast } from "@/lib/toast"; import { - createChatProject, deleteChatProject, renameChatProject, useChatProjects, @@ -42,6 +41,7 @@ import { usePinnedProjectsStore, type ProjectRecord, } from "@/features/chat"; +import { NewProjectDialog } from "./components/new-project-dialog"; import { Delete02Icon, Download01Icon, @@ -124,7 +124,6 @@ export function ProjectsPage() { ); const [creating, setCreating] = useState(false); - const [nameDraft, setNameDraft] = useState(""); const [renaming, setRenaming] = useState(null); const [renameDraft, setRenameDraft] = useState(""); const [deleting, setDeleting] = useState(null); @@ -258,21 +257,6 @@ export function ProjectsPage() { navigate({ to: "/chat", search: { project: projectId } }); } - async function commitCreate() { - const name = nameDraft.trim(); - if (!name) return; - try { - const project = await createChatProject(name); - setCreating(false); - setNameDraft(""); - openProject(project.id); - } catch (err) { - toast.error("Failed to create project", { - description: err instanceof Error ? err.message : undefined, - }); - } - } - async function commitRename() { const target = renaming; const name = renameDraft.trim(); @@ -469,14 +453,7 @@ export function ProjectsPage() { - + @@ -511,10 +488,7 @@ export function ProjectsPage() { - - -
-
+ {/* Create project (name + drag-and-drop sources) */} + {/* Rename project */} = 1024 && unit < units.length - 1) { + value /= 1024; + unit += 1; + } + const shown = + value >= 10 || unit === 0 + ? String(Math.round(value)) + : value.toFixed(1).replace(/\.0$/, ""); + return `${shown} ${units[unit]}`; +} + +const ACCEPTED_EXTS = new Set( + RAG_UPLOAD_ACCEPT.split(",").map((ext) => ext.trim().toLowerCase()), +); + +// `accept` only filters the picker, so a drop can carry anything. A folder +// arrives as an extension-less entry, which this rejects along with the types +// the backend would 400 on. +function isSupported(file: File): boolean { + const dot = file.name.lastIndexOf("."); + if (dot <= 0) return false; + return ACCEPTED_EXTS.has(file.name.slice(dot).toLowerCase()); +} + +/** Merge a selection into the staged list. Returns the names it would not take, + * so the caller can say so once instead of dropping them silently. */ +function addStagedSources( + staged: StagedSource[], + incoming: FileList | File[], +): { next: StagedSource[]; unsupported: string[]; duplicates: string[] } { + const seen = new Set(staged.map((entry) => fileSignature(entry.file))); + const next = [...staged]; + const unsupported: string[] = []; + const duplicates: string[] = []; + for (const file of Array.from(incoming)) { + if (!isSupported(file)) { + unsupported.push(file.name); + continue; + } + const signature = fileSignature(file); + if (seen.has(signature)) { + duplicates.push(file.name); + continue; + } + seen.add(signature); + next.push({ + id: `staged_${Math.random().toString(36).slice(2)}`, + file, + }); + } + return { next, unsupported, duplicates }; +} + +// Projects created with staged files, so the landing can open on Sources. +const projectsWithPendingSources = new Set(); + +function markProjectSourcesPending(projectId: string): void { + projectsWithPendingSources.add(projectId); +} + +/** Whether this project was just created with staged sources. Read-only, so it + * is safe in a render pass that React may replay. */ +export function hasProjectSourcesPending(projectId: string): boolean { + return projectsWithPendingSources.has(projectId); +} + +/** Drop the marker once the landing has committed. */ +export function consumeProjectSourcesPending(projectId: string): void { + projectsWithPendingSources.delete(projectId); +} + +/** Upload staged files to a new project. Indexing runs in the background; a + * per-file failure toasts and never blocks project creation. */ +export async function uploadStagedSources( + projectId: string, + staged: StagedSource[], +): Promise { + if (staged.length === 0) return; + invalidateProjectSources(projectId); + markProjectSourcesPending(projectId); + const { ocr, caption } = resolveVisionOverrides(); + const documentIds = new Set(); + const merged: string[] = []; + for (const { file } of staged) { + try { + const result = await uploadProjectDocument(projectId, file, ocr, caption); + // Same bytes under another name: the backend hashes content, so this is + // the document already uploaded. Say so rather than imply a new source. + if (documentIds.has(result.documentId)) merged.push(file.name); + else documentIds.add(result.documentId); + } catch (error) { + toast.error(`Couldn't upload ${file.name}`, { + description: error instanceof Error ? error.message : String(error), + }); + } + } + if (merged.length > 0) { + toast.info( + merged.length === 1 + ? `${merged[0]} matched a file already added` + : `${merged.length} files matched files already added`, + { description: "Identical contents are stored once." }, + ); + } + invalidateProjectSources(projectId); +} + +/** Create-project drop area: stages files until the project exists. */ +export function ProjectSourceDropzone({ + staged, + onChange, + disabled = false, +}: { + staged: StagedSource[]; + onChange: (next: StagedSource[]) => void; + disabled?: boolean; +}) { + const inputRef = useRef(null); + // Count enter/leave pairs: children fire dragleave on the parent. + const dragDepth = useRef(0); + const [dragging, setDragging] = useState(false); + + const addFiles = useCallback( + (files: FileList | File[]) => { + const { next, unsupported, duplicates } = addStagedSources(staged, files); + if (next.length !== staged.length) onChange(next); + if (unsupported.length > 0) { + toast.info( + unsupported.length === 1 + ? `Can't add ${unsupported[0]}` + : `Can't add ${unsupported.length} files`, + { description: `Supported types: ${RAG_UPLOAD_ACCEPT}` }, + ); + } + // Name, size and mtime can in principle match for two different files, so + // never drop one without saying so. + if (duplicates.length > 0) { + toast.info( + duplicates.length === 1 + ? `${duplicates[0]} is already added` + : `${duplicates.length} files were already added`, + ); + } + }, + [staged, onChange], + ); + + const endDrag = useCallback(() => { + dragDepth.current = 0; + setDragging(false); + }, []); + + return ( +
+

Sources

+ {/* Panel is the drop target; the inner button owns the click so staged + rows can carry their own remove buttons. */} +
{ + e.preventDefault(); + if (disabled) return; + dragDepth.current += 1; + setDragging(true); + }} + onDragOver={(e) => { + e.preventDefault(); + if (disabled) return; + e.dataTransfer.dropEffect = "copy"; + }} + onDragLeave={() => { + dragDepth.current = Math.max(0, dragDepth.current - 1); + if (dragDepth.current === 0) setDragging(false); + }} + onDrop={(e) => { + e.preventDefault(); + if (disabled) return; + endDrag(); + addFiles(Array.from(e.dataTransfer.files ?? [])); + }} + className={cn( + "rounded-[22px] border border-border transition-colors dark:border-white/10", + dragging && "border-primary/60 bg-primary/5", + disabled && "opacity-60", + )} + > + { + const files = Array.from(e.target.files ?? []); + e.target.value = ""; + addFiles(files); + }} + /> + {staged.length === 0 ? ( + + ) : ( +
+
    + {staged.map((entry) => ( +
  • + + + {entry.file.name} + + + {formatSize(entry.file.size)} + + +
  • + ))} +
+ +
+ )} +
+
+ ); +} diff --git a/studio/frontend/src/features/rag/components/use-rag-documents.ts b/studio/frontend/src/features/rag/components/use-rag-documents.ts index 8d6433d8c3..bdab7b0518 100644 --- a/studio/frontend/src/features/rag/components/use-rag-documents.ts +++ b/studio/frontend/src/features/rag/components/use-rag-documents.ts @@ -1,13 +1,8 @@ // 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 { useCallback, useEffect, useRef, useState } from "react"; -import { - CHAT_RAG_CAPTION_KEY, - CHAT_RAG_OCR_KEY, - useChatRuntimeStore, -} from "@/features/chat"; import { toast } from "@/lib/toast"; +import { useCallback, useEffect, useRef, useState } from "react"; import { deleteDocument, getJob, @@ -17,6 +12,7 @@ import { uploadThreadDocument, } from "../api/rag-api"; import type { DocumentStatus, RagDocument } from "../types/rag"; +import { resolveVisionOverrides } from "./vision-overrides"; export interface TrackedDocument extends RagDocument { progress?: number | null; @@ -263,18 +259,7 @@ export function useRagDocuments( tempId: string, ) => { try { - // Send vision-pass overrides only after the user has explicitly set them; - // otherwise backend env defaults own the ingest policy. - const state = useChatRuntimeStore.getState(); - const hasLocal = (key: string) => - typeof window !== "undefined" && - window.localStorage.getItem(key) !== null; - const ocr = hasLocal(CHAT_RAG_OCR_KEY) - ? state.ragOcrScanned - : undefined; - const caption = hasLocal(CHAT_RAG_CAPTION_KEY) - ? state.ragCaptionFigures - : undefined; + const { ocr, caption } = resolveVisionOverrides(); const result = activeScope.type === "kb" ? await uploadKnowledgeBaseDocument( diff --git a/studio/frontend/src/features/rag/components/vision-overrides.ts b/studio/frontend/src/features/rag/components/vision-overrides.ts new file mode 100644 index 0000000000..674484970b --- /dev/null +++ b/studio/frontend/src/features/rag/components/vision-overrides.ts @@ -0,0 +1,35 @@ +// 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 { + CHAT_RAG_CAPTION_KEY, + CHAT_RAG_OCR_KEY, + useChatRuntimeStore, +} from "@/features/chat"; + +function hasLocal(key: string): boolean { + if (typeof window === "undefined") return false; + try { + return window.localStorage.getItem(key) !== null; + } catch { + // Storage can be blocked outright (sandboxed context). These overrides are + // optional, so fall back to the backend defaults rather than failing the + // upload that asked for them. + return false; + } +} + +/** Ingest-time vision-pass overrides, sent only once the user has set them; + * otherwise backend env defaults own the policy. Shared by every upload path. */ +export function resolveVisionOverrides(): { + ocr: boolean | undefined; + caption: boolean | undefined; +} { + const state = useChatRuntimeStore.getState(); + return { + ocr: hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined, + caption: hasLocal(CHAT_RAG_CAPTION_KEY) + ? state.ragCaptionFigures + : undefined, + }; +} From 671d6dbf6902f0355496e5d42219daad5d2f06fb Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Sun, 26 Jul 2026 00:01:22 -0700 Subject: [PATCH 045/161] Settings: match dialog fills to the app shell surfaces (#7457) * Settings: match dialog fills to the app shell surfaces Tabs use the sidebar fill and the content pane uses the page fill, so both track the active palette in light and dark. * Pair the tab column fill with the sidebar foreground Custom themes set --foreground but not --sidebar, so search result rows could land white on white. Track the sidebar token instead. --- studio/frontend/src/features/settings/settings-dialog.tsx | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index e3ef476470..0ba59f7095 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -267,7 +267,9 @@ export function SettingsDialog() { {/* Keep tab content from expanding the dialog grid. */}
-
+ ); +} + +// Quote only values with shell metacharacters, e.g. a local path with spaces. +function quoteShellArg(value: string, windows: boolean): string { + if (SAFE_SHELL_ARG_PATTERN.test(value)) { + return value; + } + return windows ? `'${psSingle(value)}'` : `'${shSingle(value)}'`; +} + +function SubagentSection({ + agent, + baseCommand, + modelArgs, +}: { + agent: AgentDetails; + baseCommand: string; + modelArgs: string; +}) { + const t = useT(); + // modelArgs is empty when attaching to a resident model that has no id to name. + const command = `${baseCommand} --as-subagent${modelArgs ? ` ${modelArgs}` : ""}`; + const prompt = + agent.id === "opencode" + ? t("settings.agents.subagent.opencodePrompt") + : t("settings.agents.subagent.defaultPrompt"); + const commandCopy = useCopyButton(command); + const promptCopy = useCopyButton(prompt); + + if (!SUBAGENT_AGENT_IDS.has(agent.id)) { + return null; + } + + return ( +
+
+ + {t("settings.agents.subagent.title")} + +

+ {t("settings.agents.subagent.description", { agent: agent.name })} +

+
+ +
+
+ + {t("settings.agents.subagent.setupCommand")} + + +
+ + {command} + +
+ +
+
+ + {t("settings.agents.subagent.usagePrompt", { agent: agent.name })} + + +
+ + {prompt} + +
); } @@ -233,18 +610,142 @@ function CommandBlock({ command }: { command: string }) { export function AgentsTab() { const t = useT(); const serverUrl = usePlatformStore((s) => s.serverUrl); + const hfToken = useHfTokenStore((s) => s.token); const deviceType = usePlatformStore((s) => s.deviceType); - const [info, setInfo] = useState(null); - - const origin = typeof window !== "undefined" ? window.location.origin : ""; - const localDetection = canUseLocalAgentDetection(serverUrl ?? origin); - // The remote snippet runs on the client, so use the client platform, not deviceType. // Anchor the match: a bare includes("win") would also match "darwin". const [isWindowsClient] = useState(() => { const p = getClientPlatform(); return p.startsWith("win") || p.includes("windows"); }); + const origin = typeof window !== "undefined" ? window.location.origin : ""; + // Browser commands target the viewed origin; a desktop window origin is a Tauri URL + // the CLI cannot reach, so use the backend URL from /api/health (getApiBase until it + // lands). The command then runs wherever that CLI is: a loopback base is this Studio's + // own host, so deviceType decides, and it reports wsl where the browser would claim + // Windows; any other base is reached from the viewer's machine, so only the client + // platform describes that shell. + const studioBase = isTauri ? (serverUrl ?? getApiBase()) : origin; + const isWindowsShell = isLoopbackBase(studioBase) + ? deviceType === "windows" + : isWindowsClient; + const localDetection = canUseLocalAgentDetection(serverUrl ?? origin); + const [agents, setAgents] = useState( + SUPPORTED_AGENTS.map((agent) => agent.id), + ); + const [selectedAgent, setSelectedAgent] = useState(FALLBACK_AGENT.id); + const agentSelectionChanged = useRef(false); + const [detectedAgents, setDetectedAgents] = useState>(new Set()); + const [loaded, setLoaded] = useState(false); + const [models, setModels] = useState([EXAMPLE_MODEL_REPO]); + const [cachedLoadIds, setCachedLoadIds] = useState>( + {}, + ); + // Display names for scanned models, keyed by the path that identifies them. + const [modelLabels, setModelLabels] = useState>({}); + // The model /api/inference/status reports as resident, so the command attaches to it + // rather than remapping to another cached copy. + const [activeStatusModel, setActiveStatusModel] = useState( + null, + ); + // Set only for a native-grant GGUF, which is resident but has no id to pass. + const [attachOnlyModel, setAttachOnlyModel] = useState(null); + const [knownVariants, setKnownVariants] = useState>({ + [EXAMPLE_MODEL_REPO]: EXAMPLE_MODEL_VARIANT, + }); + const [selectedModel, setSelectedModel] = useState(EXAMPLE_MODEL_REPO); + const modelSelectionChanged = useRef(false); + // The model status last reported, for the discovery scan to preserve. + const activeModelRef = useRef(null); + // Only the newest status request may apply; a slow earlier one must not win. + const statusSeq = useRef(0); + // A quant picked by hand, scoped to its repo: polling and refetches must not + // overwrite it, but it must not follow the selection onto a different repo. + const chosenVariant = useRef<{ model: string; variant: string } | null>(null); + const [modelSearch, setModelSearch] = useState(""); + const [modelPickerOpen, setModelPickerOpen] = useState(false); + const [variants, setVariants] = useState([]); + const [defaultVariant, setDefaultVariant] = useState(null); + const [selectedVariant, setSelectedVariant] = useState( + EXAMPLE_MODEL_VARIANT, + ); + const [variantsLoading, setVariantsLoading] = useState(true); + const [variantsFailed, setVariantsFailed] = useState(false); + + const labelFor = (model: string) => modelLabels[model] ?? model; + const matchingModels = useMemo(() => { + const tokens = modelSearch + .trim() + .toLowerCase() + .split(SEARCH_TOKEN_PATTERN) + .filter(Boolean); + const matches = + tokens.length === 0 + ? models + : models.filter((model) => { + // Search both, so a scanned model is findable by name and by path. + const haystack = + `${model} ${modelLabels[model] ?? ""}`.toLowerCase(); + return tokens.every((token) => haystack.includes(token)); + }); + + if (tokens.length === 0 && matches.includes(selectedModel)) { + return [ + selectedModel, + ...matches.filter((model) => model !== selectedModel), + ]; + } + return matches; + }, [modelLabels, modelSearch, models, selectedModel]); + + const visibleModels = matchingModels.slice(0, MODEL_RESULT_LIMIT); + const preferredVariant = knownVariants[selectedModel] ?? null; + const selectedAgentDetails = detailsFor(selectedAgent); + // A GGUF outside the active cache does not resolve by repo id, so name its + // snapshot path; `unsloth start` now also matches a path by the basename + // /v1/models advertises for it. The resident model is exempt: it already + // loaded by id, and cached-gguf keeps the largest copy across caches, whose + // snapshot could switch cache or quant under it. + const cachedLoadId = + selectedModel === activeStatusModel + ? null + : (cachedLoadIds[selectedModel] ?? + cachedLoadIds[selectedModel.toLowerCase()] ?? + null); + const modelId = cachedLoadId ?? selectedModel; + const suffixVariant = isHuggingFaceRepo(modelId); + const commandModel = + selectedVariant && suffixVariant + ? `${modelId}:${selectedVariant}` + : modelId; + const commandModelArg = quoteShellArg(commandModel, isWindowsShell); + // A bare `unsloth start` attaches to whatever is loaded, which is the only way + // to reach a native-grant GGUF: naming it would switch the server to another model. + const attachOnly = selectedModel === attachOnlyModel; + const modelArgs = attachOnly + ? "" + : selectedVariant && !suffixVariant + ? `--model ${commandModelArg} --gguf-variant ${quoteShellArg(selectedVariant, isWindowsShell)}` + : `--model ${commandModelArg}`; + // No key is passed: the CLI caches an explicit one per base, overwriting a working + // saved key. Omitting it replays the saved key; the remote section covers first setup. + const commandOs = isWindowsShell ? "windows" : "unix"; + const commandBase = buildAgentCommand( + studioBase, + null, + commandOs, + selectedAgent, + ); + const command = attachOnly ? commandBase : `${commandBase} ${modelArgs}`; + // The fixed examples below target the same Studio, not a bare 127.0.0.1:8888. + const example = (agentId: string, flags: string) => + `${buildAgentCommand(studioBase, null, commandOs, agentId)} ${flags}`; + const { + copied, + copy: handleCopy, + reset: resetCopied, + } = useCopyButton(command); + const remoteCommand = isWindowsClient ? REMOTE_CMD_WINDOWS : REMOTE_CMD_UNIX; useEffect(() => { void fetchDeviceType({ force: true }); @@ -252,56 +753,324 @@ export function AgentsTab() { // A remote backend's PATH says nothing about the machine running the copied command. useEffect(() => { - if (!localDetection) return; + if (!localDetection) { + return; + } let cancelled = false; loadCodingAgents() .then((next) => { - if (!cancelled) setInfo(next); + if (cancelled) { + return; + } + if (next.agents.length > 0) { + setAgents(next.agents); + setSelectedAgent((current) => { + if (agentSelectionChanged.current) { + return current; + } + const detected = next.detected.find((agent) => + next.agents.includes(agent), + ); + return ( + detected ?? + (next.agents.includes(current) ? current : next.agents[0]) + ); + }); + } + setDetectedAgents(new Set(next.detected)); }) .catch(() => { // Best-effort; the tab still works without PATH detection. + }) + .finally(() => { + if (!cancelled) { + setLoaded(true); + } }); return () => { cancelled = true; }; }, [localDetection]); - // Derive visibility from localDetection instead of clearing info in the effect. - const visibleInfo = localDetection ? info : null; - const detected = new Set(visibleInfo?.detected ?? []); - const remoteCommand = isWindowsClient ? REMOTE_CMD_WINDOWS : REMOTE_CMD_UNIX; + useEffect(() => { + let cancelled = false; + Promise.all([ + listModels().catch(() => null), + listCachedGguf().catch(() => []), + listLocalModels().catch(() => null), + ]) + .then(([info, cachedGgufs, local]) => { + if (cancelled) { + return; + } + const localEntries = localGgufEntries(local?.models ?? []); + const discovered = discoverGgufModels(info?.models ?? [], [ + ...cachedGgufs.map((cached) => cached.repo_id), + ...localEntries.map((entry) => entry.id), + ]); + // Keep the snapshot load_id for --model while listing the model by repo id. + const loadIds: Record = {}; + for (const cached of cachedGgufs) { + if (cached.load_id && cached.load_id !== cached.repo_id) { + // Key both spellings: the merge above keeps whichever casing arrived + // first, which may not be this endpoint's. + loadIds[cached.repo_id] = cached.load_id; + loadIds[cached.repo_id.toLowerCase()] = cached.load_id; + } + } + const labels: Record = {}; + for (const entry of localEntries) { + if (entry.label !== entry.id) { + labels[entry.id] = entry.label; + } + } + // Status is applied on its own schedule now, so keep whatever model it has + // already adopted rather than dropping it when this slower scan lands. + setModels(() => { + const active = activeModelRef.current; + return active && !discovered.models.includes(active) + ? [active, ...discovered.models] + : discovered.models; + }); + setCachedLoadIds(loadIds); + setModelLabels(labels); + setKnownVariants((current) => ({ + ...current, + ...discovered.variants, + })); + }) + .catch(() => { + // The example model keeps the builder useful if discovery fails. + }); + return () => { + cancelled = true; + }; + }, []); - // `codex` needs a GGUF model (unsloth_cli's _require_gguf_for_codex exits otherwise), so flag - // its row instead of offering a failing command. Same three signals the API usage panel uses. - const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); - const activeNativePathToken = useChatRuntimeStore( - (s) => s.activeNativePathToken, + // List the resident model and follow it, unless the user picked one explicitly. + const adoptActiveModel = useCallback( + (active: { model: string; variant: string | null }) => { + setModels((current) => + current.includes(active.model) ? current : [active.model, ...current], + ); + if (active.variant) { + setKnownVariants((current) => ({ + ...current, + [active.model]: active.variant as string, + })); + } + if (!modelSelectionChanged.current) { + setSelectedModel(active.model); + if (chosenVariant.current?.model !== active.model) { + setSelectedVariant(active.variant); + } + } + }, + [], ); - const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); - const isGguf = - activeGgufVariant != null || - activeNativePathToken != null || - ggufContextLength != null; - // Build from the reachable base: a bare `unsloth start` only probes 127.0.0.1:8888, but the - // desktop falls back across 8888-8908 and Studio may be remote. The browser must use its own - // origin, since /api/health reports the backend's localhost (the user's, behind a tunnel); - // the desktop has no window origin and falls back to getApiBase() while serverUrl loads. - // No --api-key: the CLI caches an explicit key per base, so a placeholder would overwrite a - // working saved one. Omitting it replays the saved key; the remote section covers first setup. - const commandBase = isTauri ? (serverUrl ?? getApiBase()) : origin; - // The command runs wherever the CLI is. For a loopback base that is this Studio's - // own host, so use deviceType, which reports wsl where the browser would claim - // Windows and emit $env: syntax bash rejects. A remote base is reached from the - // viewer's machine instead, so only the client platform describes that shell. - const commandOs = - (isLoopbackBase(commandBase) ? deviceType === "windows" : isWindowsClient) - ? "windows" - : "unix"; - const agentCommand = (agentId: string) => - buildAgentCommand(commandBase, null, commandOs, agentId); - const example = (agentId: string, flags: string) => - `${agentCommand(agentId)} ${flags}`; + // A native-grant label only stands for whatever was resident at the time, so once + // that model is replaced the label cannot name anything and has to go, even when + // it was picked by hand: leaving it selected would emit it as --model. + const retireAttachOnly = useCallback((label: string, replacement: string) => { + setModels((current) => current.filter((model) => model !== label)); + setSelectedModel((current) => { + if (current !== label) { + return current; + } + // Drop the quant in the same transition: it belonged to the label, and an + // explicit pick stops adoptActiveModel from correcting it afterwards. + chosenVariant.current = null; + setSelectedVariant(null); + return replacement; + }); + }, []); + + // The resident GGUF went away (unloaded, or replaced by a transformer model). + // Following it means letting go too, or the command would name a stale model and + // switch the shared server back. A native-grant label is not even loadable, so it + // leaves the list entirely. An explicit pick still wins. + const dropActiveModel = useCallback( + (attachOnly: string | null, wasActive: string | null) => { + if (attachOnly) { + setModels((current) => current.filter((model) => model !== attachOnly)); + // Even a deliberate pick has to go: the label stood for a withheld path, so + // naming it would emit --model