diff --git a/README.md b/README.md index 26a578656c..7046a2af7c 100644 --- a/README.md +++ b/README.md @@ -109,18 +109,19 @@ You can use the same Docker image as Unsloth Studio. For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth).
To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel). -## ✨ Free Notebooks +## πŸ“’ Free Notebooks -Train for free with our notebooks. Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Add dataset, run, then deploy your trained model. +Train for free with our notebooks. You can use our new [free Unsloth Studio notebook](https://colab.research.google.com/github/unslothai/unsloth/blob/main/studio/Unsloth_Studio_Colab.ipynb) to run and train models for free in a web UI. +Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Add dataset, run, then deploy your trained model. | Model | Free Notebooks | Performance | Memory use | |-----------|---------|--------|----------| +| **Gemma 4 (E2B)** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma4_(E2B)-Vision.ipynb) | 1.5x faster | 50% less | | **Qwen3.5 (4B)** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_5_(4B)_Vision.ipynb) | 1.5x faster | 60% less | | **gpt-oss (20B)** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-Fine-tuning.ipynb) | 2x faster | 70% less | | **Qwen3.5 GSPO** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_5_(4B)_Vision_GRPO.ipynb) | 2x faster | 70% less | | **gpt-oss (20B): GRPO** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb) | 2x faster | 80% less | | **Qwen3: Advanced GRPO** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Qwen3_(4B)-GRPO.ipynb) | 2x faster | 70% less | -| **Gemma 3 (4B) Vision** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Gemma3_(4B)-Vision.ipynb) | 1.7x faster | 60% less | | **embeddinggemma (300M)** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/EmbeddingGemma_(300M).ipynb) | 2x faster | 20% less | | **Mistral Ministral 3 (3B)** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Ministral_3_VL_(3B)_Vision.ipynb) | 1.5x faster | 60% less | | **Llama 3.1 (8B) Alpaca** | [▢️ Start for free](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.1_(8B)-Alpaca.ipynb) | 2x faster | 70% less | @@ -132,6 +133,7 @@ Train for free with our notebooks. Read our [guide](https://unsloth.ai/docs/get- - See detailed documentation for Unsloth [here](https://unsloth.ai/docs) ## πŸ¦₯ Unsloth News +- **Gemma 4**: Run and train Google’s new models directly in Unsloth Studio! [Blog](https://unsloth.ai/docs/models/gemma-4) - **Introducing Unsloth Studio**: our new web UI for running and training LLMs. [Blog](https://unsloth.ai/docs/new/studio) - **Qwen3.5** - 0.8B, 2B, 4B, 9B, 27B, 35-A3B, 112B-A10B are now supported. [Guide + notebooks](https://unsloth.ai/docs/models/qwen3.5/fine-tune) - Train **MoE LLMs 12x faster** with 35% less VRAM - DeepSeek, GLM, Qwen and gpt-oss. [Blog](https://unsloth.ai/docs/new/faster-moe) diff --git a/install_gemma4_mlx.sh b/install_gemma4_mlx.sh index e06339e204..b653af9154 100755 --- a/install_gemma4_mlx.sh +++ b/install_gemma4_mlx.sh @@ -132,29 +132,17 @@ step "install" "installing mlx, mlx-lm..." uv pip install --python "$_VENV_PY" -q mlx mlx-lm 2>/dev/null substep "done" -TRANSFORMERS_WHL="transformers-5.5.0-py3-none-any.whl" -TRANSFORMERS_GH="git+https://github.com/huggingface/transformers.git@v5.5-release" - step "install" "installing transformers>=5.5.0..." -if uv pip install --python "$_VENV_PY" -q "$TRANSFORMERS_GH" 2>/dev/null; then - substep "installed from huggingface/transformers v5.5-release" -elif uv pip install --python "$_VENV_PY" -q "transformers>=5.5.0" 2>/dev/null; then +if uv pip install --python "$_VENV_PY" -q "transformers>=5.5.0" 2>/dev/null; then substep "installed from PyPI" else - substep "not on PyPI, trying unsloth branch..." - _whl_tmp=$(mktemp -d)/"${TRANSFORMERS_WHL}" - if curl -fsSL "${REPO_URL}/${TRANSFORMERS_WHL}" -o "$_whl_tmp" 2>/dev/null && \ - uv pip install --python "$_VENV_PY" -q "$_whl_tmp" 2>/dev/null; then - substep "installed from branch ${BRANCH}" - elif [ -f "./${TRANSFORMERS_WHL}" ]; then - substep "using local ./${TRANSFORMERS_WHL}" - uv pip install --python "$_VENV_PY" -q "./${TRANSFORMERS_WHL}" + substep "PyPI install failed (Python <3.10?), trying GitHub..." + if uv pip install --python "$_VENV_PY" -q "git+https://github.com/huggingface/transformers.git@v5.5-release" 2>/dev/null; then + substep "installed from huggingface/transformers v5.5-release" else - rm -f "$_whl_tmp" 2>/dev/null - step "install" "skipping transformers β€” could not find >=5.5.0" "$C_WARN" - substep "tried: huggingface/transformers v5.5-release, PyPI, branch ${BRANCH}, local ./${TRANSFORMERS_WHL}" + step "warning" "could not install transformers>=5.5.0" "$C_WARN" + substep "tried: PyPI, huggingface/transformers v5.5-release" fi - rm -f "$_whl_tmp" 2>/dev/null fi # ── Find mlx-lm models directory ───────────────────────────── diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index f520041add..1b4b5381e8 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -93,6 +93,14 @@ "min_p": 0.0, "repetition_penalty": 1.0 }, + "gemma-4": { + "temperature": 1.0, + "top_p": 0.95, + "top_k": 64, + "min_p": 0.0, + "repetition_penalty": 1.0, + "presence_penalty": 0.0 + }, "gemma-3n": { "temperature": 1.0, "top_p": 0.95, @@ -366,7 +374,7 @@ "qwen2.5-coder", "qwen2.5-vl", "qwen2.5-omni", "qwen2.5-math", "qwen2.5", "qwen2-vl", "qwen2", "qwq", - "gemma-3n", "gemma-3", "medgemma", "gemma-2", + "gemma-4", "gemma-3n", "gemma-3", "medgemma", "gemma-2", "llama-4", "llama-3.3", "llama-3.2", "llama-3.1", "llama-3", "phi-4", "phi-3", "mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral", diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml new file mode 100644 index 0000000000..c80506d9f5 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B-it.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-26B-A4B-it +# Also applies to: google/gemma-4-26B-A4B-it, unsloth/gemma-4-26B-A4B-it-GGUF + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml new file mode 100644 index 0000000000..9e579be503 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-26B-A4B.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-26B-A4B (base/pretrained) +# Also applies to: google/gemma-4-26B-A4B + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml new file mode 100644 index 0000000000..cec4ea95e1 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B-it.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-31B-it +# Also applies to: google/gemma-4-31B-it, unsloth/gemma-4-31B-it-GGUF + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml new file mode 100644 index 0000000000..717cdd5e63 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-31B.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-31B (base/pretrained) +# Also applies to: google/gemma-4-31B + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml new file mode 100644 index 0000000000..43e3d78a23 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B-it.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-E2B-it +# Also applies to: google/gemma-4-E2B-it, unsloth/gemma-4-E2B-it-GGUF + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml new file mode 100644 index 0000000000..bd86cef751 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E2B.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-E2B (base/pretrained) +# Also applies to: google/gemma-4-E2B + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml new file mode 100644 index 0000000000..a8ef51836b --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B-it.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-E4B-it +# Also applies to: google/gemma-4-E4B-it, unsloth/gemma-4-E4B-it-GGUF + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml new file mode 100644 index 0000000000..740cc99df5 --- /dev/null +++ b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-4-E4B.yaml @@ -0,0 +1,47 @@ +# Model defaults for unsloth/gemma-4-E4B (base/pretrained) +# Also applies to: google/gemma-4-E4B + +training: + trust_remote_code: false + max_seq_length: 2048 + num_epochs: 0 + learning_rate: 2e-4 + batch_size: 2 + gradient_accumulation_steps: 4 + warmup_steps: 5 + max_steps: 30 + save_steps: 30 + weight_decay: 0.001 + random_seed: 3407 + packing: false + train_on_completions: true + gradient_checkpointing: "unsloth" + optim: "adamw_8bit" + lr_scheduler_type: "linear" + +lora: + lora_r: 8 + lora_alpha: 8 + lora_dropout: 0.0 + target_modules: + - "all-linear" + use_rslora: false + use_loftq: false + finetune_vision_layers: true + finetune_language_layers: true + finetune_attention_modules: true + finetune_mlp_modules: true + +logging: + enable_wandb: false + wandb_project: "llm-finetuning" + enable_tensorboard: false + tensorboard_dir: "runs" + log_frequency: 10 + +inference: + trust_remote_code: false + temperature: 1.0 + top_p: 0.95 + top_k: 64 + min_p: 0.0 diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 550358ae61..7fff36aefd 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -167,12 +167,7 @@ def _validate_recipe_runtime_support( recipe: dict[str, Any], model_providers: list[Any], ) -> None: - if not _recipe_has_llm_columns(recipe): - raise ValueError( - "Recipe Studio currently requires at least one AI generation step." - ) - - if not model_providers: + if _recipe_has_llm_columns(recipe) and not model_providers: raise ValueError("Add a Provider connection block before running this recipe.") @@ -266,6 +261,21 @@ def create_data_designer( model_providers = build_model_providers(recipe) _validate_recipe_runtime_support(recipe, model_providers) + # DataDesigner requires at least one model provider in its registry even + # when the pipeline contains no LLM columns. Supply a lightweight stub + # so sampler/expression-only recipes can run without a real provider. + if not model_providers: + from data_designer.config.models import ModelProvider + + model_providers = [ + ModelProvider( + name = "_unused", + endpoint = "http://localhost", + provider_type = "openai", + api_key = None, + ) + ] + return DataDesigner( artifact_path = artifact_path, model_providers = model_providers, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 44c700bf3d..c84ac640df 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -22,6 +22,7 @@ import threading import time from pathlib import Path from typing import Generator, Optional +from urllib.parse import urlparse import httpx @@ -52,8 +53,14 @@ _REPROMPT_MAX_CHARS = 2000 _SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$") _SHARD_RE = re.compile(r"^(.*)-\d{5}-of-\d{5}\.gguf$") -# Model size extraction (shared with routes/inference.py) -from utils.models import extract_model_size_b as _extract_model_size_b + +# Model size extraction β€” lazy import to avoid pulling in transformers +# at module level. See PR description for the full explanation. +def _extract_model_size_b(model_id: str): + from utils.models import extract_model_size_b + + return extract_model_size_b(model_id) + # ── Pre-compiled patterns for tool XML stripping ───────────── _TOOL_CLOSED_PATS = [ @@ -102,6 +109,7 @@ class LlamaCppBackend: self._supports_tools: bool = False self._cache_type_kv: Optional[str] = None self._reasoning_default: bool = True + self._speculative_type: Optional[str] = None # KV-cache estimation fields (populated by _read_gguf_metadata) self._n_layers: Optional[int] = None self._n_kv_heads: Optional[int] = None @@ -191,6 +199,10 @@ class LlamaCppBackend: def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv + @property + def speculative_type(self) -> Optional[str]: + return self._speculative_type + # ── Binary discovery ────────────────────────────────────────── @staticmethod @@ -1048,6 +1060,7 @@ class LlamaCppBackend: n_ctx: int = 4096, chat_template_override: Optional[str] = None, cache_type_kv: Optional[str] = None, + speculative_type: Optional[str] = None, n_threads: Optional[int] = None, n_gpu_layers: Optional[int] = None, # Accepted for caller compat, unused ) -> bool: @@ -1308,6 +1321,46 @@ class LlamaCppBackend: else: self._cache_type_kv = None + # Speculative decoding (n-gram self-speculation, zero VRAM cost) + # ngram-mod: ~16 MB shared hash pool, constant memory/complexity, + # variable draft lengths. Helps most when the model repeats + # existing text (code refactoring, summarization, reasoning). + # For general chat with low repetition, overhead is ~5 ms. + # + # Benchmarks from llama.cpp PRs #18471, #19164: + # Scenario | Without | With | Speedup + # gpt-oss-120b code refactor | 181 t/s | 446 t/s | 2.5x + # Qwen3-235B offloaded | 12 t/s | 21 t/s | 1.8x + # gpt-oss-120b repeat (92% accept)| 181 t/s | 814 t/s | 4.5x + # + # Params from llama.cpp docs (docs/speculative.md): + # --spec-ngram-size-n 24 (small n not recommended) + # --draft-min 48 --draft-max 64 (MoEs need long drafts; + # dense models can reduce these) + # ref: https://github.com/ggml-org/llama.cpp/blob/master/docs/speculative.md + # ref: https://github.com/ggml-org/llama.cpp/pull/19164 + # ref: https://github.com/ggml-org/llama.cpp/pull/18471 + _valid_spec_types = {"ngram-simple", "ngram-mod"} + if speculative_type and speculative_type in _valid_spec_types: + if not is_vision: # spec decoding disabled for vision models + cmd.extend(["--spec-type", speculative_type]) + if speculative_type == "ngram-mod": + cmd.extend( + [ + "--spec-ngram-size-n", + "24", + "--draft-min", + "48", + "--draft-max", + "64", + ] + ) + self._speculative_type = speculative_type + else: + self._speculative_type = None + else: + self._speculative_type = None + # Apply custom chat template override if provided if chat_template_override: import tempfile @@ -1546,6 +1599,7 @@ class LlamaCppBackend: self._reasoning_always_on = False self._supports_tools = False self._cache_type_kv = None + self._speculative_type = None self._n_layers = None self._n_kv_heads = None self._n_heads = None @@ -2264,7 +2318,7 @@ class LlamaCppBackend: Agentic loop: let the model call tools, execute them, and continue. Yields dicts with: - {"type": "status", "text": "Searching: ..."} -- tool status updates + {"type": "status", "text": "Searching: ..."/"Reading: ..."} -- tool status updates {"type": "content", "text": "token"} -- streamed content tokens (cumulative) {"type": "reasoning", "text": "token"} -- streamed reasoning tokens (cumulative) """ @@ -2831,7 +2885,18 @@ class LlamaCppBackend: arguments = raw_args if tool_name == "web_search": - status_text = f"Searching: {arguments.get('query', '')}" + _ws_url = (arguments.get("url") or "").strip() + if _ws_url: + _parsed = urlparse(_ws_url) + if _parsed.scheme in ("http", "https") and _parsed.hostname: + _ws_host = _parsed.hostname + if _ws_host.startswith("www."): + _ws_host = _ws_host[4:] + status_text = f"Reading: {_ws_host}" + else: + status_text = "Reading page..." + else: + status_text = f"Searching: {arguments.get('query', '')}" elif tool_name == "python": preview = ( (arguments.get("code") or "").strip().split("\n")[0][:60] diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index b23372b766..87cc933d4b 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -14,6 +14,8 @@ import os os.environ["UNSLOTH_IS_PRESENT"] = "1" import random +import re +import shlex import ssl import subprocess import sys @@ -27,14 +29,239 @@ logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes +# Pre-import modules used in _sandbox_preexec at module level so that +# the preexec_fn closure does not trigger the import machinery in the +# forked child (which can deadlock in multi-threaded servers). +_libc = None +if sys.platform == "linux": + try: + import ctypes + import ctypes.util + + _libc_name = ctypes.util.find_library("c") + if _libc_name: + _libc = ctypes.CDLL(_libc_name, use_errno = True) + except (OSError, AttributeError): + pass + +_resource = None +if sys.platform != "win32": + try: + import resource as _resource + except ImportError: + pass + # Strict raster-image allowlist for sandbox file serving. # No .svg (XSS risk via embedded scripts), no .html, no .pdf. _IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"}) _MAX_OUTPUT_CHARS = 8000 # truncate long output -_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"} +_BLOCKED_COMMANDS_COMMON = frozenset( + { + "rm", + "sudo", + "su", + "dd", + "chmod", + "chown", + "mkfs", + "shutdown", + "reboot", + "passwd", + "mount", + "umount", + "fdisk", + "kill", + "killall", + "pkill", + } +) +_BLOCKED_COMMANDS_WIN = frozenset( + { + "rmdir", + "takeown", + "icacls", + "runas", + "powershell", + "pwsh", + } +) +_BLOCKED_COMMANDS = ( + _BLOCKED_COMMANDS_COMMON | _BLOCKED_COMMANDS_WIN + if sys.platform == "win32" + else _BLOCKED_COMMANDS_COMMON +) + + +def _find_blocked_commands(command: str) -> set[str]: + """Detect blocked commands using shlex tokenization and regex scanning. + + Catches: full paths (/usr/bin/sudo), quoted strings ("sudo"), + split-quotes (su""do), backslash escapes (\\rm), and command-position + words after ;, |, &&, $(). + """ + blocked = set() + + # 1. shlex tokenization (handles quotes, escapes, concatenation) + try: + tokens = ( + shlex.split(command) + if sys.platform != "win32" + else shlex.split(command, posix = False) + ) + except ValueError: + tokens = command.split() + + for token in tokens: + base = os.path.basename(token).lower() + # Strip common Windows executable extensions so that + # runas.exe, shutdown.bat, etc. match the blocklist. + stem, ext = os.path.splitext(base) + if ext in {".exe", ".com", ".bat", ".cmd"}: + base = stem + if base in _BLOCKED_COMMANDS: + blocked.add(base) + + # 2. Regex: catch blocked words at shell command boundaries + # (semicolons, pipes, &&, ||, backticks, $(), <(), subshells, newlines) + # Uses a single combined pattern for all blocked words. + # Handles optional Unix path prefix (/usr/bin/) and Windows drive + # letter prefix (C:\Windows\...\). + lowered = command.lower() + if _BLOCKED_COMMANDS: + words_alt = "|".join(re.escape(w) for w in sorted(_BLOCKED_COMMANDS)) + pattern = ( + rf"(?:^|[;&|`\n(]\s*|[$]\(\s*|<\(\s*)" + rf"(?:[\w./\\-]*/|[a-zA-Z]:[/\\][\w./\\-]*)?" + rf"({words_alt})(?:\.(?:exe|com|bat|cmd))?\b" + ) + blocked.update(re.findall(pattern, lowered)) + + # 3. Check for nested shell invocations (bash -c 'sudo whoami', + # bash -lc '...', bash --login -c '...', cmd /c '...'). + # When a -c or /c flag is found, look backwards for a shell name + # (skipping intermediate flags like --login, -l, -x) and recursively + # scan the nested command string. + _SHELLS = {"bash", "sh", "zsh", "dash", "ksh", "csh", "tcsh", "fish"} + _SHELLS_WIN = {"cmd", "cmd.exe"} + for i, token in enumerate(tokens): + tok_lower = token.lower() + # Match -c exactly, or combined flags ending in c (e.g. -lc, -xc) + is_unix_c = tok_lower == "-c" or ( + tok_lower.startswith("-") + and tok_lower.endswith("c") + and not tok_lower.startswith("--") + ) + is_win_c = tok_lower == "/c" + if not (is_unix_c or is_win_c) or i < 1 or i + 1 >= len(tokens): + continue + # Look backwards past any flags to find the shell binary. + # On Unix, flags start with - (skip those). On Windows, flags + # start with / but so do absolute paths, so only skip short + # single-char /X flags (not /bin/bash style paths). + for j in range(i - 1, -1, -1): + prev = tokens[j] + if prev.startswith("-"): + continue # skip Unix flags like --login, -l + if is_win_c and prev.startswith("/") and len(prev) <= 3: + continue # skip Windows flags like /s, /q (not /bin/bash) + prev_base = os.path.basename(prev).lower() + if is_unix_c and prev_base in _SHELLS: + blocked |= _find_blocked_commands(tokens[i + 1]) + elif is_win_c and prev_base in _SHELLS_WIN: + blocked |= _find_blocked_commands(tokens[i + 1]) + break # stop at first non-flag token + + return blocked + + +def _build_safe_env(workdir: str) -> dict[str, str]: + """Build a minimal, credential-free environment for sandboxed subprocesses. + + Strips HF_TOKEN, WANDB_API_KEY, AWS_*, GH_TOKEN, LD_PRELOAD, DYLD_*, etc. + Preserves the active Python interpreter and virtualenv directories in PATH + so that pip, uv, and packages installed in the Studio runtime remain + accessible. + """ + # Start with the directory containing the running Python interpreter + # so that subprocess calls to 'python', 'pip', etc. resolve to the + # same environment the Studio server is running in. + exe_dir = os.path.dirname(sys.executable) + path_entries = [exe_dir] if exe_dir else [] + + # If a virtualenv is active, include its bin/Scripts directory. + venv = os.environ.get("VIRTUAL_ENV") + if venv: + venv_bin = os.path.join(venv, "Scripts" if sys.platform == "win32" else "bin") + if venv_bin not in path_entries: + path_entries.append(venv_bin) + + if sys.platform == "win32": + sysroot = os.environ.get("SystemRoot", r"C:\Windows") + path_entries.extend([os.path.join(sysroot, "System32"), sysroot]) + else: + path_entries.extend(["/usr/local/bin", "/usr/bin", "/bin"]) + + # Deduplicate while preserving order + deduped = list(dict.fromkeys(p for p in path_entries if p)) + + env = { + "PATH": os.pathsep.join(deduped), + "HOME": workdir, + "TMPDIR": workdir, + "LANG": os.environ.get("LANG", "C.UTF-8"), + "TERM": "dumb", + "PYTHONIOENCODING": "utf-8", + } + if venv: + env["VIRTUAL_ENV"] = venv + # Windows needs SystemRoot for Python/subprocess to work + if sys.platform == "win32": + env["SystemRoot"] = os.environ.get("SystemRoot", r"C:\Windows") + return env + + +def _sandbox_preexec(): + """Pre-exec hook: drop privilege escalation ability and set resource limits. + + On Linux, applies PR_SET_NO_NEW_PRIVS so sudo/su/pkexec fail at the + kernel level. On Linux and macOS, sets RLIMIT_FSIZE. + No-op on Windows (use creationflags instead). + + Note: RLIMIT_NPROC is intentionally NOT set because Linux enforces it + per real UID, not per process tree, so it would starve the Studio + server and other sessions sharing the same user account. + + All modules and handles are resolved at import time (module level) so + this function does not trigger Python imports in the forked child, + avoiding potential deadlocks in multi-threaded servers. + """ + if _libc is not None: + try: + # PR_SET_NO_NEW_PRIVS = 38, arg2 = 1 (enable) + _libc.prctl(38, 1, 0, 0, 0) + except (OSError, AttributeError): + pass # Not available (container, old kernel, etc.) + + if _resource is not None: + try: + # Limit file size to 100MB (prevents disk filling) + _resource.setrlimit( + _resource.RLIMIT_FSIZE, (100 * 1024 * 1024, 100 * 1024 * 1024) + ) + except (ValueError, OSError): + pass + + +def _get_shell_cmd(command: str) -> list[str]: + """Return the platform-appropriate shell invocation for a command string.""" + if sys.platform == "win32": + return ["cmd", "/c", command] + return ["bash", "-c", command] + # Per-session working directories so each chat thread gets its own sandbox. -# Falls back to a shared ~/studio_sandbox/ for API callers without a session_id. +# Falls back to a shared ~/studio_sandbox/_default for API callers without a +# session_id. _workdirs: dict[str, str] = {} @@ -55,7 +282,7 @@ def _get_workdir(session_id: str | None = None) -> str: if not os.path.realpath(workdir).startswith(os.path.realpath(sandbox_root)): workdir = os.path.join(sandbox_root, "_invalid") else: - workdir = sandbox_root + workdir = os.path.join(sandbox_root, "_default") os.makedirs(workdir, exist_ok = True) _workdirs[key] = workdir return _workdirs[key] @@ -428,6 +655,7 @@ def _check_signal_escape_patterns(code: str): signal_tampering = [] exception_catching = [] + shell_escapes = [] warnings = [] def _ast_name_matches(node, names): @@ -445,10 +673,84 @@ def _check_signal_escape_patterns(code: str): return full_name in names return False + # Dangerous os/subprocess functions that can execute shell commands + _SHELL_EXEC_FUNCS = frozenset( + { + "os.system", + "os.popen", + "os.popen2", + "os.popen3", + "os.popen4", + "os.execl", + "os.execle", + "os.execlp", + "os.execlpe", + "os.execv", + "os.execve", + "os.execvp", + "os.execvpe", + "os.spawnl", + "os.spawnle", + "os.spawnlp", + "os.spawnlpe", + "os.spawnv", + "os.spawnve", + "os.spawnvp", + "os.spawnvpe", + "os.posix_spawn", + "os.posix_spawnp", + "subprocess.run", + "subprocess.call", + "subprocess.check_call", + "subprocess.check_output", + "subprocess.Popen", + "subprocess.getoutput", + "subprocess.getstatusoutput", + } + ) + + def _extract_string_from_node(node): + """Extract a plain string value from an AST node, if it is a constant.""" + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + return None + + def _extract_strings_from_list(node): + """Extract string elements from an AST List or Tuple node.""" + if isinstance(node, (ast.List, ast.Tuple)): + parts = [] + for elt in node.elts: + s = _extract_string_from_node(elt) + if s is not None: + parts.append(s) + return parts + return [] + + # Keyword argument names that carry command content (as opposed to + # control flags like check=True, text=True, capture_output=True). + _CMD_KWARGS = frozenset({"args", "command", "executable", "path", "file"}) + + def _check_args_for_blocked(args_nodes): + """Check if any call arguments contain blocked commands.""" + found = set() + for arg in args_nodes: + s = _extract_string_from_node(arg) + if s is not None: + found |= _find_blocked_commands(s) + strs = _extract_strings_from_list(arg) + for s in strs: + found |= _find_blocked_commands(s) + return found + class SignalEscapeVisitor(ast.NodeVisitor): def __init__(self): self.imports_signal = False self.signal_aliases = {"signal"} + self.os_aliases = {"os"} + self.subprocess_aliases = {"subprocess"} + # Maps bare function names to their fully-qualified form + # for from-import tracking (e.g. "system" -> "os.system") + self.shell_exec_aliases: dict[str, str] = {} self.loop_depth = 0 def visit_Import(self, node): @@ -457,6 +759,10 @@ def _check_signal_escape_patterns(code: str): self.imports_signal = True if alias.asname: self.signal_aliases.add(alias.asname) + elif alias.name == "os": + self.os_aliases.add(alias.asname or "os") + elif alias.name == "subprocess": + self.subprocess_aliases.add(alias.asname or "subprocess") self.generic_visit(node) def visit_ImportFrom(self, node): @@ -474,6 +780,16 @@ def _check_signal_escape_patterns(code: str): "alarm", ): self.signal_aliases.add(alias.asname or alias.name) + elif node.module in ("os", "subprocess"): + if node.module == "os": + self.os_aliases.add("os") + else: + self.subprocess_aliases.add("subprocess") + # Track from-imports of dangerous functions + for alias in node.names: + fq = f"{node.module}.{alias.name}" + if fq in _SHELL_EXEC_FUNCS: + self.shell_exec_aliases[alias.asname or alias.name] = fq self.generic_visit(node) def visit_While(self, node): @@ -538,6 +854,111 @@ def _check_signal_escape_patterns(code: str): "description": "Modifies signal mask (may block SIGALRM)", } ) + + # --- Shell escape detection --- + # Resolve the fully qualified function name for os.*/subprocess.* + shell_func = None + if isinstance(func, ast.Attribute): + if isinstance(func.value, ast.Name): + if func.value.id in self.os_aliases: + shell_func = f"os.{func.attr}" + elif func.value.id in self.subprocess_aliases: + shell_func = f"subprocess.{func.attr}" + elif isinstance(func, ast.Name): + # Check from-import aliases: from os import system; system(...) + shell_func = self.shell_exec_aliases.get(func.id) + + if shell_func and shell_func in _SHELL_EXEC_FUNCS: + # Expand **kwargs dicts to inspect their keys + expanded_kwargs: dict[str, ast.AST] = {} + has_opaque_kwargs = False + for kw in node.keywords: + if kw.arg is not None: + expanded_kwargs[kw.arg] = kw.value + elif isinstance(kw.value, ast.Dict): + for k, v in zip(kw.value.keys, kw.value.values): + key = _extract_string_from_node(k) if k else None + if key is not None: + expanded_kwargs[key] = v + else: + has_opaque_kwargs = True + + cmd_kw_values = [ + v for k, v in expanded_kwargs.items() if k in _CMD_KWARGS + ] + all_call_args = list(node.args) + cmd_kw_values + blocked_in_args = _check_args_for_blocked(all_call_args) + + if has_opaque_kwargs: + # Can't inspect dynamic **kwargs -- flag as unsafe + shell_escapes.append( + { + "type": "shell_escape_dynamic", + "line": node.lineno, + "description": ( + f"{shell_func}() called with dynamic **kwargs" + ), + } + ) + elif blocked_in_args: + shell_escapes.append( + { + "type": "shell_escape", + "line": node.lineno, + "description": ( + f"{shell_func}() invokes blocked command(s): " + f"{', '.join(sorted(blocked_in_args))}" + ), + } + ) + else: + # Only flag dynamic args for functions that interpret + # strings as shell commands, or when shell= might be + # enabled. Treat any non-literal-False shell= value + # as potentially True (conservative). + _STRING_SHELL_FUNCS = frozenset( + { + "os.system", + "os.popen", + "os.popen2", + "os.popen3", + "os.popen4", + "subprocess.getoutput", + "subprocess.getstatusoutput", + } + ) + shell_node = expanded_kwargs.get("shell") + shell_safe = shell_node is None or ( + isinstance(shell_node, ast.Constant) + and shell_node.value is False + ) + if shell_func in _STRING_SHELL_FUNCS or not shell_safe: + + def _is_safe_literal(n): + if _extract_string_from_node(n) is not None: + return True + if isinstance(n, (ast.List, ast.Tuple)): + return all( + _extract_string_from_node(e) is not None + for e in n.elts + ) + return False + + has_non_literal = any( + not _is_safe_literal(a) for a in all_call_args + ) + if has_non_literal: + shell_escapes.append( + { + "type": "shell_escape_dynamic", + "line": node.lineno, + "description": ( + f"{shell_func}() called with non-literal " + f"shell command (potential shell escape)" + ), + } + ) + self.generic_visit(node) def visit_ExceptHandler(self, node): @@ -553,7 +974,12 @@ def _check_signal_escape_patterns(code: str): } ) elif isinstance(node.type, ast.Name): - if node.type.id in ("TimeoutError", "BaseException", "Exception"): + # Only flag BaseException and TimeoutError, NOT Exception. + # except Exception does not catch SystemExit or + # KeyboardInterrupt, so it cannot suppress timeout + # enforcement. Flagging Exception causes false positives + # on normal error-handling patterns. + if node.type.id in ("TimeoutError", "BaseException"): exception_catching.append( { "type": f"catches_{node.type.id}_in_loop", @@ -564,7 +990,7 @@ def _check_signal_escape_patterns(code: str): elif isinstance(node.type, ast.Tuple): for elt in node.type.elts: if isinstance(elt, ast.Name): - if elt.id in ("TimeoutError", "BaseException", "Exception"): + if elt.id in ("TimeoutError", "BaseException"): exception_catching.append( { "type": f"catches_{elt.id}_in_loop", @@ -580,10 +1006,15 @@ def _check_signal_escape_patterns(code: str): if visitor.imports_signal and not signal_tampering: warnings.append("Code imports 'signal' module - review manually for safety") - is_safe = len(signal_tampering) == 0 and len(exception_catching) == 0 + is_safe = ( + len(signal_tampering) == 0 + and len(exception_catching) == 0 + and len(shell_escapes) == 0 + ) return is_safe, { "signal_tampering": signal_tampering, "exception_catching": exception_catching, + "shell_escapes": shell_escapes, "warnings": warnings, } @@ -604,10 +1035,18 @@ def _check_code_safety(code: str) -> str | None: reasons = [ item.get("description", "") for item in info.get("signal_tampering", []) ] - return ( - f"Error: unsafe code detected ({'; '.join(reasons)}). " - f"Please remove signal manipulation from your code." - ) + shell_reasons = [ + item.get("description", "") for item in info.get("shell_escapes", []) + ] + exception_reasons = [ + item.get("description", "") for item in info.get("exception_catching", []) + ] + all_reasons = [r for r in reasons + shell_reasons + exception_reasons if r] + if all_reasons: + return ( + f"Error: unsafe code detected ({'; '.join(all_reasons)}). " + f"Please remove unsafe patterns from your code." + ) return None @@ -662,13 +1101,20 @@ def _python_exec( with os.fdopen(fd, "w") as f: f.write(code) - proc = subprocess.Popen( - [sys.executable, tmp_path], + safe_env = _build_safe_env(workdir) + popen_kwargs = dict( stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, cwd = workdir, + env = safe_env, ) + if sys.platform != "win32": + popen_kwargs["preexec_fn"] = _sandbox_preexec + else: + popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW + + proc = subprocess.Popen([sys.executable, tmp_path], **popen_kwargs) # Spawn cancel watcher if we have a cancel event if cancel_event is not None: @@ -734,21 +1180,27 @@ def _bash_exec( if not command or not command.strip(): return "No command provided." - # Block dangerous commands - tokens = set(command.lower().split()) - blocked = tokens & _BASH_BLOCKED_WORDS + # Block dangerous commands (shlex + regex based) + blocked = _find_blocked_commands(command) if blocked: return f"Blocked command(s) for safety: {', '.join(sorted(blocked))}" try: workdir = _get_workdir(session_id) - proc = subprocess.Popen( - ["bash", "-c", command], + safe_env = _build_safe_env(workdir) + popen_kwargs = dict( stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, cwd = workdir, + env = safe_env, ) + if sys.platform != "win32": + popen_kwargs["preexec_fn"] = _sandbox_preexec + else: + popen_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW + + proc = subprocess.Popen(_get_shell_cmd(command), **popen_kwargs) if cancel_event is not None: watcher = threading.Thread( diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index e2513f43de..7f7291a56d 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -145,6 +145,8 @@ def _get_hf_download_state( blobs_dirs: list[Path] = [] if model_names: + from utils.paths import resolve_cached_repo_id_case + for name in model_names: if not name: continue @@ -154,6 +156,7 @@ def _get_hf_download_state( # relative paths, and Windows paths. if name.startswith(("/", ".", "~")) or "\\" in name: continue + name = resolve_cached_repo_id_case(name) # HF cache dir format: models--org--name (slashes -> --) cache_dir_name = "models--" + name.replace("/", "--") blobs_dir = cache / cache_dir_name / "blobs" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3094df4169..cf08ecbc12 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -48,6 +48,10 @@ class LoadRequest(BaseModel): 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. Not supported for GGUF models.", ) + speculative_type: Optional[str] = Field( + None, + description = "Speculative decoding mode for GGUF models (e.g. 'ngram-simple', 'ngram-mod'). Ignored for non-GGUF and vision models.", + ) class UnloadRequest(BaseModel): @@ -163,6 +167,10 @@ class LoadResponse(BaseModel): None, description = "Jinja2 chat template string (from GGUF metadata or tokenizer)", ) + speculative_type: Optional[str] = Field( + None, + description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled", + ) class UnloadResponse(BaseModel): @@ -225,6 +233,10 @@ class InferenceStatusResponse(BaseModel): None, description = "Model's native context length from GGUF metadata (not capped by VRAM)", ) + speculative_type: Optional[str] = Field( + None, + description = "Active speculative decoding mode (e.g. 'ngram-simple', 'ngram-mod'), or None if disabled", + ) # ===================================================================== diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ced24c1d5f..30ff7da49c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -179,6 +179,7 @@ async def load_model( supports_reasoning = llama_backend.supports_reasoning, reasoning_always_on = llama_backend.reasoning_always_on, chat_template = llama_backend.chat_template, + speculative_type = llama_backend.speculative_type, ) else: if ( @@ -263,6 +264,7 @@ async def load_model( n_ctx = request.max_seq_length, chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, + speculative_type = request.speculative_type, ) else: # Local mode: llama-server loads via -m @@ -275,6 +277,7 @@ async def load_model( n_ctx = request.max_seq_length, chat_template_override = request.chat_template_override, cache_type_kv = request.cache_type_kv, + speculative_type = request.speculative_type, ) if not success: @@ -317,6 +320,7 @@ async def load_model( supports_tools = llama_backend.supports_tools, cache_type_kv = llama_backend.cache_type_kv, chat_template = llama_backend.chat_template, + speculative_type = llama_backend.speculative_type, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -652,6 +656,7 @@ async def get_status( context_length = llama_backend.context_length, max_context_length = llama_backend.max_context_length, native_context_length = llama_backend.native_context_length, + speculative_type = llama_backend.speculative_type, ) # Otherwise, report Unsloth backend status diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 445cf0e7f4..1e31a91e26 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -49,8 +49,10 @@ try: ) from core.inference import get_inference_backend from utils.paths import ( + is_local_path, outputs_root, exports_root, + resolve_cached_repo_id_case, resolve_output_dir, resolve_export_dir, ) @@ -77,8 +79,10 @@ except ImportError: ) from core.inference import get_inference_backend from utils.paths import ( + is_local_path, outputs_root, exports_root, + resolve_cached_repo_id_case, resolve_output_dir, resolve_export_dir, ) @@ -597,10 +601,15 @@ async def get_model_config( This endpoint wraps the backend load_model_defaults function. """ try: - from utils.models.model_config import is_local_path - if not is_local_path(model_name): - model_name = model_name.lower() + resolved = resolve_cached_repo_id_case(model_name) + if resolved != model_name: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved, + model_name, + ) + model_name = resolved logger.info(f"Getting model config for: {model_name}") from utils.models.model_config import detect_audio_type diff --git a/studio/backend/tests/test_cache_case_resolution.py b/studio/backend/tests/test_cache_case_resolution.py new file mode 100644 index 0000000000..60963b4f7d --- /dev/null +++ b/studio/backend/tests/test_cache_case_resolution.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from pathlib import Path +import sys +import types + +# Keep this test runnable in lightweight environments where optional logging +# deps are not installed. +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +from utils.paths.path_utils import ( + resolve_cached_repo_id_case, + get_cache_case_resolution_stats, + reset_cache_case_resolution_state, +) +import utils.paths.path_utils as path_utils + + +def _mk_cache_repo(cache_root: Path, repo_id: str) -> Path: + repo_dir = cache_root / f"models--{repo_id.replace('/', '--')}" + repo_dir.mkdir(parents = True, exist_ok = True) + return repo_dir + + +def test_resolve_cached_repo_id_case_exact_hit(tmp_path, monkeypatch): + reset_cache_case_resolution_state() + _mk_cache_repo(tmp_path, "Org/Model") + monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path) + + resolved = resolve_cached_repo_id_case("Org/Model") + + assert resolved == "Org/Model" + stats = get_cache_case_resolution_stats() + assert stats["calls"] == 1 + assert stats["exact_hits"] == 1 + assert stats["variant_hits"] == 0 + + +def test_resolve_cached_repo_id_case_variant_hit(tmp_path, monkeypatch): + reset_cache_case_resolution_state() + _mk_cache_repo(tmp_path, "Org/Model") + monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path) + + resolved = resolve_cached_repo_id_case("org/model") + + assert resolved == "Org/Model" + stats = get_cache_case_resolution_stats() + assert stats["variant_hits"] == 1 + assert stats["tie_breaks"] == 0 + + +def test_resolve_cached_repo_id_case_tie_break_deterministic(tmp_path, monkeypatch): + reset_cache_case_resolution_state() + _mk_cache_repo(tmp_path, "Org/Model") + _mk_cache_repo(tmp_path, "org/model") + monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path) + + resolved = resolve_cached_repo_id_case("oRg/mOdEl") + + # Deterministic rule: lexical sort of candidate repo ids. + assert resolved == "Org/Model" + stats = get_cache_case_resolution_stats() + assert stats["variant_hits"] == 1 + assert stats["tie_breaks"] == 1 + + +def test_resolve_cached_repo_id_case_no_cache_fallback(tmp_path, monkeypatch): + reset_cache_case_resolution_state() + monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path) + + resolved = resolve_cached_repo_id_case("Org/Missing") + + assert resolved == "Org/Missing" + stats = get_cache_case_resolution_stats() + assert stats["fallbacks"] == 1 + assert stats["variant_hits"] == 0 + assert stats["exact_hits"] == 0 + + +def test_resolve_cached_repo_id_case_memoization(tmp_path, monkeypatch): + reset_cache_case_resolution_state() + _mk_cache_repo(tmp_path, "Org/Model") + monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path) + + first = resolve_cached_repo_id_case("org/model") + second = resolve_cached_repo_id_case("org/model") + + assert first == "Org/Model" + assert second == "Org/Model" + stats = get_cache_case_resolution_stats() + assert stats["calls"] == 2 + assert stats["variant_hits"] == 1 + assert stats["memo_hits"] == 1 + + +def test_resolve_cached_repo_id_case_late_cache_population(tmp_path, monkeypatch): + """Regression guard: memoized fallback should not hide a later cache variant.""" + reset_cache_case_resolution_state() + monkeypatch.setattr(path_utils, "_hf_hub_cache_dir", lambda: tmp_path) + + first = resolve_cached_repo_id_case("org/model") + assert first == "org/model" + + # Simulate cache being populated after first miss (e.g. another code path/download). + _mk_cache_repo(tmp_path, "Org/Model") + + second = resolve_cached_repo_id_case("org/model") + + # Desired behavior: second lookup should pick up the now-existing variant. + assert second == "Org/Model" diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py new file mode 100644 index 0000000000..3481e29948 --- /dev/null +++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py @@ -0,0 +1,81 @@ +# 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 asyncio +import sys +import types + +# Keep this test runnable in lightweight environments where optional logging +# deps are not installed. +if "structlog" not in sys.modules: + + class _DummyLogger: + def __getattr__(self, _name): + return lambda *args, **kwargs: None + + sys.modules["structlog"] = types.SimpleNamespace( + BoundLogger = _DummyLogger, + get_logger = lambda *args, **kwargs: _DummyLogger(), + ) + +import routes.models as models_route +import utils.models.model_config as model_config_module + + +def test_get_model_config_resolves_cached_case_before_model_checks(monkeypatch): + calls: dict[str, str] = {} + + class _DummyModelConfig: + is_lora = False + base_model = None + + def _record_load(model_name): + calls["load_model_defaults"] = model_name + return {} + + def _record_vision(model_name, hf_token = None): + calls["is_vision_model"] = model_name + return False + + def _record_embedding(model_name, hf_token = None): + calls["is_embedding_model"] = model_name + return False + + def _record_audio(model_name, hf_token = None): + calls["detect_audio_type"] = model_name + return None + + def _record_from_identifier(cls, model_name): + calls["from_identifier"] = model_name + return _DummyModelConfig() + + monkeypatch.setattr(models_route, "is_local_path", lambda _: False) + monkeypatch.setattr( + models_route, "resolve_cached_repo_id_case", lambda _: "Org/Model" + ) + monkeypatch.setattr(models_route, "load_model_defaults", _record_load) + monkeypatch.setattr(models_route, "is_vision_model", _record_vision) + monkeypatch.setattr(models_route, "is_embedding_model", _record_embedding) + monkeypatch.setattr(model_config_module, "detect_audio_type", _record_audio) + monkeypatch.setattr( + models_route.ModelConfig, + "from_identifier", + classmethod(_record_from_identifier), + ) + monkeypatch.setattr(models_route, "_get_max_position_embeddings", lambda _: 4096) + monkeypatch.setattr(models_route, "_get_model_size_bytes", lambda *_args, **_kw: 0) + + result = asyncio.run( + models_route.get_model_config( + model_name = "org/model", + hf_token = None, + current_subject = "test-subject", + ) + ) + + assert result.model_name == "Org/Model" + assert calls["load_model_defaults"] == "Org/Model" + assert calls["is_vision_model"] == "Org/Model" + assert calls["is_embedding_model"] == "Org/Model" + assert calls["detect_audio_type"] == "Org/Model" + assert calls["from_identifier"] == "Org/Model" diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index df1058abf6..f97ea993eb 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -5,13 +5,14 @@ Model and LoRA configuration handling """ -from transformers import AutoConfig from dataclasses import dataclass from typing import Optional, Dict, Any from utils.paths import ( normalize_path, is_local_path, is_model_cached, + get_cache_path, + resolve_cached_repo_id_case, outputs_root, exports_root, resolve_output_dir, @@ -158,6 +159,38 @@ MODEL_NAME_MAPPING = { "unsloth/gemma-3n-E4B-unsloth-bnb-4bit", "google/gemma-3n-E4B", ], + "unsloth_gemma-4-31B-it.yaml": [ + "unsloth/gemma-4-31B-it", + "google/gemma-4-31B-it", + ], + "unsloth_gemma-4-26B-A4B-it.yaml": [ + "unsloth/gemma-4-26B-A4B-it", + "google/gemma-4-26B-A4B-it", + ], + "unsloth_gemma-4-E2B-it.yaml": [ + "unsloth/gemma-4-E2B-it", + "google/gemma-4-E2B-it", + ], + "unsloth_gemma-4-E4B-it.yaml": [ + "unsloth/gemma-4-E4B-it", + "google/gemma-4-E4B-it", + ], + "unsloth_gemma-4-31B.yaml": [ + "unsloth/gemma-4-31B", + "google/gemma-4-31B", + ], + "unsloth_gemma-4-26B-A4B.yaml": [ + "unsloth/gemma-4-26B-A4B", + "google/gemma-4-26B-A4B", + ], + "unsloth_gemma-4-E2B.yaml": [ + "unsloth/gemma-4-E2B", + "google/gemma-4-E2B", + ], + "unsloth_gemma-4-E4B.yaml": [ + "unsloth/gemma-4-E4B", + "google/gemma-4-E4B", + ], "unsloth_gpt-oss-20b.yaml": [ "openai/gpt-oss-20b", "unsloth/gpt-oss-20b-unsloth-bnb-4bit", @@ -422,6 +455,7 @@ def load_model_config( """ Load model config with optional authentication control. """ + from transformers import AutoConfig if token: # Explicit token provided - use it @@ -711,12 +745,8 @@ def _detect_audio_from_tokenizer( # 1) Check local HF cache first (works for gated/offline models) try: - from huggingface_hub.constants import HF_HUB_CACHE - - cache_dir = Path(HF_HUB_CACHE) - repo_dir_name = f"models--{model_name.replace('/', '--')}" - repo_dir = cache_dir / repo_dir_name - if repo_dir.exists(): + repo_dir = get_cache_path(model_name) + if repo_dir is not None and repo_dir.exists(): snapshots_dir = repo_dir / "snapshots" if snapshots_dir.exists(): for snapshot in snapshots_dir.iterdir(): @@ -1627,11 +1657,18 @@ class ModelConfig: identifier = f"unsloth/{identifier}" path = identifier - # Enforce lowercase for remote Hugging Face identifiers to prevent cache duplication - # Hugging Face Hub APIs are case-insensitive remotely, but case-sensitive locally (repo_folder_name). + # Preserve requested casing, but if a case-variant already exists in local HF cache, + # reuse that exact repo_id spelling to avoid one-time re-downloads after #2592. if not is_local: - identifier = identifier.lower() - path = path.lower() + resolved_identifier = resolve_cached_repo_id_case(identifier) + if resolved_identifier != identifier: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved_identifier, + identifier, + ) + identifier = resolved_identifier + path = resolved_identifier # Auto-detect GGUF models (check before LoRA/vision detection) if is_local: @@ -1852,6 +1889,12 @@ class ModelConfig: identifier = f"unsloth/{identifier}" path = identifier + if not is_local: + resolved_identifier = resolve_cached_repo_id_case(identifier) + if resolved_identifier != identifier: + identifier = resolved_identifier + path = resolved_identifier + # --- Logic for Base Model and Vision Detection --- base_model = None is_vision = False diff --git a/studio/backend/utils/paths/__init__.py b/studio/backend/utils/paths/__init__.py index 44a7c8e287..11709ae56e 100644 --- a/studio/backend/utils/paths/__init__.py +++ b/studio/backend/utils/paths/__init__.py @@ -5,7 +5,15 @@ Path utilities for model and dataset handling """ -from .path_utils import normalize_path, is_local_path, is_model_cached, get_cache_path +from .path_utils import ( + normalize_path, + is_local_path, + is_model_cached, + get_cache_path, + resolve_cached_repo_id_case, + get_cache_case_resolution_stats, + reset_cache_case_resolution_state, +) from .storage_roots import ( studio_root, assets_root, @@ -40,6 +48,9 @@ __all__ = [ "is_local_path", "is_model_cached", "get_cache_path", + "resolve_cached_repo_id_case", + "get_cache_case_resolution_stats", + "reset_cache_case_resolution_state", "studio_root", "assets_root", "datasets_root", diff --git a/studio/backend/utils/paths/path_utils.py b/studio/backend/utils/paths/path_utils.py index b38db18286..9ef9a2dd92 100644 --- a/studio/backend/utils/paths/path_utils.py +++ b/studio/backend/utils/paths/path_utils.py @@ -14,6 +14,20 @@ from loggers import get_logger logger = get_logger(__name__) +# Per-process cache to avoid repeated cache-dir scans for the same identifier. +_CACHE_CASE_RESOLUTION_MEMO: dict[str, str] = {} + +# Lightweight instrumentation counters for operational visibility. +_CACHE_CASE_RESOLUTION_STATS: dict[str, int] = { + "calls": 0, + "memo_hits": 0, + "exact_hits": 0, + "variant_hits": 0, + "tie_breaks": 0, + "fallbacks": 0, + "errors": 0, +} + def _is_wsl() -> bool: """Detect if we are running inside WSL (Windows Subsystem for Linux).""" @@ -94,8 +108,9 @@ def is_local_path(path: str) -> bool: def get_cache_path(model_name: str) -> Optional[Path]: """Get HuggingFace cache path for a model if it exists.""" - cache_dir = Path.home() / ".cache" / "huggingface" / "hub" - model_cache_name = model_name.replace("/", "--") + cache_dir = _hf_hub_cache_dir() + resolved_name = resolve_cached_repo_id_case(model_name) + model_cache_name = resolved_name.replace("/", "--") model_cache_path = cache_dir / f"models--{model_cache_name}" return model_cache_path if model_cache_path.exists() else None @@ -113,3 +128,102 @@ def is_model_cached(model_name: str) -> bool: return True return False + + +def _hf_hub_cache_dir() -> Path: + """Return HF cache root honoring HF_HUB_CACHE when available.""" + try: + from huggingface_hub.constants import HF_HUB_CACHE + + return Path(HF_HUB_CACHE) + except Exception as exc: + logger.debug( + "Could not read huggingface_hub HF_HUB_CACHE, using default hub path: %s", + exc, + ) + return Path.home() / ".cache" / "huggingface" / "hub" + + +def resolve_cached_repo_id_case(model_name: str, use_memo: bool = True) -> str: + """Resolve repo_id to the exact casing already present in local HF cache. + + Policy: prefer the requested/canonical repo_id, but if a case-variant already + exists in local HF cache, reuse that exact cached spelling. This avoids + duplicate downloads while preserving user intent whenever possible. + """ + _CACHE_CASE_RESOLUTION_STATS["calls"] += 1 + + if not model_name or "/" not in model_name: + _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1 + return model_name + + cache_dir = _hf_hub_cache_dir() + if not cache_dir.exists(): + _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1 + return model_name + + expected_dir = f"models--{model_name.replace('/', '--')}" + + # Always check the exact-case path first so a newly-appeared exact match + # wins over any previously memoized variant. + exact_path = cache_dir / expected_dir + if exact_path.is_dir(): + if use_memo: + _CACHE_CASE_RESOLUTION_MEMO[model_name] = model_name + _CACHE_CASE_RESOLUTION_STATS["exact_hits"] += 1 + return model_name + + # Validate memoized entries still exist on disk before returning them. + # This prevents stale results when cache dirs are deleted/recreated. + if use_memo: + cached = _CACHE_CASE_RESOLUTION_MEMO.get(model_name) + if cached is not None: + cached_path = cache_dir / f"models--{cached.replace('/', '--')}" + if cached_path.is_dir(): + _CACHE_CASE_RESOLUTION_STATS["memo_hits"] += 1 + return cached + # Stale entry -- drop it and re-scan below. + _CACHE_CASE_RESOLUTION_MEMO.pop(model_name, None) + + expected_lower = expected_dir.lower() + try: + candidates: list[str] = [] + for entry in cache_dir.iterdir(): + if not entry.is_dir(): + continue + if entry.name.lower() != expected_lower: + continue + if not entry.name.startswith("models--"): + continue + repo_part = entry.name[len("models--") :] + if not repo_part: + continue + candidates.append(repo_part.replace("--", "/")) + + if candidates: + # Deterministic tie-break if multiple case variants coexist. + resolved = sorted(candidates)[0] + if len(candidates) > 1: + _CACHE_CASE_RESOLUTION_STATS["tie_breaks"] += 1 + _CACHE_CASE_RESOLUTION_STATS["variant_hits"] += 1 + if use_memo: + _CACHE_CASE_RESOLUTION_MEMO[model_name] = resolved + return resolved + except Exception as exc: + _CACHE_CASE_RESOLUTION_STATS["errors"] += 1 + logger.debug(f"Could not resolve cached repo_id case for '{model_name}': {exc}") + + _CACHE_CASE_RESOLUTION_STATS["fallbacks"] += 1 + return model_name + + +def get_cache_case_resolution_stats() -> dict[str, int]: + """Return a copy of case-resolution instrumentation counters.""" + return dict(_CACHE_CASE_RESOLUTION_STATS) + + +def reset_cache_case_resolution_state() -> None: + """Clear resolver memo and counters (primarily for tests).""" + _CACHE_CASE_RESOLUTION_MEMO.clear() + for key in _CACHE_CASE_RESOLUTION_STATS: + _CACHE_CASE_RESOLUTION_STATS[key] = 0 diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx index d3a86846de..da4c0cbbb9 100644 --- a/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx +++ b/studio/frontend/src/components/assistant-ui/tool-ui-web-search.tsx @@ -52,6 +52,18 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({ status, }) => { const query = (args as { query?: string })?.query ?? ""; + const url = ((args as { url?: string })?.url ?? "").trim(); + const isUrlFetch = !!url; + const displayDomain = (() => { + if (!url) return ""; + try { + const parsed = new URL(url); + if (parsed.protocol !== "http:" && parsed.protocol !== "https:") return ""; + return parsed.hostname.replace(/^www\./, ""); + } catch { + return ""; + } + })(); const isRunning = status?.type === "running"; const sources = result ? parseSearchResults( @@ -75,7 +87,13 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({ return ( @@ -83,7 +101,12 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({ {isRunning ? (
- Searching for “{query}”… + + {isUrlFetch + ? <>Reading {displayDomain || "page"}… + : <>Searching for “{query}”… + } +
) : sources.length > 0 ? (
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index e287daf33a..3d1bce2905 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -421,6 +421,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { return { async *run({ messages, abortSignal, unstable_threadId }) { let runtime = useChatRuntimeStore.getState(); + // Capture the thread ID once at the start so it stays stable even if + // the user switches chats while waiting for model load / auto-load. + const resolvedThreadId = + (unstable_threadId ?? runtime.activeThreadId) || undefined; // Wait for in-progress model load to finish before inferring if (runtime.modelLoading) { @@ -473,14 +477,14 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } runtime.clearPendingAudio(); } - const useAdapter = await resolveUseAdapter(unstable_threadId); + const useAdapter = await resolveUseAdapter(resolvedThreadId); // ── Audio model path (non-streaming) ───────────────────── const activeModel = runtime.models.find( (m) => m.id === params.checkpoint, ); if (activeModel?.isAudio && !activeModel?.hasAudioInput) { - const threadKey = unstable_threadId || "__default"; + const threadKey = resolvedThreadId || "__default"; runtime.setThreadRunning(threadKey, true); try { yield { @@ -527,7 +531,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { return; } - const threadKey = unstable_threadId || "__default"; + const threadKey = resolvedThreadId || "__default"; let waitingFirstChunk = true; let firstTokenSettled = false; const streamStartTime = Date.now(); @@ -600,7 +604,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const mins = useChatRuntimeStore.getState().toolCallTimeout; return mins >= 9999 ? 9999 : mins * 60; })(), - session_id: unstable_threadId || undefined, + session_id: resolvedThreadId, } : {}), }, @@ -641,7 +645,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { let parsedResult: string | { text: string; images: string[]; sessionId: string }; if (imgIdx !== -1) { const text = rawResult.slice(0, imgIdx); - const sessionId = unstable_threadId || ""; + // Fall back to "_default" to match the backend sandbox directory + // used when no session_id is provided (see tools.py _get_workdir). + const sessionId = resolvedThreadId || "_default"; try { const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[]; parsedResult = { text, images, sessionId }; diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 08450c7ec7..1dbff145ee 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -592,6 +592,9 @@ export function ChatPage(): ReactElement { }, []); const handleNewCompare = useCallback(() => { setView({ mode: "compare", pairId: crypto.randomUUID() }); + // Clear activeThreadId so compare panes do not inherit the single-chat + // thread ID as a fallback for session_id routing. + useChatRuntimeStore.getState().setActiveThreadId(null); useChatRuntimeStore.getState().setContextUsage(null); }, []); @@ -619,6 +622,9 @@ export function ChatPage(): ReactElement { const enterCompare = useCallback(() => { setViewBeforeCompare((prev) => prev ?? view); setView({ mode: "compare", pairId: crypto.randomUUID() }); + // Clear activeThreadId so compare panes do not inherit the single-chat + // thread ID as a fallback for session_id routing. + useChatRuntimeStore.getState().setActiveThreadId(null); useChatRuntimeStore.getState().setContextUsage(null); }, [view]); @@ -626,9 +632,13 @@ export function ChatPage(): ReactElement { if (!viewBeforeCompare) return; setView(viewBeforeCompare); setViewBeforeCompare(null); - // Restore context usage from the active thread's last assistant message + // Restore context usage from the active thread's last assistant message. + // Use the thread ID from the saved view rather than the store, because + // activeThreadId may have been cleared on compare entry. const store = useChatRuntimeStore.getState(); - const threadId = store.activeThreadId; + const threadId = + ("threadId" in viewBeforeCompare ? viewBeforeCompare.threadId : null) ?? + store.activeThreadId; if (threadId) { void db.messages .where("threadId") @@ -735,6 +745,7 @@ export function ChatPage(): ReactElement { await selectModelRef.current({ id: targetLora.id, isLora: true }); if (canceled) return; setView({ mode: "compare", pairId: crypto.randomUUID() }); + useChatRuntimeStore.getState().setActiveThreadId(null); useChatRuntimeStore.getState().setContextUsage(null); clearHandoff(); console.info("[chat-handoff] loaded lora + opened compare"); diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index e5b0814343..550df2bf7c 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -280,6 +280,15 @@ export function ChatSettingsPanel({ }: ChatSettingsPanelProps) { const isMobile = useIsMobile(); const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; + const speculativeType = useChatRuntimeStore((s) => s.speculativeType); + const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType); + const loadedSpeculativeType = useChatRuntimeStore( + (s) => s.loadedSpeculativeType, + ); + const currentModels = useChatRuntimeStore((s) => s.models); + const currentCheckpoint = params.checkpoint; + const currentModelIsVision = + currentModels.find((m) => m.id === currentCheckpoint)?.isVision ?? false; const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); const ggufMaxContextLength = useChatRuntimeStore( (s) => s.ggufMaxContextLength, @@ -299,7 +308,8 @@ export function ChatSettingsPanel({ const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null; const kvDirty = kvCacheDtype !== loadedKvCacheDtype; const ctxDirty = customContextLength !== null; - const modelSettingsDirty = kvDirty || ctxDirty; + const specDirty = speculativeType !== loadedSpeculativeType; + const modelSettingsDirty = kvDirty || ctxDirty || specDirty; const [customPresets, setCustomPresets] = useState(() => loadSavedCustomPresets(), ); @@ -580,6 +590,32 @@ export function ChatSettingsPanel({
+ {!currentModelIsVision && ( +
+
+
+ Speculative Decoding +
+
+ Speed up generation with no VRAM cost. +
+
+ +
+ )} {modelSettingsDirty && (
{execution.dataset.length === 0 ? ( -

No rows returned.

+ isExecutionInProgress(execution.status) ? ( +
+ +
+

+ Generating data… +

+

+ Check the Overview tab for live terminal logs. +

+
+
+ ) : ( +

No rows returned.

+ ) ) : tableColumns.length === 0 ? (

All columns hidden. Use Columns to show at least one. diff --git a/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx b/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx index 06b4b20629..8e2e710299 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx @@ -116,10 +116,12 @@ export function ExecutionOverviewTab({ />

-

- LLM columns - {formatMetricValue(llmColumnCount)} -

+ {llmColumnCount > 0 && ( +

+ LLM columns + {formatMetricValue(llmColumnCount)} +

+ )}

Null rate {nullRate?.toFixed(1) ?? "--"}% @@ -164,40 +166,42 @@ export function ExecutionOverviewTab({

-
-
-

Model usage

- -
- {modelUsageRows.length === 0 ? ( -

No model usage yet.

- ) : ( -
- - - - Model - Input - Output - - - - {modelUsageRows.map((usage) => ( - - {usage.model} - - {formatMetricValue(usage.input)} - - - {formatMetricValue(usage.output)} - - - ))} - -
+ {(llmColumnCount > 0 || modelUsageRows.length > 0) && ( +
+
+

Model usage

+
- )} -
+ {modelUsageRows.length === 0 ? ( +

No model usage yet.

+ ) : ( +
+ + + + Model + Input + Output + + + + {modelUsageRows.map((usage) => ( + + {usage.model} + + {formatMetricValue(usage.input)} + + + {formatMetricValue(usage.output)} + + + ))} + +
+
+ )} +
+ )}
)}
diff --git a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx index f30cf61012..f442d8f98e 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx @@ -59,7 +59,7 @@ export function ExecutionsView({ typeof value === "number" && Number.isFinite(value) ? `${value.toLocaleString()} s` : "--"; - const [detailTab, setDetailTab] = useState("overview"); + const [detailTab, setDetailTab] = useState("data"); const [hiddenDatasetColumnsByExecution, setHiddenDatasetColumnsByExecution] = useState< Record >({}); @@ -341,12 +341,16 @@ export function ExecutionsView({ }, [selectedExecution]); useEffect(() => { - if (!terminalRef.current) { + setDetailTab("data"); + }, [selectedExecution?.id]); + + useEffect(() => { + if (detailTab !== "overview" || !terminalRef.current) { return; } shouldStickTerminalToBottomRef.current = true; terminalRef.current.scrollTop = terminalRef.current.scrollHeight; - }, [selectedExecution?.id]); + }, [detailTab, selectedExecution?.id]); useEffect(() => { if (!terminalRef.current) { @@ -440,9 +444,9 @@ export function ExecutionsView({
+ Data Overview Columns - Data Raw
diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 5e8fe314b4..8d06c7d0e1 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -65,9 +65,10 @@ def env_int(name: str, default: int, *, minimum: int | None = None) -> int: # errors. Only use "master" temporarily when the latest release is missing # support for a new model architecture. DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", "latest") -# Force all installs to use mainline llama.cpp from ggml-org. -# Previously: DEFAULT_PUBLISHED_REPO = os.environ.get("UNSLOTH_LLAMA_RELEASE_REPO", "unslothai/llama.cpp") -DEFAULT_PUBLISHED_REPO = "ggml-org/llama.cpp" +# Default published repo for prebuilt release resolution. Linux uses +# Unsloth prebuilts; setup.sh/setup.ps1 pass --published-repo explicitly +# for macOS/Windows to override with ggml-org/llama.cpp when needed. +DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" DEFAULT_PUBLISHED_TAG = os.environ.get("UNSLOTH_LLAMA_RELEASE_TAG") DEFAULT_PUBLISHED_MANIFEST_ASSET = os.environ.get( "UNSLOTH_LLAMA_RELEASE_MANIFEST_ASSET", "llama-prebuilt-manifest.json" @@ -89,6 +90,12 @@ GITHUB_AUTH_HOSTS = {"api.github.com", "github.com"} RETRYABLE_HTTP_STATUS = {408, 429, 500, 502, 503, 504} HTTP_FETCH_ATTEMPTS = 4 HTTP_FETCH_BASE_DELAY_SECONDS = 0.75 +JSON_FETCH_ATTEMPTS = 3 +DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES = env_int( + "UNSLOTH_LLAMA_GITHUB_RELEASE_SCAN_MAX_PAGES", + 5, + minimum = 1, +) SERVER_PORT_BIND_ATTEMPTS = 3 SERVER_BIND_RETRY_WINDOW_SECONDS = 5.0 TTY_PROGRESS_START_DELAY_SECONDS = 0.5 @@ -97,6 +104,58 @@ DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS = env_int( 2, minimum = 1, ) +FORCE_COMPILE_DEFAULT_REF = os.environ.get("UNSLOTH_LLAMA_FORCE_COMPILE_REF", "master") + +DIRECT_LINUX_BUNDLE_PROFILES: dict[str, dict[str, Any]] = { + "cuda12-older": { + "runtime_line": "cuda12", + "coverage_class": "older", + "supported_sms": ["70", "75", "80", "86", "89"], + "min_sm": 70, + "max_sm": 89, + "rank": 10, + }, + "cuda12-newer": { + "runtime_line": "cuda12", + "coverage_class": "newer", + "supported_sms": ["86", "89", "90", "100", "120"], + "min_sm": 86, + "max_sm": 120, + "rank": 20, + }, + "cuda12-portable": { + "runtime_line": "cuda12", + "coverage_class": "portable", + "supported_sms": ["70", "75", "80", "86", "89", "90", "100", "120"], + "min_sm": 70, + "max_sm": 120, + "rank": 30, + }, + "cuda13-older": { + "runtime_line": "cuda13", + "coverage_class": "older", + "supported_sms": ["75", "80", "86", "89"], + "min_sm": 75, + "max_sm": 89, + "rank": 40, + }, + "cuda13-newer": { + "runtime_line": "cuda13", + "coverage_class": "newer", + "supported_sms": ["86", "89", "90", "100", "120"], + "min_sm": 86, + "max_sm": 120, + "rank": 50, + }, + "cuda13-portable": { + "runtime_line": "cuda13", + "coverage_class": "portable", + "supported_sms": ["75", "80", "86", "89", "90", "100", "120"], + "min_sm": 75, + "max_sm": 120, + "rank": 60, + }, +} @dataclass @@ -753,32 +812,48 @@ def download_bytes( def fetch_json(url: str) -> Any: - try: - data = download_bytes( - url, - timeout = 30, - headers = github_api_headers(url) - if is_github_api_url(url) - else auth_headers(url), - ) - except urllib.error.HTTPError as exc: - if exc.code == 403 and is_github_api_url(url): - hint = "" - if not (os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")): - hint = "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits" - raise RuntimeError(f"GitHub API returned 403 for {url}{hint}") from exc - raise - if not data: - raise RuntimeError(f"downloaded empty JSON payload from {url}") - try: - payload = json.loads(data.decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as exc: - raise RuntimeError(f"downloaded invalid JSON from {url}: {exc}") from exc - if not isinstance(payload, dict) and not isinstance(payload, list): - raise RuntimeError( - f"downloaded unexpected JSON type from {url}: {type(payload).__name__}" - ) - return payload + attempts = JSON_FETCH_ATTEMPTS if is_github_api_url(url) else 1 + last_decode_exc: Exception | None = None + for attempt in range(1, attempts + 1): + try: + data = download_bytes( + url, + timeout = 30, + headers = github_api_headers(url) + if is_github_api_url(url) + else auth_headers(url), + ) + except urllib.error.HTTPError as exc: + if exc.code == 403 and is_github_api_url(url): + hint = "" + if not (os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")): + hint = ( + "; set GH_TOKEN or GITHUB_TOKEN to avoid GitHub API rate limits" + ) + raise RuntimeError(f"GitHub API returned 403 for {url}{hint}") from exc + raise + if not data: + last_decode_exc = RuntimeError(f"downloaded empty JSON payload from {url}") + else: + try: + payload = json.loads(data.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + last_decode_exc = RuntimeError( + f"downloaded invalid JSON from {url}: {exc}" + ) + else: + if not isinstance(payload, dict) and not isinstance(payload, list): + raise RuntimeError( + f"downloaded unexpected JSON type from {url}: {type(payload).__name__}" + ) + return payload + if attempt >= attempts: + assert last_decode_exc is not None + raise last_decode_exc + log(f"json fetch failed ({attempt}/{attempts}) for {url}; retrying") + sleep_backoff(attempt) + assert last_decode_exc is not None + raise last_decode_exc def download_file(url: str, destination: Path) -> None: @@ -838,12 +913,16 @@ def download_file_verified( url: str, destination: Path, *, - expected_sha256: str, + expected_sha256: str | None, label: str, ) -> None: normalized_expected = normalize_sha256_digest(expected_sha256) if not normalized_expected: - raise PrebuiltFallback(f"{label} did not have a valid approved sha256") + download_file(url, destination) + log( + f"downloaded {label} without a published sha256; relying on install validation" + ) + return for attempt in range(1, 3): download_file(url, destination) @@ -898,7 +977,12 @@ def github_release(repo: str, tag: str) -> dict[str, Any]: return payload -def github_releases(repo: str, *, per_page: int = 100) -> list[dict[str, Any]]: +def github_releases( + repo: str, + *, + per_page: int = 100, + max_pages: int = 0, +) -> list[dict[str, Any]]: releases: list[dict[str, Any]] = [] page = 1 while True: @@ -912,6 +996,8 @@ def github_releases(repo: str, *, per_page: int = 100) -> list[dict[str, Any]]: if len(payload) < per_page: break page += 1 + if max_pages > 0 and page > max_pages: + break return releases @@ -925,6 +1011,372 @@ def latest_upstream_release_tag() -> str: return tag +def is_release_tag_like(value: str | None) -> bool: + return isinstance(value, str) and bool(re.fullmatch(r"b\d+", value.strip())) + + +def release_time_sort_key(release: dict[str, Any]) -> tuple[str, int]: + published_at = release.get("published_at") + created_at = release.get("created_at") + release_id = release.get("id") + timestamp = ( + published_at + if isinstance(published_at, str) and published_at + else created_at + if isinstance(created_at, str) and created_at + else "" + ) + try: + normalized_id = int(release_id) + except (TypeError, ValueError): + normalized_id = 0 + return (timestamp, normalized_id) + + +def iter_release_payloads_by_time( + repo: str, + published_release_tag: str = "", + requested_tag: str = "", +) -> Iterable[dict[str, Any]]: + if published_release_tag: + yield github_release(repo, published_release_tag) + return + + if ( + requested_tag + and requested_tag != "latest" + and is_release_tag_like(requested_tag) + ): + try: + yield github_release(repo, requested_tag) + return + except urllib.error.HTTPError as exc: + if exc.code == 404: + log( + f"release tag {requested_tag} not found in {repo}; scanning recent releases" + ) + else: + raise + except Exception: + raise + + releases = [ + release + for release in github_releases( + repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES + ) + if isinstance(release, dict) + and not release.get("draft") + and not release.get("prerelease") + ] + releases.sort(key = release_time_sort_key, reverse = True) + for release in releases: + yield release + + +def direct_release_matches_request( + *, release_tag: str, llama_tag: str, requested_tag: str +) -> bool: + if requested_tag == "latest": + return True + for candidate in (release_tag, llama_tag): + if refs_match(candidate, requested_tag): + return True + return False + + +def synthetic_checksums_for_release( + repo: str, release_tag: str, upstream_tag: str +) -> ApprovedReleaseChecksums: + return ApprovedReleaseChecksums( + repo = repo, + release_tag = release_tag, + upstream_tag = upstream_tag, + artifacts = {}, + ) + + +def parse_direct_linux_release_bundle( + repo: str, release: dict[str, Any] +) -> PublishedReleaseBundle | None: + release_tag = release.get("tag_name") + if not isinstance(release_tag, str) or not release_tag: + return None + + assets = release_asset_map(release) + artifacts: list[PublishedLlamaArtifact] = [] + inferred_labels: list[str] = [] + + linux_asset_re = re.compile( + r"^app-(?P