Merge branch 'main' into pip
This commit is contained in:
Daniel Han
commit
9ad3b761ee
41 changed files with 2450 additions and 451 deletions
|
|
@ -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). <br>
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -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 ─────────────────────────────
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
||||
# =====================================================================
|
||||
|
|
|
|||
|
|
@ -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 <path>
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
120
studio/backend/tests/test_cache_case_resolution.py
Normal file
120
studio/backend/tests/test_cache_case_resolution.py
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
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"
|
||||
|
|
@ -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"
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<ToolFallbackRoot open={open} onOpenChange={setOpen}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={query ? `Searched "${query}"` : "Web Search"}
|
||||
toolName={
|
||||
isUrlFetch
|
||||
? displayDomain ? `Read ${displayDomain}` : "Read page"
|
||||
: query
|
||||
? `Searched "${query}"`
|
||||
: "Web Search"
|
||||
}
|
||||
status={status}
|
||||
icon={GlobeIcon}
|
||||
/>
|
||||
|
|
@ -83,7 +101,12 @@ const WebSearchToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
{isRunning ? (
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<LoaderIcon className="size-3.5 animate-spin" />
|
||||
<span>Searching for “{query}”…</span>
|
||||
<span>
|
||||
{isUrlFetch
|
||||
? <>Reading {displayDomain || "page"}…</>
|
||||
: <>Searching for “{query}”…</>
|
||||
}
|
||||
</span>
|
||||
</div>
|
||||
) : sources.length > 0 ? (
|
||||
<div className="flex flex-wrap gap-1.5">
|
||||
|
|
|
|||
|
|
@ -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 };
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
|
|
|
|||
|
|
@ -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<Preset[]>(() =>
|
||||
loadSavedCustomPresets(),
|
||||
);
|
||||
|
|
@ -580,6 +590,32 @@ export function ChatSettingsPanel({
|
|||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
{!currentModelIsVision && (
|
||||
<div className="flex items-center justify-between gap-3">
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium">
|
||||
Speculative Decoding
|
||||
</div>
|
||||
<div className="text-[11px] text-muted-foreground">
|
||||
Speed up generation with no VRAM cost.
|
||||
</div>
|
||||
</div>
|
||||
<Select
|
||||
value={speculativeType ?? "off"}
|
||||
onValueChange={(v) => {
|
||||
setSpeculativeType(v === "off" ? null : v);
|
||||
}}
|
||||
>
|
||||
<SelectTrigger className="h-7 w-[120px] text-xs">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="ngram-mod">On</SelectItem>
|
||||
<SelectItem value="off">Off</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
)}
|
||||
{modelSettingsDirty && (
|
||||
<div className="flex flex-wrap gap-1.5 pt-1">
|
||||
<button
|
||||
|
|
@ -594,6 +630,7 @@ export function ChatSettingsPanel({
|
|||
onClick={() => {
|
||||
setCustomContextLength(null);
|
||||
setKvCacheDtype(loadedKvCacheDtype);
|
||||
setSpeculativeType(loadedSpeculativeType);
|
||||
}}
|
||||
className="rounded-md border px-2.5 py-1 text-[11px] font-medium text-muted-foreground transition-colors hover:bg-accent"
|
||||
>
|
||||
|
|
|
|||
|
|
@ -250,6 +250,7 @@ export function useChatModelRuntime() {
|
|||
const ggufNativeContextLength = statusRes.is_gguf
|
||||
? (statusRes.native_context_length ?? null)
|
||||
: null;
|
||||
const currentSpecType = statusRes.speculative_type ?? null;
|
||||
useChatRuntimeStore.setState({
|
||||
supportsReasoning,
|
||||
reasoningAlwaysOn,
|
||||
|
|
@ -257,6 +258,8 @@ export function useChatModelRuntime() {
|
|||
ggufContextLength: currentGgufContextLength,
|
||||
ggufMaxContextLength,
|
||||
ggufNativeContextLength,
|
||||
speculativeType: currentSpecType,
|
||||
loadedSpeculativeType: currentSpecType,
|
||||
});
|
||||
|
||||
// Set reasoning default for Qwen3.5 small models
|
||||
|
|
@ -393,7 +396,7 @@ export function useChatModelRuntime() {
|
|||
previousWasUnloaded = true;
|
||||
}
|
||||
|
||||
const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength } = useChatRuntimeStore.getState();
|
||||
const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength, speculativeType } = useChatRuntimeStore.getState();
|
||||
// GGUF: use custom context length, or 0 = model's native context
|
||||
// Non-GGUF: use the Max Seq Length slider value
|
||||
const effectiveMaxSeqLength = customContextLength != null
|
||||
|
|
@ -409,6 +412,7 @@ export function useChatModelRuntime() {
|
|||
trust_remote_code: paramsBeforeLoad.trustRemoteCode ?? false,
|
||||
chat_template_override: chatTemplateOverride,
|
||||
cache_type_kv: kvCacheDtype,
|
||||
speculative_type: speculativeType,
|
||||
});
|
||||
|
||||
// If cancelled while loading, don't update UI to show
|
||||
|
|
@ -431,6 +435,7 @@ export function useChatModelRuntime() {
|
|||
}
|
||||
}
|
||||
const loadedKv = loadResponse.cache_type_kv ?? null;
|
||||
const loadedSpec = loadResponse.speculative_type ?? null;
|
||||
const nativeCtx = loadResponse.is_gguf
|
||||
? (loadResponse.context_length ?? 131072)
|
||||
: null;
|
||||
|
|
@ -457,6 +462,8 @@ export function useChatModelRuntime() {
|
|||
codeToolsEnabled: loadResponse.supports_tools ?? false,
|
||||
kvCacheDtype: loadedKv,
|
||||
loadedKvCacheDtype: loadedKv,
|
||||
speculativeType: loadedSpec,
|
||||
loadedSpeculativeType: loadedSpec,
|
||||
customContextLength: keepCustomCtx,
|
||||
defaultChatTemplate: loadResponse.chat_template ?? null,
|
||||
chatTemplateOverride: null,
|
||||
|
|
|
|||
|
|
@ -679,9 +679,33 @@ function ThreadNewChatSwitch({
|
|||
const isLoading = useAuiState(({ threads }) => threads.isLoading);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isLoading) {
|
||||
aui.threads().switchToNewThread();
|
||||
if (isLoading) {
|
||||
return;
|
||||
}
|
||||
|
||||
let cancelled = false;
|
||||
// Clear immediately so the adapter never picks up a stale thread ID
|
||||
// from a previous chat while we initialize the new one.
|
||||
useChatRuntimeStore.getState().setActiveThreadId(null);
|
||||
|
||||
void (async () => {
|
||||
try {
|
||||
aui.threads().switchToNewThread();
|
||||
const { remoteId } = await aui.threadListItem().initialize();
|
||||
if (!cancelled) {
|
||||
useChatRuntimeStore.getState().setActiveThreadId(remoteId);
|
||||
}
|
||||
} catch (error) {
|
||||
if (!cancelled) {
|
||||
useChatRuntimeStore.getState().setActiveThreadId(null);
|
||||
}
|
||||
console.error("Failed to initialize new chat thread", error);
|
||||
}
|
||||
})();
|
||||
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
}, [aui, isLoading, nonce]);
|
||||
|
||||
return null;
|
||||
|
|
@ -730,7 +754,7 @@ export function ChatRuntimeProvider({
|
|||
|
||||
return (
|
||||
<AssistantRuntimeProvider runtime={runtime} aui={aui}>
|
||||
<ActiveThreadSync enabled={modelType === "base" && !pairId} />
|
||||
<ActiveThreadSync enabled={modelType === "base" && !pairId && !newThreadNonce} />
|
||||
{initialThreadId && <ThreadAutoSwitch threadId={initialThreadId} />}
|
||||
{!initialThreadId && newThreadNonce && (
|
||||
<ThreadNewChatSwitch nonce={newThreadNonce} />
|
||||
|
|
|
|||
|
|
@ -165,6 +165,8 @@ type ChatRuntimeStore = {
|
|||
toolCallTimeout: number;
|
||||
kvCacheDtype: string | null;
|
||||
loadedKvCacheDtype: string | null;
|
||||
speculativeType: string | null;
|
||||
loadedSpeculativeType: string | null;
|
||||
customContextLength: number | null;
|
||||
defaultChatTemplate: string | null;
|
||||
chatTemplateOverride: string | null;
|
||||
|
|
@ -198,6 +200,7 @@ type ChatRuntimeStore = {
|
|||
setMaxToolCallsPerMessage: (value: number) => void;
|
||||
setToolCallTimeout: (value: number) => void;
|
||||
setKvCacheDtype: (dtype: string | null) => void;
|
||||
setSpeculativeType: (type: string | null) => void;
|
||||
setCustomContextLength: (v: number | null) => void;
|
||||
setChatTemplateOverride: (template: string | null) => void;
|
||||
setPendingAudio: (base64: string, name: string) => void;
|
||||
|
|
@ -230,6 +233,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5),
|
||||
kvCacheDtype: null,
|
||||
loadedKvCacheDtype: null,
|
||||
speculativeType: "ngram-mod",
|
||||
loadedSpeculativeType: null,
|
||||
customContextLength: null,
|
||||
defaultChatTemplate: null,
|
||||
chatTemplateOverride: null,
|
||||
|
|
@ -302,6 +307,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
toolStatus: null,
|
||||
kvCacheDtype: null,
|
||||
loadedKvCacheDtype: null,
|
||||
speculativeType: "ngram-mod",
|
||||
loadedSpeculativeType: null,
|
||||
customContextLength: null,
|
||||
defaultChatTemplate: null,
|
||||
chatTemplateOverride: null,
|
||||
|
|
@ -327,6 +334,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
|||
return { toolCallTimeout };
|
||||
}),
|
||||
setKvCacheDtype: (kvCacheDtype) => set({ kvCacheDtype }),
|
||||
setSpeculativeType: (speculativeType) => set({ speculativeType }),
|
||||
setCustomContextLength: (customContextLength) => set({ customContextLength }),
|
||||
setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }),
|
||||
setPendingAudio: (base64, name) =>
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ export interface LoadModelRequest {
|
|||
trust_remote_code?: boolean;
|
||||
chat_template_override?: string | null;
|
||||
cache_type_kv?: string | null;
|
||||
speculative_type?: string | null;
|
||||
}
|
||||
|
||||
export interface ValidateModelResponse {
|
||||
|
|
@ -93,6 +94,7 @@ export interface LoadModelResponse {
|
|||
supports_tools?: boolean;
|
||||
cache_type_kv?: string | null;
|
||||
chat_template?: string | null;
|
||||
speculative_type?: string | null;
|
||||
}
|
||||
|
||||
export interface UnloadModelRequest {
|
||||
|
|
@ -123,6 +125,7 @@ export interface InferenceStatusResponse {
|
|||
context_length?: number | null;
|
||||
max_context_length?: number | null;
|
||||
native_context_length?: number | null;
|
||||
speculative_type?: string | null;
|
||||
}
|
||||
|
||||
export interface AudioGenerationResponse {
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import {
|
|||
DropdownMenuLabel,
|
||||
DropdownMenuTrigger,
|
||||
} from "@/components/ui/dropdown-menu";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { isExecutionInProgress } from "../../executions/execution-helpers";
|
||||
import type { RecipeExecutionRecord } from "../../execution-types";
|
||||
|
|
@ -121,7 +122,21 @@ export function ExecutionDataTab({
|
|||
</div>
|
||||
</div>
|
||||
{execution.dataset.length === 0 ? (
|
||||
<p className="text-xs text-muted-foreground">No rows returned.</p>
|
||||
isExecutionInProgress(execution.status) ? (
|
||||
<div className="flex flex-col items-center justify-center gap-3 py-12 text-center">
|
||||
<Spinner className="size-5" />
|
||||
<div className="space-y-1">
|
||||
<p className="text-sm font-medium text-muted-foreground">
|
||||
Generating data…
|
||||
</p>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Check the Overview tab for live terminal logs.
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
) : (
|
||||
<p className="text-xs text-muted-foreground">No rows returned.</p>
|
||||
)
|
||||
) : tableColumns.length === 0 ? (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
All columns hidden. Use Columns to show at least one.
|
||||
|
|
|
|||
|
|
@ -116,10 +116,12 @@ export function ExecutionOverviewTab({
|
|||
/>
|
||||
</div>
|
||||
<div className="space-y-1.5 text-xs">
|
||||
<p className="flex items-center justify-between gap-3">
|
||||
<span className="text-muted-foreground">LLM columns</span>
|
||||
<span className="font-semibold">{formatMetricValue(llmColumnCount)}</span>
|
||||
</p>
|
||||
{llmColumnCount > 0 && (
|
||||
<p className="flex items-center justify-between gap-3">
|
||||
<span className="text-muted-foreground">LLM columns</span>
|
||||
<span className="font-semibold">{formatMetricValue(llmColumnCount)}</span>
|
||||
</p>
|
||||
)}
|
||||
<p className="flex items-center justify-between gap-3">
|
||||
<span className="text-muted-foreground">Null rate</span>
|
||||
<span className="font-semibold">{nullRate?.toFixed(1) ?? "--"}%</span>
|
||||
|
|
@ -164,40 +166,42 @@ export function ExecutionOverviewTab({
|
|||
</div>
|
||||
</div>
|
||||
</div>
|
||||
<div className="rounded-xl border border-border/60 bg-card/55 p-3">
|
||||
<div className="mb-2 flex items-center justify-between">
|
||||
<p className="text-xs text-muted-foreground">Model usage</p>
|
||||
<HugeiconsIcon icon={Flag02Icon} className="size-4 text-muted-foreground" />
|
||||
</div>
|
||||
{modelUsageRows.length === 0 ? (
|
||||
<p className="text-xs text-muted-foreground">No model usage yet.</p>
|
||||
) : (
|
||||
<div className="overflow-hidden rounded-lg border border-border/60 bg-card/50">
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead>Model</TableHead>
|
||||
<TableHead className="text-right">Input</TableHead>
|
||||
<TableHead className="text-right">Output</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{modelUsageRows.map((usage) => (
|
||||
<TableRow key={usage.model}>
|
||||
<TableCell className="max-w-[320px] truncate">{usage.model}</TableCell>
|
||||
<TableCell className="text-right">
|
||||
{formatMetricValue(usage.input)}
|
||||
</TableCell>
|
||||
<TableCell className="text-right">
|
||||
{formatMetricValue(usage.output)}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
</TableBody>
|
||||
</Table>
|
||||
{(llmColumnCount > 0 || modelUsageRows.length > 0) && (
|
||||
<div className="rounded-xl border border-border/60 bg-card/55 p-3">
|
||||
<div className="mb-2 flex items-center justify-between">
|
||||
<p className="text-xs text-muted-foreground">Model usage</p>
|
||||
<HugeiconsIcon icon={Flag02Icon} className="size-4 text-muted-foreground" />
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
{modelUsageRows.length === 0 ? (
|
||||
<p className="text-xs text-muted-foreground">No model usage yet.</p>
|
||||
) : (
|
||||
<div className="overflow-hidden rounded-lg border border-border/60 bg-card/50">
|
||||
<Table>
|
||||
<TableHeader>
|
||||
<TableRow>
|
||||
<TableHead>Model</TableHead>
|
||||
<TableHead className="text-right">Input</TableHead>
|
||||
<TableHead className="text-right">Output</TableHead>
|
||||
</TableRow>
|
||||
</TableHeader>
|
||||
<TableBody>
|
||||
{modelUsageRows.map((usage) => (
|
||||
<TableRow key={usage.model}>
|
||||
<TableCell className="max-w-[320px] truncate">{usage.model}</TableCell>
|
||||
<TableCell className="text-right">
|
||||
{formatMetricValue(usage.input)}
|
||||
</TableCell>
|
||||
<TableCell className="text-right">
|
||||
{formatMetricValue(usage.output)}
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))}
|
||||
</TableBody>
|
||||
</Table>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
<div className="overflow-hidden rounded-xl corner-squircle border">
|
||||
|
|
|
|||
|
|
@ -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<string, string[]>
|
||||
>({});
|
||||
|
|
@ -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({
|
|||
<Tabs value={detailTab} onValueChange={setDetailTab}>
|
||||
<div className="flex items-center justify-between gap-2">
|
||||
<TabsList className="border border-border/60 bg-card/40">
|
||||
<TabsTrigger value="data">Data</TabsTrigger>
|
||||
<TabsTrigger value="overview">Overview</TabsTrigger>
|
||||
<TabsTrigger value="columns">Columns</TabsTrigger>
|
||||
<TabsTrigger value="data">Data</TabsTrigger>
|
||||
<TabsTrigger value="raw">Raw</TabsTrigger>
|
||||
</TabsList>
|
||||
<div className="flex items-center gap-2">
|
||||
|
|
|
|||
|
|
@ -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<label>.+)-(?P<target>linux-x64(?:-cpu)?|linux-x64-(?:cuda12|cuda13)-(?:older|newer|portable))\.tar\.gz$"
|
||||
)
|
||||
for asset_name in sorted(assets):
|
||||
match = linux_asset_re.fullmatch(asset_name)
|
||||
if not match:
|
||||
continue
|
||||
inferred_labels.append(match.group("label"))
|
||||
target = match.group("target")
|
||||
if target in {"linux-x64", "linux-x64-cpu"}:
|
||||
artifacts.append(
|
||||
PublishedLlamaArtifact(
|
||||
asset_name = asset_name,
|
||||
install_kind = "linux-cpu",
|
||||
runtime_line = None,
|
||||
coverage_class = None,
|
||||
supported_sms = [],
|
||||
min_sm = None,
|
||||
max_sm = None,
|
||||
bundle_profile = None,
|
||||
rank = 1000,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
bundle_profile = target.removeprefix("linux-x64-")
|
||||
profile = DIRECT_LINUX_BUNDLE_PROFILES.get(bundle_profile)
|
||||
if profile is None:
|
||||
continue
|
||||
artifacts.append(
|
||||
PublishedLlamaArtifact(
|
||||
asset_name = asset_name,
|
||||
install_kind = "linux-cuda",
|
||||
runtime_line = str(profile["runtime_line"]),
|
||||
coverage_class = str(profile["coverage_class"]),
|
||||
supported_sms = [str(value) for value in profile["supported_sms"]],
|
||||
min_sm = int(profile["min_sm"]),
|
||||
max_sm = int(profile["max_sm"]),
|
||||
bundle_profile = bundle_profile,
|
||||
rank = int(profile["rank"]),
|
||||
)
|
||||
)
|
||||
|
||||
if not artifacts:
|
||||
return None
|
||||
|
||||
upstream_tag = (
|
||||
release_tag
|
||||
if is_release_tag_like(release_tag)
|
||||
else inferred_labels[0]
|
||||
if len(set(inferred_labels)) == 1 and inferred_labels
|
||||
else release_tag
|
||||
)
|
||||
selection_log = [
|
||||
f"published_release: repo={repo}",
|
||||
f"published_release: tag={release_tag}",
|
||||
f"published_release: upstream_tag={upstream_tag}",
|
||||
"published_release: direct_asset_scan=linux",
|
||||
]
|
||||
return PublishedReleaseBundle(
|
||||
repo = repo,
|
||||
release_tag = release_tag,
|
||||
upstream_tag = upstream_tag,
|
||||
assets = assets,
|
||||
manifest_asset_name = DEFAULT_PUBLISHED_MANIFEST_ASSET,
|
||||
artifacts = artifacts,
|
||||
selection_log = selection_log,
|
||||
)
|
||||
|
||||
|
||||
def direct_linux_release_plan(
|
||||
release: dict[str, Any],
|
||||
host: HostInfo,
|
||||
repo: str,
|
||||
requested_tag: str,
|
||||
) -> InstallReleasePlan | None:
|
||||
bundle = parse_direct_linux_release_bundle(repo, release)
|
||||
if bundle is None:
|
||||
return None
|
||||
if not direct_release_matches_request(
|
||||
release_tag = bundle.release_tag,
|
||||
llama_tag = bundle.upstream_tag,
|
||||
requested_tag = requested_tag,
|
||||
):
|
||||
return None
|
||||
|
||||
attempts: list[AssetChoice] = []
|
||||
if host.has_usable_nvidia:
|
||||
selection = linux_cuda_choice_from_release(host, bundle)
|
||||
if selection is not None:
|
||||
attempts.extend(selection.attempts)
|
||||
cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu")
|
||||
if cpu_choice is not None:
|
||||
attempts.append(cpu_choice)
|
||||
if not attempts:
|
||||
raise PrebuiltFallback("no compatible Linux prebuilt asset was found")
|
||||
return InstallReleasePlan(
|
||||
requested_tag = requested_tag,
|
||||
llama_tag = bundle.upstream_tag,
|
||||
release_tag = bundle.release_tag,
|
||||
attempts = attempts,
|
||||
approved_checksums = synthetic_checksums_for_release(
|
||||
repo,
|
||||
bundle.release_tag,
|
||||
bundle.upstream_tag,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def direct_upstream_release_plan(
|
||||
release: dict[str, Any],
|
||||
host: HostInfo,
|
||||
repo: str,
|
||||
requested_tag: str,
|
||||
) -> InstallReleasePlan | None:
|
||||
release_tag = release.get("tag_name")
|
||||
if not isinstance(release_tag, str) or not release_tag:
|
||||
return None
|
||||
if not direct_release_matches_request(
|
||||
release_tag = release_tag,
|
||||
llama_tag = release_tag,
|
||||
requested_tag = requested_tag,
|
||||
):
|
||||
return None
|
||||
|
||||
assets = release_asset_map(release)
|
||||
attempts: list[AssetChoice] = []
|
||||
if host.is_windows and host.is_x86_64:
|
||||
if host.has_usable_nvidia:
|
||||
torch_preference = detect_torch_cuda_runtime_preference(host)
|
||||
attempts.extend(
|
||||
windows_cuda_attempts(
|
||||
host,
|
||||
release_tag,
|
||||
assets,
|
||||
torch_preference.runtime_line,
|
||||
torch_preference.selection_log,
|
||||
)
|
||||
)
|
||||
cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip"
|
||||
cpu_url = assets.get(cpu_asset)
|
||||
if cpu_url:
|
||||
attempts.append(
|
||||
AssetChoice(
|
||||
repo = repo,
|
||||
tag = release_tag,
|
||||
name = cpu_asset,
|
||||
url = cpu_url,
|
||||
source_label = "upstream",
|
||||
install_kind = "windows-cpu",
|
||||
)
|
||||
)
|
||||
elif host.is_macos and host.is_arm64:
|
||||
asset_name = f"llama-{release_tag}-bin-macos-arm64.tar.gz"
|
||||
asset_url = assets.get(asset_name)
|
||||
if asset_url:
|
||||
attempts.append(
|
||||
AssetChoice(
|
||||
repo = repo,
|
||||
tag = release_tag,
|
||||
name = asset_name,
|
||||
url = asset_url,
|
||||
source_label = "upstream",
|
||||
install_kind = "macos-arm64",
|
||||
)
|
||||
)
|
||||
elif host.is_macos and host.is_x86_64:
|
||||
asset_name = f"llama-{release_tag}-bin-macos-x64.tar.gz"
|
||||
asset_url = assets.get(asset_name)
|
||||
if asset_url:
|
||||
attempts.append(
|
||||
AssetChoice(
|
||||
repo = repo,
|
||||
tag = release_tag,
|
||||
name = asset_name,
|
||||
url = asset_url,
|
||||
source_label = "upstream",
|
||||
install_kind = "macos-x64",
|
||||
)
|
||||
)
|
||||
elif host.is_linux and host.is_x86_64 and not host.has_usable_nvidia:
|
||||
asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz"
|
||||
asset_url = assets.get(asset_name)
|
||||
if asset_url:
|
||||
attempts.append(
|
||||
AssetChoice(
|
||||
repo = repo,
|
||||
tag = release_tag,
|
||||
name = asset_name,
|
||||
url = asset_url,
|
||||
source_label = "upstream",
|
||||
install_kind = "linux-cpu",
|
||||
)
|
||||
)
|
||||
if not attempts:
|
||||
raise PrebuiltFallback("no compatible upstream prebuilt asset was found")
|
||||
return InstallReleasePlan(
|
||||
requested_tag = requested_tag,
|
||||
llama_tag = release_tag,
|
||||
release_tag = release_tag,
|
||||
attempts = attempts,
|
||||
approved_checksums = synthetic_checksums_for_release(
|
||||
repo,
|
||||
release_tag,
|
||||
release_tag,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def resolve_simple_install_release_plans(
|
||||
llama_tag: str,
|
||||
host: HostInfo,
|
||||
published_repo: str,
|
||||
published_release_tag: str,
|
||||
*,
|
||||
max_release_fallbacks: int = DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS,
|
||||
) -> tuple[str, list[InstallReleasePlan]]:
|
||||
repo = published_repo or DEFAULT_PUBLISHED_REPO
|
||||
requested_tag = normalized_requested_llama_tag(llama_tag)
|
||||
allow_older_release_fallback = (
|
||||
requested_tag == "latest" and not published_release_tag
|
||||
)
|
||||
release_limit = max(1, max_release_fallbacks)
|
||||
plans: list[InstallReleasePlan] = []
|
||||
last_error: PrebuiltFallback | None = None
|
||||
|
||||
try:
|
||||
releases = iter_release_payloads_by_time(
|
||||
repo, published_release_tag, requested_tag
|
||||
)
|
||||
for release in releases:
|
||||
try:
|
||||
if host.is_linux and repo == "unslothai/llama.cpp":
|
||||
plan = direct_linux_release_plan(release, host, repo, requested_tag)
|
||||
else:
|
||||
plan = direct_upstream_release_plan(
|
||||
release, host, repo, requested_tag
|
||||
)
|
||||
if plan is None:
|
||||
continue
|
||||
except PrebuiltFallback as exc:
|
||||
last_error = exc
|
||||
if not allow_older_release_fallback:
|
||||
raise
|
||||
release_tag = release.get("tag_name") or "unknown"
|
||||
log(
|
||||
"published release skipped for install planning: "
|
||||
f"{repo}@{release_tag} ({exc})"
|
||||
)
|
||||
continue
|
||||
|
||||
plans.append(plan)
|
||||
if not allow_older_release_fallback or len(plans) >= release_limit:
|
||||
break
|
||||
except PrebuiltFallback:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise PrebuiltFallback(
|
||||
f"failed to inspect published releases in {repo}: {exc}"
|
||||
) from exc
|
||||
|
||||
if plans:
|
||||
return requested_tag, plans
|
||||
if last_error is not None:
|
||||
raise last_error
|
||||
raise PrebuiltFallback(
|
||||
f"no installable published llama.cpp releases were found in {repo}"
|
||||
)
|
||||
|
||||
|
||||
def normalized_requested_llama_tag(requested_tag: str | None) -> str:
|
||||
if isinstance(requested_tag, str):
|
||||
normalized = requested_tag.strip()
|
||||
|
|
@ -1435,7 +1887,7 @@ def iter_published_release_bundles(
|
|||
releases = (
|
||||
[github_release(repo, published_release_tag)]
|
||||
if published_release_tag
|
||||
else github_releases(repo)
|
||||
else github_releases(repo, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES)
|
||||
)
|
||||
for release in releases:
|
||||
if not published_release_tag and (
|
||||
|
|
@ -1669,7 +2121,9 @@ def latest_published_linux_cuda_tag(host: HostInfo, published_repo: str) -> str
|
|||
|
||||
|
||||
def iter_upstream_releases() -> Iterable[dict[str, Any]]:
|
||||
for release in github_releases(UPSTREAM_REPO):
|
||||
for release in github_releases(
|
||||
UPSTREAM_REPO, max_pages = DEFAULT_GITHUB_RELEASE_SCAN_MAX_PAGES
|
||||
):
|
||||
if release.get("draft") or release.get("prerelease"):
|
||||
continue
|
||||
yield release
|
||||
|
|
@ -2799,7 +3253,7 @@ def hydrate_source_tree(
|
|||
work_dir: Path,
|
||||
*,
|
||||
source_repo: str = UPSTREAM_REPO,
|
||||
expected_sha256: str,
|
||||
expected_sha256: str | None,
|
||||
source_label: str | None = None,
|
||||
exact_source: bool = False,
|
||||
) -> None:
|
||||
|
|
@ -3231,10 +3685,6 @@ def install_from_archives(
|
|||
) -> tuple[Path, Path]:
|
||||
main_archive = work_dir / choice.name
|
||||
log(f"downloading {choice.name} from {choice.source_label} release")
|
||||
if not choice.expected_sha256:
|
||||
raise PrebuiltFallback(
|
||||
f"approved checksum was missing for selected asset {choice.name}"
|
||||
)
|
||||
download_file_verified(
|
||||
choice.url,
|
||||
main_archive,
|
||||
|
|
@ -3962,7 +4412,7 @@ def require_approved_source_hash(
|
|||
|
||||
def preferred_source_archive(
|
||||
checksums: ApprovedReleaseChecksums, llama_tag: str
|
||||
) -> tuple[str, str, ApprovedArtifactHash, bool]:
|
||||
) -> tuple[str, str, ApprovedArtifactHash | None, bool]:
|
||||
exact_source = exact_source_archive_hash(checksums)
|
||||
exact_repo = repo_slug_from_source(checksums.source_repo) or repo_slug_from_source(
|
||||
checksums.source_repo_url
|
||||
|
|
@ -3974,7 +4424,7 @@ def preferred_source_archive(
|
|||
exact_source,
|
||||
True,
|
||||
)
|
||||
legacy = require_approved_source_hash(checksums, llama_tag)
|
||||
legacy = checksums.artifacts.get(source_archive_logical_name(llama_tag))
|
||||
return (
|
||||
UPSTREAM_REPO,
|
||||
llama_tag,
|
||||
|
|
@ -3990,6 +4440,8 @@ def selected_source_archive_metadata(
|
|||
_source_repo, _source_ref, source_archive, _exact_source = preferred_source_archive(
|
||||
checksums, llama_tag
|
||||
)
|
||||
if source_archive is None:
|
||||
return source_archive_logical_name(llama_tag), None
|
||||
return source_archive.asset_name, source_archive.sha256
|
||||
|
||||
|
||||
|
|
@ -4153,8 +4605,6 @@ def expected_install_fingerprint(
|
|||
choice: AssetChoice,
|
||||
approved_checksums: ApprovedReleaseChecksums,
|
||||
) -> str | None:
|
||||
if not choice.expected_sha256:
|
||||
return None
|
||||
source_asset_name, source_sha256 = selected_source_archive_metadata(
|
||||
approved_checksums,
|
||||
llama_tag,
|
||||
|
|
@ -4352,7 +4802,7 @@ def validate_prebuilt_choice(
|
|||
install_dir,
|
||||
work_dir,
|
||||
source_repo = source_repo,
|
||||
expected_sha256 = source_archive.sha256,
|
||||
expected_sha256 = source_archive.sha256 if source_archive is not None else None,
|
||||
source_label = (
|
||||
f"llama.cpp source tree for {source_repo}@{source_ref}"
|
||||
if exact_source
|
||||
|
|
@ -4477,7 +4927,12 @@ def validate_prebuilt_attempts(
|
|||
|
||||
|
||||
def install_prebuilt(
|
||||
install_dir: Path, llama_tag: str, published_repo: str, published_release_tag: str
|
||||
install_dir: Path,
|
||||
llama_tag: str,
|
||||
published_repo: str,
|
||||
published_release_tag: str,
|
||||
*,
|
||||
simple_policy: bool = False,
|
||||
) -> None:
|
||||
host = detect_host()
|
||||
choice: AssetChoice | None = None
|
||||
|
|
@ -4491,12 +4946,20 @@ def install_prebuilt(
|
|||
log(
|
||||
f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install"
|
||||
)
|
||||
requested_tag, release_plans = resolve_install_release_plans(
|
||||
llama_tag,
|
||||
host,
|
||||
published_repo,
|
||||
published_release_tag,
|
||||
)
|
||||
if simple_policy:
|
||||
requested_tag, release_plans = resolve_simple_install_release_plans(
|
||||
llama_tag,
|
||||
host,
|
||||
published_repo,
|
||||
published_release_tag,
|
||||
)
|
||||
else:
|
||||
requested_tag, release_plans = resolve_install_release_plans(
|
||||
llama_tag,
|
||||
host,
|
||||
published_repo,
|
||||
published_release_tag,
|
||||
)
|
||||
if release_plans and existing_install_matches_plan(
|
||||
install_dir, host, release_plans[0]
|
||||
):
|
||||
|
|
@ -4600,6 +5063,11 @@ def parse_args() -> argparse.Namespace:
|
|||
"until a usable published llama.cpp release bundle is found."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--simple-policy",
|
||||
action = "store_true",
|
||||
help = "Use the simplified platform-specific prebuilt selection policy.",
|
||||
)
|
||||
resolve_group = parser.add_mutually_exclusive_group()
|
||||
resolve_group.add_argument(
|
||||
"--resolve-llama-tag",
|
||||
|
|
@ -4719,6 +5187,7 @@ def main() -> int:
|
|||
llama_tag = args.llama_tag,
|
||||
published_repo = args.published_repo,
|
||||
published_release_tag = args.published_release_tag or "",
|
||||
simple_policy = args.simple_policy,
|
||||
)
|
||||
return EXIT_SUCCESS
|
||||
|
||||
|
|
|
|||
160
studio/setup.ps1
160
studio/setup.ps1
|
|
@ -34,6 +34,7 @@ $PackageDir = Split-Path -Parent $ScriptDir
|
|||
$DefaultLlamaPrForce = ""
|
||||
$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"
|
||||
$DefaultLlamaTag = "latest"
|
||||
$DefaultLlamaForceCompileRef = "master"
|
||||
|
||||
# Verbose can be enabled either by CLI flag or by UNSLOTH_VERBOSE=1.
|
||||
$script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq '1')
|
||||
|
|
@ -81,6 +82,31 @@ function New-UnslothTemporaryFile {
|
|||
return Get-Item -LiteralPath $tempPath
|
||||
}
|
||||
|
||||
function Get-InstalledLlamaPrebuiltRelease {
|
||||
param([string]$InstallDir)
|
||||
|
||||
$metadataPath = Join-Path $InstallDir "UNSLOTH_PREBUILT_INFO.json"
|
||||
if (-not (Test-Path $metadataPath)) {
|
||||
return $null
|
||||
}
|
||||
|
||||
try {
|
||||
$payload = Get-Content $metadataPath -Raw | ConvertFrom-Json
|
||||
} catch {
|
||||
return $null
|
||||
}
|
||||
|
||||
if (-not $payload.published_repo -or -not $payload.release_tag) {
|
||||
return $null
|
||||
}
|
||||
|
||||
$message = "installed release: $($payload.published_repo)@$($payload.release_tag)"
|
||||
if ($payload.tag -and $payload.tag -ne $payload.release_tag) {
|
||||
$message += " (tag $($payload.tag))"
|
||||
}
|
||||
return $message
|
||||
}
|
||||
|
||||
# Find nvcc on PATH, CUDA_PATH, or standard toolkit dirs.
|
||||
# Returns the path to nvcc.exe, or $null if not found.
|
||||
function Find-Nvcc {
|
||||
|
|
@ -133,7 +159,9 @@ function Find-Nvcc {
|
|||
|
||||
# 3. Scan standard toolkit directory
|
||||
if (Test-Path $toolkitBase) {
|
||||
$latest = Get-ChildItem -Directory $toolkitBase | Sort-Object Name | Select-Object -Last 1
|
||||
$latest = Get-ChildItem -Directory $toolkitBase | Where-Object {
|
||||
$_.Name -match '^v(\d+)\.(\d+)'
|
||||
} | Sort-Object { [version]($_.Name -replace '^v','') } -Descending | Select-Object -First 1
|
||||
if ($latest -and (Test-Path (Join-Path $latest.FullName 'bin\nvcc.exe'))) {
|
||||
return (Join-Path $latest.FullName 'bin\nvcc.exe')
|
||||
}
|
||||
|
|
@ -1618,6 +1646,7 @@ if ($LlamaSource.EndsWith('.git')) { $LlamaSource = $LlamaSource.Substring(0, $L
|
|||
$ResolvedSourceUrl = $LlamaSource
|
||||
$ResolvedSourceRef = $RequestedLlamaTag
|
||||
$ResolvedSourceRefKind = "tag"
|
||||
$ResolvedLlamaTag = $RequestedLlamaTag
|
||||
|
||||
if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
|
||||
$NeedLlamaSourceBuild = $true
|
||||
|
|
@ -1693,92 +1722,27 @@ if ($LlamaPr) {
|
|||
$ResolvedSourceRefKind = "pull"
|
||||
$NeedLlamaSourceBuild = $true
|
||||
$SkipPrebuiltInstall = $true
|
||||
} elseif ($SkipPrebuiltInstall) {
|
||||
# Custom source or other override already forced source build; skip the
|
||||
# prebuilt release resolution. When building from a custom fork, the fork
|
||||
# may not carry upstream bNNNN tags.
|
||||
if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
|
||||
$ResolvedLlamaTag = $RequestedLlamaTag
|
||||
} elseif ($LlamaSource -eq "https://github.com/ggml-org/llama.cpp") {
|
||||
$resolveTagArgs = @("--resolve-llama-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
|
||||
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveTagArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
|
||||
$fallbackResult = Invoke-LlamaHelper -Arguments $resolveTagArgs
|
||||
$fallbackOutput = $fallbackResult.Output
|
||||
$fallbackExit = $fallbackResult.ExitCode
|
||||
$ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) {
|
||||
try {
|
||||
(($fallbackOutput | Out-String) | ConvertFrom-Json).llama_tag
|
||||
} catch {
|
||||
$RequestedLlamaTag
|
||||
}
|
||||
} else {
|
||||
$RequestedLlamaTag
|
||||
}
|
||||
} else {
|
||||
$ResolvedLlamaTag = $RequestedLlamaTag
|
||||
}
|
||||
} else {
|
||||
$resolveInstallArgs = @("--resolve-install-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
|
||||
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveInstallArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
|
||||
$resolveErrorLog = New-UnslothTemporaryFile
|
||||
$resolveResult = Invoke-LlamaHelper -Arguments $resolveInstallArgs -StderrPath $resolveErrorLog
|
||||
$resolveOutput = $resolveResult.Output
|
||||
$resolveExit = $resolveResult.ExitCode
|
||||
$ResolvedLlamaTag = if ($resolveOutput) {
|
||||
try {
|
||||
(($resolveOutput | Out-String) | ConvertFrom-Json).llama_tag
|
||||
} catch {
|
||||
""
|
||||
}
|
||||
} else { "" }
|
||||
if ($resolveExit -ne 0 -or [string]::IsNullOrWhiteSpace($ResolvedLlamaTag)) {
|
||||
Write-Host ""
|
||||
substep "Failed to resolve a published llama.cpp release via $HelperReleaseRepo" "Yellow"
|
||||
Write-LlamaFailureLog -Output (Get-Content -Raw $resolveErrorLog)
|
||||
# Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
|
||||
# so the resolver prefers the latest usable Unsloth-published upstream tag
|
||||
# before falling back to the bleeding-edge ggml-org/llama.cpp tag.
|
||||
$resolveFallbackArgs = @("--resolve-llama-tag", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
|
||||
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveFallbackArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
|
||||
$fallbackResult = Invoke-LlamaHelper -Arguments $resolveFallbackArgs
|
||||
$fallbackOutput = $fallbackResult.Output
|
||||
$fallbackExit = $fallbackResult.ExitCode
|
||||
$ResolvedLlamaTag = if ($fallbackExit -eq 0 -and $fallbackOutput) {
|
||||
try {
|
||||
(($fallbackOutput | Out-String) | ConvertFrom-Json).llama_tag
|
||||
} catch {
|
||||
$RequestedLlamaTag
|
||||
}
|
||||
} else {
|
||||
$RequestedLlamaTag
|
||||
}
|
||||
$NeedLlamaSourceBuild = $true
|
||||
$SkipPrebuiltInstall = $true
|
||||
}
|
||||
Remove-Item $resolveErrorLog -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
|
||||
Write-Host ""
|
||||
substep "Resolved llama.cpp release tag: $ResolvedLlamaTag"
|
||||
|
||||
if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
|
||||
Write-Host ""
|
||||
substep "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt llama.cpp install" "Yellow"
|
||||
$NeedLlamaSourceBuild = $true
|
||||
} elseif ($SkipPrebuiltInstall) {
|
||||
Write-Host ""
|
||||
substep "Skipping prebuilt install -- falling back to source build" "Yellow"
|
||||
} else {
|
||||
Write-Host ""
|
||||
substep "installing prebuilt llama.cpp bundle (preferred path)..."
|
||||
if (Test-Path $LlamaCppDir) {
|
||||
substep "Existing llama.cpp install detected -- validating staged prebuilt update before replacement"
|
||||
}
|
||||
if ($SkipPrebuiltInstall) {
|
||||
substep "Skipping prebuilt install because prebuilt tag resolution failed -- falling back to source build" "Yellow"
|
||||
} else {
|
||||
$prebuiltArgs = @(
|
||||
$prebuiltArgs = @(
|
||||
"$PSScriptRoot\install_llama_prebuilt.py",
|
||||
"--install-dir", $LlamaCppDir,
|
||||
"--llama-tag", $RequestedLlamaTag,
|
||||
"--published-repo", $HelperReleaseRepo
|
||||
"--published-repo", $HelperReleaseRepo,
|
||||
"--simple-policy"
|
||||
)
|
||||
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
|
||||
$prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG)
|
||||
|
|
@ -1817,6 +1781,10 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
|
|||
} else {
|
||||
step "llama.cpp" "prebuilt installed and validated"
|
||||
}
|
||||
$installedRelease = Get-InstalledLlamaPrebuiltRelease -InstallDir $LlamaCppDir
|
||||
if ($installedRelease) {
|
||||
substep $installedRelease
|
||||
}
|
||||
} elseif ($prebuiltExit -eq 3) {
|
||||
step "llama.cpp" "install blocked by active llama.cpp process" "Yellow"
|
||||
Write-LlamaFailureLog -Output $prebuiltOutput
|
||||
|
|
@ -1834,7 +1802,6 @@ if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
|
|||
substep "Prebuilt llama.cpp path unavailable or failed validation -- falling back to source build" "Yellow"
|
||||
$NeedLlamaSourceBuild = $true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# ==========================================================================
|
||||
|
|
@ -1981,24 +1948,43 @@ if (-not $NeedLlamaSourceBuild) {
|
|||
}
|
||||
|
||||
if (-not $LlamaPr) {
|
||||
if ($LlamaSource -eq "https://github.com/ggml-org/llama.cpp") {
|
||||
$resolveSourceArgs = @("--resolve-source-build", $RequestedLlamaTag, "--published-repo", $HelperReleaseRepo, "--output-format", "json")
|
||||
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) { $resolveSourceArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG) }
|
||||
$sourcePlanResult = Invoke-LlamaHelper -Arguments $resolveSourceArgs
|
||||
$sourcePlanOutput = $sourcePlanResult.Output
|
||||
$sourcePlanExit = $sourcePlanResult.ExitCode
|
||||
if ($sourcePlanExit -eq 0 -and $sourcePlanOutput) {
|
||||
try {
|
||||
$sourcePlan = ($sourcePlanOutput | Out-String) | ConvertFrom-Json
|
||||
$ResolvedSourceUrl = $sourcePlan.source_url
|
||||
$ResolvedSourceRefKind = $sourcePlan.source_ref_kind
|
||||
$ResolvedSourceRef = $sourcePlan.source_ref
|
||||
} catch {
|
||||
$ResolvedSourceUrl = $LlamaSource
|
||||
if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {
|
||||
if ($RequestedLlamaTag -eq "latest") {
|
||||
$ResolvedSourceRef = if ($env:UNSLOTH_LLAMA_FORCE_COMPILE_REF) {
|
||||
$env:UNSLOTH_LLAMA_FORCE_COMPILE_REF
|
||||
} else {
|
||||
$DefaultLlamaForceCompileRef
|
||||
}
|
||||
$ResolvedSourceRefKind = "branch"
|
||||
} else {
|
||||
$ResolvedSourceRef = $RequestedLlamaTag
|
||||
$ResolvedSourceRefKind = "tag"
|
||||
}
|
||||
} elseif ($RequestedLlamaTag -eq "latest") {
|
||||
$resolveTagArgs = @("--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp", "--output-format", "json")
|
||||
$resolveTagResult = Invoke-LlamaHelper -Arguments $resolveTagArgs
|
||||
$resolveTagOutput = $resolveTagResult.Output
|
||||
$resolveTagExit = $resolveTagResult.ExitCode
|
||||
if ($resolveTagExit -eq 0 -and $resolveTagOutput) {
|
||||
try {
|
||||
$ResolvedSourceRef = (($resolveTagOutput | Out-String) | ConvertFrom-Json).llama_tag
|
||||
} catch {
|
||||
$ResolvedSourceRef = ""
|
||||
}
|
||||
} else {
|
||||
$ResolvedSourceRef = ""
|
||||
}
|
||||
if ([string]::IsNullOrWhiteSpace($ResolvedSourceRef)) {
|
||||
$ResolvedSourceRef = "latest"
|
||||
}
|
||||
$ResolvedSourceRefKind = "tag"
|
||||
} else {
|
||||
$ResolvedSourceRef = $RequestedLlamaTag
|
||||
$ResolvedSourceRefKind = "tag"
|
||||
}
|
||||
if ([string]::IsNullOrWhiteSpace($ResolvedSourceUrl)) { $ResolvedSourceUrl = $LlamaSource }
|
||||
if ([string]::IsNullOrWhiteSpace($ResolvedSourceRef)) { $ResolvedSourceRef = $ResolvedLlamaTag }
|
||||
if ([string]::IsNullOrWhiteSpace($ResolvedSourceRef)) { $ResolvedSourceRef = $RequestedLlamaTag }
|
||||
}
|
||||
|
||||
# -- Step A: Clone or pull llama.cpp --
|
||||
|
|
|
|||
272
studio/setup.sh
272
studio/setup.sh
|
|
@ -25,6 +25,7 @@ RULE=$(printf '\342\224\200%.0s' {1..52})
|
|||
_DEFAULT_LLAMA_PR_FORCE=""
|
||||
_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"
|
||||
_DEFAULT_LLAMA_TAG="latest"
|
||||
_DEFAULT_LLAMA_FORCE_COMPILE_REF="master"
|
||||
|
||||
# ── Colors (same palette as startup_banner / install_python_stack) ──
|
||||
if [ -n "${NO_COLOR:-}" ]; then
|
||||
|
|
@ -121,6 +122,45 @@ print_llama_error_log() {
|
|||
tail -n 120 "$log_file" | sed 's/^/ | /' >&2
|
||||
}
|
||||
|
||||
installed_llama_prebuilt_release() {
|
||||
local install_dir=${1:-}
|
||||
local metadata_path="$install_dir/UNSLOTH_PREBUILT_INFO.json"
|
||||
[ -f "$metadata_path" ] || return 0
|
||||
python - "$metadata_path" <<'PY' 2>/dev/null || true
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
payload = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
raise SystemExit(0)
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
raise SystemExit(0)
|
||||
|
||||
repo = str(payload.get("published_repo") or "").strip()
|
||||
release_tag = str(payload.get("release_tag") or "").strip()
|
||||
llama_tag = str(payload.get("tag") or "").strip()
|
||||
if not repo or not release_tag:
|
||||
raise SystemExit(0)
|
||||
|
||||
message = f"installed release: {repo}@{release_tag}"
|
||||
if llama_tag and llama_tag != release_tag:
|
||||
message += f" (tag {llama_tag})"
|
||||
print(message)
|
||||
PY
|
||||
}
|
||||
|
||||
print_installed_llama_prebuilt_release() {
|
||||
local install_dir=${1:-}
|
||||
local installed_release
|
||||
installed_release="$(installed_llama_prebuilt_release "$install_dir")"
|
||||
if [ -n "$installed_release" ]; then
|
||||
substep "$installed_release"
|
||||
fi
|
||||
}
|
||||
|
||||
# ── Banner ──
|
||||
echo ""
|
||||
printf " ${C_TITLE}%s${C_RST}\n" "🦥 Unsloth Studio Setup"
|
||||
|
|
@ -487,30 +527,27 @@ _NEED_LLAMA_SOURCE_BUILD=false
|
|||
_LLAMA_CPP_DEGRADED=false
|
||||
_LLAMA_FORCE_COMPILE="${UNSLOTH_LLAMA_FORCE_COMPILE:-0}"
|
||||
_REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}"
|
||||
# Force all installs to use mainline llama.cpp from ggml-org.
|
||||
_HELPER_RELEASE_REPO="ggml-org/llama.cpp"
|
||||
_HOST_SYSTEM="$(uname -s 2>/dev/null || true)"
|
||||
if [ "$_HOST_SYSTEM" = "Darwin" ]; then
|
||||
_HELPER_RELEASE_REPO="ggml-org/llama.cpp"
|
||||
else
|
||||
_HELPER_RELEASE_REPO="unslothai/llama.cpp"
|
||||
fi
|
||||
_LLAMA_PR="${UNSLOTH_LLAMA_PR:-}"
|
||||
|
||||
_SKIP_PREBUILT_INSTALL=false
|
||||
_LLAMA_PR_FORCE="${UNSLOTH_LLAMA_PR_FORCE:-${_DEFAULT_LLAMA_PR_FORCE}}"
|
||||
# Force mainline source -- no env var override for now.
|
||||
_LLAMA_SOURCE="${_DEFAULT_LLAMA_SOURCE}"
|
||||
_LLAMA_SOURCE="${_LLAMA_SOURCE%.git}" # normalize: strip trailing .git
|
||||
_RESOLVED_SOURCE_URL="$_LLAMA_SOURCE"
|
||||
_RESOLVED_SOURCE_REF="$_REQUESTED_LLAMA_TAG"
|
||||
_RESOLVED_SOURCE_REF_KIND="tag"
|
||||
_RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
|
||||
|
||||
if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
|
||||
_NEED_LLAMA_SOURCE_BUILD=true
|
||||
_SKIP_PREBUILT_INSTALL=true
|
||||
fi
|
||||
|
||||
# Non-default source URL forces source build (fork has different code than prebuilt).
|
||||
if [ "$_LLAMA_SOURCE" != "https://github.com/ggml-org/llama.cpp" ]; then
|
||||
step "llama.cpp" "custom source: $_LLAMA_SOURCE -- forcing source build" "$C_WARN"
|
||||
_NEED_LLAMA_SOURCE_BUILD=true
|
||||
_SKIP_PREBUILT_INSTALL=true
|
||||
fi
|
||||
|
||||
# Baked-in PR_FORCE promotes to _LLAMA_PR when user hasn't set one.
|
||||
if [ -z "$_LLAMA_PR" ] && [ -n "$_LLAMA_PR_FORCE" ] && \
|
||||
[[ "$_LLAMA_PR_FORCE" =~ ^[0-9]+$ ]] && [ "$_LLAMA_PR_FORCE" -gt 0 ]; then
|
||||
|
|
@ -530,149 +567,68 @@ if [ -n "$_LLAMA_PR" ]; then
|
|||
_RESOLVED_SOURCE_REF_KIND="pull"
|
||||
_NEED_LLAMA_SOURCE_BUILD=true
|
||||
_SKIP_PREBUILT_INSTALL=true
|
||||
elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then
|
||||
# Custom source or other override already forced source build; skip
|
||||
# the prebuilt release resolution entirely. When building from a custom
|
||||
# fork, the fork may not carry upstream bNNNN tags, so resolve the tag
|
||||
# only when the source is the default ggml-org repo.
|
||||
if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
|
||||
_RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
|
||||
elif [ "$_LLAMA_SOURCE" = "https://github.com/ggml-org/llama.cpp" ]; then
|
||||
_RESOLVE_TAG_ARGS=(--resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
|
||||
_RESOLVE_TAG_ARGS+=(--output-format json)
|
||||
if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
|
||||
_RESOLVE_TAG_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
|
||||
fi
|
||||
set +e
|
||||
_RESOLVE_TAG_JSON="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_TAG_ARGS[@]}" 2>/dev/null)"
|
||||
_RESOLVE_UPSTREAM_STATUS=$?
|
||||
set -e
|
||||
if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "${_RESOLVE_TAG_JSON:-}" ]; then
|
||||
_RESOLVED_LLAMA_TAG="$(
|
||||
printf '%s' "$_RESOLVE_TAG_JSON" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true
|
||||
)"
|
||||
else
|
||||
_RESOLVED_LLAMA_TAG=""
|
||||
fi
|
||||
if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
|
||||
_RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
|
||||
fi
|
||||
else
|
||||
_RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
|
||||
fi
|
||||
else
|
||||
_RESOLVE_INSTALL_ARGS=(--resolve-install-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
|
||||
_RESOLVE_INSTALL_ARGS+=(--output-format json)
|
||||
if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
|
||||
_RESOLVE_INSTALL_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
|
||||
fi
|
||||
_RESOLVE_LLAMA_LOG="$(mktemp)"
|
||||
set +e
|
||||
_RESOLVE_INSTALL_JSON="$(
|
||||
python "$SCRIPT_DIR/install_llama_prebuilt.py" \
|
||||
"${_RESOLVE_INSTALL_ARGS[@]}" 2>"$_RESOLVE_LLAMA_LOG"
|
||||
)"
|
||||
_RESOLVE_LLAMA_STATUS=$?
|
||||
set -e
|
||||
if [ "$_RESOLVE_LLAMA_STATUS" -eq 0 ]; then
|
||||
_RESOLVED_LLAMA_TAG="$(
|
||||
printf '%s' "${_RESOLVE_INSTALL_JSON:-}" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true
|
||||
)"
|
||||
else
|
||||
_RESOLVED_LLAMA_TAG=""
|
||||
fi
|
||||
if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
|
||||
step "llama.cpp" "failed to resolve a published llama.cpp release via $_HELPER_RELEASE_REPO" "$C_WARN"
|
||||
print_llama_error_log "$_RESOLVE_LLAMA_LOG"
|
||||
set +e
|
||||
# Resolve the llama.cpp tag for source-build fallback. Pass --published-repo
|
||||
# so the resolver prefers the latest usable Unsloth-published upstream tag
|
||||
# before falling back to the bleeding-edge ggml-org/llama.cpp tag.
|
||||
_RESOLVE_FALLBACK_ARGS=(--resolve-llama-tag "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
|
||||
_RESOLVE_FALLBACK_ARGS+=(--output-format json)
|
||||
if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
|
||||
_RESOLVE_FALLBACK_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
|
||||
fi
|
||||
_RESOLVE_FALLBACK_JSON="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_FALLBACK_ARGS[@]}" 2>/dev/null)"
|
||||
_RESOLVE_UPSTREAM_STATUS=$?
|
||||
set -e
|
||||
if [ "$_RESOLVE_UPSTREAM_STATUS" -eq 0 ] && [ -n "${_RESOLVE_FALLBACK_JSON:-}" ]; then
|
||||
_RESOLVED_LLAMA_TAG="$(
|
||||
printf '%s' "$_RESOLVE_FALLBACK_JSON" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true
|
||||
)"
|
||||
else
|
||||
_RESOLVED_LLAMA_TAG=""
|
||||
fi
|
||||
if [ -z "$_RESOLVED_LLAMA_TAG" ]; then
|
||||
_RESOLVED_LLAMA_TAG="$_REQUESTED_LLAMA_TAG"
|
||||
fi
|
||||
_NEED_LLAMA_SOURCE_BUILD=true
|
||||
_SKIP_PREBUILT_INSTALL=true
|
||||
fi
|
||||
rm -f "$_RESOLVE_LLAMA_LOG"
|
||||
fi
|
||||
|
||||
substep "resolved llama.cpp tag: $_RESOLVED_LLAMA_TAG"
|
||||
verbose_substep "requested llama.cpp tag: $_REQUESTED_LLAMA_TAG (repo: $_HELPER_RELEASE_REPO)"
|
||||
|
||||
if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
|
||||
step "llama.cpp" "UNSLOTH_LLAMA_FORCE_COMPILE=1 -- skipping prebuilt" "$C_WARN"
|
||||
_NEED_LLAMA_SOURCE_BUILD=true
|
||||
elif [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then
|
||||
substep "prebuilt install skipped -- falling back to source build"
|
||||
else
|
||||
substep "installing prebuilt llama.cpp..."
|
||||
if [ -d "$LLAMA_CPP_DIR" ]; then
|
||||
substep "existing install detected -- validating update"
|
||||
fi
|
||||
if [ "${_SKIP_PREBUILT_INSTALL:-false}" = true ]; then
|
||||
substep "prebuilt tag resolution failed -- falling back to source build"
|
||||
_PREBUILT_CMD=(
|
||||
python "$SCRIPT_DIR/install_llama_prebuilt.py"
|
||||
--install-dir "$LLAMA_CPP_DIR"
|
||||
--llama-tag "$_REQUESTED_LLAMA_TAG"
|
||||
--published-repo "$_HELPER_RELEASE_REPO"
|
||||
--simple-policy
|
||||
)
|
||||
if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
|
||||
_PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
|
||||
fi
|
||||
_PREBUILT_LOG="$(mktemp)"
|
||||
set +e
|
||||
if _is_verbose; then
|
||||
"${_PREBUILT_CMD[@]}" 2>&1 | tee "$_PREBUILT_LOG"
|
||||
_PREBUILT_STATUS=${PIPESTATUS[0]}
|
||||
else
|
||||
_PREBUILT_CMD=(
|
||||
python "$SCRIPT_DIR/install_llama_prebuilt.py"
|
||||
--install-dir "$LLAMA_CPP_DIR"
|
||||
--llama-tag "$_REQUESTED_LLAMA_TAG"
|
||||
--published-repo "$_HELPER_RELEASE_REPO"
|
||||
)
|
||||
if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
|
||||
_PREBUILT_CMD+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
|
||||
fi
|
||||
_PREBUILT_LOG="$(mktemp)"
|
||||
set +e
|
||||
if _is_verbose; then
|
||||
"${_PREBUILT_CMD[@]}" 2>&1 | tee "$_PREBUILT_LOG"
|
||||
_PREBUILT_STATUS=${PIPESTATUS[0]}
|
||||
else
|
||||
"${_PREBUILT_CMD[@]}" >"$_PREBUILT_LOG" 2>&1
|
||||
_PREBUILT_STATUS=$?
|
||||
fi
|
||||
set -e
|
||||
"${_PREBUILT_CMD[@]}" >"$_PREBUILT_LOG" 2>&1
|
||||
_PREBUILT_STATUS=$?
|
||||
fi
|
||||
set -e
|
||||
|
||||
if [ "$_PREBUILT_STATUS" -eq 0 ]; then
|
||||
if grep -Fq "already matches" "$_PREBUILT_LOG"; then
|
||||
step "llama.cpp" "prebuilt up to date and validated"
|
||||
else
|
||||
step "llama.cpp" "prebuilt installed and validated"
|
||||
fi
|
||||
verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR"
|
||||
rm -f "$_PREBUILT_LOG"
|
||||
elif [ "$_PREBUILT_STATUS" -eq 3 ]; then
|
||||
step "llama.cpp" "install blocked by active llama.cpp process" "$C_WARN"
|
||||
print_llama_error_log "$_PREBUILT_LOG"
|
||||
rm -f "$_PREBUILT_LOG"
|
||||
if [ -d "$LLAMA_CPP_DIR" ]; then
|
||||
substep "existing install was restored"
|
||||
fi
|
||||
substep "close Studio or other llama.cpp users and retry"
|
||||
exit 3
|
||||
if [ "$_PREBUILT_STATUS" -eq 0 ]; then
|
||||
if grep -Fq "already matches" "$_PREBUILT_LOG"; then
|
||||
step "llama.cpp" "prebuilt up to date and validated"
|
||||
else
|
||||
step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
|
||||
print_llama_error_log "$_PREBUILT_LOG"
|
||||
rm -f "$_PREBUILT_LOG"
|
||||
if [ -d "$LLAMA_CPP_DIR" ]; then
|
||||
substep "prebuilt update failed; existing install restored"
|
||||
fi
|
||||
substep "falling back to source build"
|
||||
_NEED_LLAMA_SOURCE_BUILD=true
|
||||
step "llama.cpp" "prebuilt installed and validated"
|
||||
fi
|
||||
print_installed_llama_prebuilt_release "$LLAMA_CPP_DIR"
|
||||
verbose_substep "llama.cpp install dir: $LLAMA_CPP_DIR"
|
||||
rm -f "$_PREBUILT_LOG"
|
||||
elif [ "$_PREBUILT_STATUS" -eq 3 ]; then
|
||||
step "llama.cpp" "install blocked by active llama.cpp process" "$C_WARN"
|
||||
print_llama_error_log "$_PREBUILT_LOG"
|
||||
rm -f "$_PREBUILT_LOG"
|
||||
if [ -d "$LLAMA_CPP_DIR" ]; then
|
||||
substep "existing install was restored"
|
||||
fi
|
||||
substep "close Studio or other llama.cpp users and retry"
|
||||
exit 3
|
||||
else
|
||||
step "llama.cpp" "prebuilt install failed (continuing)" "$C_WARN"
|
||||
print_llama_error_log "$_PREBUILT_LOG"
|
||||
rm -f "$_PREBUILT_LOG"
|
||||
if [ -d "$LLAMA_CPP_DIR" ]; then
|
||||
substep "prebuilt update failed; existing install restored"
|
||||
fi
|
||||
substep "falling back to source build"
|
||||
_NEED_LLAMA_SOURCE_BUILD=true
|
||||
fi
|
||||
fi
|
||||
|
||||
|
|
@ -746,33 +702,41 @@ else
|
|||
[ -f "$LLAMA_SERVER_BIN" ] || _LLAMA_CPP_DEGRADED=true
|
||||
else
|
||||
if [ -z "$_LLAMA_PR" ]; then
|
||||
if [ "$_LLAMA_SOURCE" = "https://github.com/ggml-org/llama.cpp" ]; then
|
||||
_RESOLVE_SOURCE_ARGS=(--resolve-source-build "$_REQUESTED_LLAMA_TAG" --published-repo "$_HELPER_RELEASE_REPO")
|
||||
_RESOLVE_SOURCE_ARGS+=(--output-format json)
|
||||
if [ -n "${UNSLOTH_LLAMA_RELEASE_TAG:-}" ]; then
|
||||
_RESOLVE_SOURCE_ARGS+=(--published-release-tag "$UNSLOTH_LLAMA_RELEASE_TAG")
|
||||
_RESOLVED_SOURCE_URL="$_LLAMA_SOURCE"
|
||||
if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then
|
||||
if [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then
|
||||
_RESOLVED_SOURCE_REF="${UNSLOTH_LLAMA_FORCE_COMPILE_REF:-${_DEFAULT_LLAMA_FORCE_COMPILE_REF}}"
|
||||
_RESOLVED_SOURCE_REF_KIND="branch"
|
||||
else
|
||||
_RESOLVED_SOURCE_REF="$_REQUESTED_LLAMA_TAG"
|
||||
_RESOLVED_SOURCE_REF_KIND="tag"
|
||||
fi
|
||||
elif [ "$_REQUESTED_LLAMA_TAG" = "latest" ]; then
|
||||
_RESOLVE_TAG_ARGS=(--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp" --output-format json)
|
||||
set +e
|
||||
_SOURCE_BUILD_PLAN="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_SOURCE_ARGS[@]}" 2>/dev/null)"
|
||||
_RESOLVE_SOURCE_STATUS=$?
|
||||
_RESOLVE_TAG_JSON="$(python "$SCRIPT_DIR/install_llama_prebuilt.py" "${_RESOLVE_TAG_ARGS[@]}" 2>/dev/null)"
|
||||
_RESOLVE_TAG_STATUS=$?
|
||||
set -e
|
||||
if [ "$_RESOLVE_SOURCE_STATUS" -eq 0 ] && [ -n "$_SOURCE_BUILD_PLAN" ]; then
|
||||
_RESOLVED_SOURCE_URL="$(
|
||||
printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_url",""))' 2>/dev/null || true
|
||||
)"
|
||||
_RESOLVED_SOURCE_REF_KIND="$(
|
||||
printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_ref_kind",""))' 2>/dev/null || true
|
||||
)"
|
||||
if [ "$_RESOLVE_TAG_STATUS" -eq 0 ] && [ -n "${_RESOLVE_TAG_JSON:-}" ]; then
|
||||
_RESOLVED_SOURCE_REF="$(
|
||||
printf '%s' "$_SOURCE_BUILD_PLAN" | python -c 'import json,sys; print(json.load(sys.stdin).get("source_ref",""))' 2>/dev/null || true
|
||||
printf '%s' "$_RESOLVE_TAG_JSON" | python -c 'import json,sys; print(json.load(sys.stdin).get("llama_tag",""))' 2>/dev/null || true
|
||||
)"
|
||||
else
|
||||
_RESOLVED_SOURCE_REF=""
|
||||
fi
|
||||
if [ -z "$_RESOLVED_SOURCE_REF" ]; then
|
||||
_RESOLVED_SOURCE_REF="latest"
|
||||
fi
|
||||
_RESOLVED_SOURCE_REF_KIND="tag"
|
||||
else
|
||||
_RESOLVED_SOURCE_REF="$_REQUESTED_LLAMA_TAG"
|
||||
_RESOLVED_SOURCE_REF_KIND="tag"
|
||||
fi
|
||||
if [ -z "$_RESOLVED_SOURCE_URL" ]; then
|
||||
_RESOLVED_SOURCE_URL="$_LLAMA_SOURCE"
|
||||
fi
|
||||
if [ -z "$_RESOLVED_SOURCE_REF" ]; then
|
||||
_RESOLVED_SOURCE_REF="$_RESOLVED_LLAMA_TAG"
|
||||
_RESOLVED_SOURCE_REF="$_REQUESTED_LLAMA_TAG"
|
||||
fi
|
||||
fi
|
||||
verbose_substep "source build repo: $_RESOLVED_SOURCE_URL"
|
||||
|
|
|
|||
|
|
@ -14,10 +14,12 @@ Run: pytest tests/studio/install/test_pr4562_bugfixes.py -v
|
|||
"""
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import urllib.parse
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
|
@ -347,6 +349,47 @@ class TestResolveRequestedLlamaTag:
|
|||
}
|
||||
|
||||
|
||||
class TestFetchJsonRetries:
|
||||
def test_fetch_json_retries_invalid_github_api_json(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
calls = {"count": 0}
|
||||
|
||||
def fake_download_bytes(url, **kwargs):
|
||||
calls["count"] += 1
|
||||
if calls["count"] == 1:
|
||||
return b'{"incomplete":"payload'
|
||||
return json.dumps([{"tag_name": "b8635"}]).encode("utf-8")
|
||||
|
||||
monkeypatch.setattr(MOD, "download_bytes", fake_download_bytes)
|
||||
monkeypatch.setattr(MOD, "sleep_backoff", lambda _attempt: None)
|
||||
|
||||
payload = MOD.fetch_json(
|
||||
"https://api.github.com/repos/ggml-org/llama.cpp/releases?per_page=100&page=1"
|
||||
)
|
||||
|
||||
assert isinstance(payload, list)
|
||||
assert payload[0]["tag_name"] == "b8635"
|
||||
assert calls["count"] == 2
|
||||
|
||||
def test_github_releases_honors_max_pages(self, monkeypatch: pytest.MonkeyPatch):
|
||||
seen_pages: list[int] = []
|
||||
|
||||
def fake_fetch_json(url: str):
|
||||
parsed = urllib.parse.urlparse(url)
|
||||
params = urllib.parse.parse_qs(parsed.query)
|
||||
page = int(params["page"][0])
|
||||
seen_pages.append(page)
|
||||
return [{"tag_name": f"b{page:04d}"} for _ in range(100)]
|
||||
|
||||
monkeypatch.setattr(MOD, "fetch_json", fake_fetch_json)
|
||||
|
||||
releases = MOD.github_releases("ggml-org/llama.cpp", max_pages = 2)
|
||||
|
||||
assert seen_pages == [1, 2]
|
||||
assert len(releases) == 200
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# TEST GROUP C: setup.sh logic (bash subprocess tests)
|
||||
# =========================================================================
|
||||
|
|
@ -643,25 +686,35 @@ class TestSourceCodePatterns:
|
|||
'_RESOLVED_SOURCE_REF" != "latest"' in content
|
||||
), "Should guard against literal 'latest' tag"
|
||||
|
||||
def test_setup_sh_source_build_uses_helper_resolution(self):
|
||||
"""Shell source fallback should consult the helper for repo/ref planning."""
|
||||
def test_setup_sh_source_build_uses_helper_latest_tag_only(self):
|
||||
"""Shell source fallback should only use helper latest-tag resolution."""
|
||||
content = SETUP_SH.read_text()
|
||||
assert "--resolve-source-build" in content
|
||||
assert "--resolve-source-build" not in content
|
||||
assert "--resolve-install-tag" not in content
|
||||
assert (
|
||||
'--resolve-llama-tag latest --published-repo "ggml-org/llama.cpp"'
|
||||
in content
|
||||
)
|
||||
assert "--output-format json" in content
|
||||
assert "_RESOLVED_SOURCE_URL" in content
|
||||
assert "_RESOLVED_SOURCE_REF_KIND" in content
|
||||
assert "_RESOLVED_SOURCE_REF" in content
|
||||
|
||||
def test_setup_sh_latest_resolution_uses_helper_only(self):
|
||||
"""Shell fallback should rely on helper output, not raw GitHub API tag_name."""
|
||||
def test_setup_sh_prebuilt_install_uses_simple_policy_only(self):
|
||||
"""Shell prebuilt path should use the simplified helper install entrypoint."""
|
||||
content = SETUP_SH.read_text()
|
||||
assert "--resolve-install-tag" in content
|
||||
assert "--resolve-llama-tag" in content
|
||||
assert 'tail -n 1 "$_RESOLVE_LLAMA_LOG"' not in content
|
||||
assert "json.load" in content
|
||||
assert "--simple-policy" in content
|
||||
assert "--resolve-install-tag" not in content
|
||||
assert "_HELPER_RELEASE_REPO}/releases/latest" not in content
|
||||
assert "ggml-org/llama.cpp/releases/latest" not in content
|
||||
|
||||
def test_setup_sh_reports_installed_prebuilt_release(self):
|
||||
"""Shell wrapper should report the installed prebuilt release from metadata."""
|
||||
content = SETUP_SH.read_text()
|
||||
assert "UNSLOTH_PREBUILT_INFO.json" in content
|
||||
assert "installed release:" in content
|
||||
assert 'print_installed_llama_prebuilt_release "$LLAMA_CPP_DIR"' in content
|
||||
|
||||
def test_setup_sh_macos_arm64_uses_metal_flags(self):
|
||||
"""Apple Silicon source builds should explicitly enable Metal like upstream."""
|
||||
content = SETUP_SH.read_text()
|
||||
|
|
@ -759,20 +812,34 @@ class TestSourceCodePatterns:
|
|||
f"Found 'git pull' in llama.cpp build section at line {i+1}"
|
||||
)
|
||||
|
||||
def test_setup_ps1_latest_resolution_uses_helper_only(self):
|
||||
"""PS1 fallback should rely on helper output, not raw GitHub API tag_name."""
|
||||
def test_setup_ps1_prebuilt_install_uses_simple_policy_only(self):
|
||||
"""PS1 prebuilt path should use the simplified helper install entrypoint."""
|
||||
content = SETUP_PS1.read_text()
|
||||
assert "--resolve-install-tag" in content
|
||||
assert "--resolve-llama-tag" in content
|
||||
assert '--output-format", "json"' in content
|
||||
assert "ConvertFrom-Json" in content
|
||||
assert '"--simple-policy"' in content
|
||||
assert "--resolve-install-tag" not in content
|
||||
assert "$HelperReleaseRepo/releases/latest" not in content
|
||||
assert "ggml-org/llama.cpp/releases/latest" not in content
|
||||
|
||||
def test_setup_ps1_source_build_uses_helper_resolution(self):
|
||||
"""PS1 source fallback should consult the helper for repo/ref planning."""
|
||||
def test_setup_ps1_reports_installed_prebuilt_release(self):
|
||||
"""PS1 wrapper should report the installed prebuilt release from metadata."""
|
||||
content = SETUP_PS1.read_text()
|
||||
assert "--resolve-source-build" in content
|
||||
assert "Get-InstalledLlamaPrebuiltRelease" in content
|
||||
assert "UNSLOTH_PREBUILT_INFO.json" in content
|
||||
assert "installed release:" in content
|
||||
assert (
|
||||
"$installedRelease = Get-InstalledLlamaPrebuiltRelease -InstallDir $LlamaCppDir"
|
||||
in content
|
||||
)
|
||||
|
||||
def test_setup_ps1_source_build_uses_helper_latest_tag_only(self):
|
||||
"""PS1 source fallback should only use helper latest-tag resolution."""
|
||||
content = SETUP_PS1.read_text()
|
||||
assert "--resolve-source-build" not in content
|
||||
assert "--resolve-install-tag" not in content
|
||||
assert (
|
||||
'"--resolve-llama-tag", "latest", "--published-repo", "ggml-org/llama.cpp"'
|
||||
in content
|
||||
)
|
||||
assert '--output-format", "json"' in content
|
||||
assert "$ResolvedSourceUrl" in content
|
||||
assert "$ResolvedSourceRefKind" in content
|
||||
|
|
@ -794,18 +861,25 @@ class TestSourceCodePatterns:
|
|||
"""Helper resolution should suppress terminating NativeCommandError on PS 5.1."""
|
||||
content = SETUP_PS1.read_text()
|
||||
helper_idx = content.index("function Invoke-LlamaHelper")
|
||||
block = content[helper_idx : helper_idx + 1200]
|
||||
block = content[helper_idx : helper_idx + 2200]
|
||||
assert "$previousErrorActionPreference = $ErrorActionPreference" in block
|
||||
assert '$ErrorActionPreference = "Continue"' in block
|
||||
assert "$ErrorActionPreference = $previousErrorActionPreference" in block
|
||||
|
||||
def test_setup_ps1_uses_local_tempfile_helper(self):
|
||||
"""PS1 should not depend on New-TemporaryFile being available."""
|
||||
"""PS1 should not depend on New-TemporaryFile being available anywhere."""
|
||||
content = SETUP_PS1.read_text()
|
||||
assert "function New-UnslothTemporaryFile" in content
|
||||
assert "$resolveErrorLog = New-UnslothTemporaryFile" in content
|
||||
assert "$resolveErrorLog = New-TemporaryFile" not in content
|
||||
|
||||
def test_setup_ps1_find_nvcc_uses_version_sort_for_latest_toolkit(self):
|
||||
"""The unconstrained nvcc fallback should not sort toolkit dirs lexicographically."""
|
||||
content = SETUP_PS1.read_text()
|
||||
assert "Sort-Object Name | Select-Object -Last 1" not in content
|
||||
assert (
|
||||
"Sort-Object { [version]($_.Name -replace '^v','') } -Descending" in content
|
||||
)
|
||||
|
||||
def test_binary_env_linux_has_binary_parent(self):
|
||||
"""The Linux branch of binary_env should include binary_path.parent."""
|
||||
content = MODULE_PATH.read_text()
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@
|
|||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
__version__ = "2026.4.1"
|
||||
__version__ = "2026.4.2"
|
||||
|
||||
__all__ = [
|
||||
"SUPPORTS_BFLOAT16",
|
||||
|
|
|
|||
|
|
@ -114,6 +114,17 @@ FORCE_FLOAT32 = [
|
|||
|
||||
global DISABLE_COMPILE_MODEL_NAMES
|
||||
# Must be alphabetically sorted for each entry
|
||||
|
||||
|
||||
def _strip_unsloth_bnb_4bit_suffix(model_name: str) -> str:
|
||||
"""Remove Unsloth 4bit suffixes without lowercasing (HF cache dirs are case-sensitive)."""
|
||||
s = model_name
|
||||
for suffix in ("-unsloth-bnb-4bit", "-bnb-4bit"):
|
||||
if len(s) >= len(suffix) and s.lower().endswith(suffix.lower()):
|
||||
s = s[: -len(suffix)]
|
||||
return s
|
||||
|
||||
|
||||
DISABLE_COMPILE_MODEL_NAMES = [
|
||||
"aya_vision",
|
||||
"modernbert",
|
||||
|
|
@ -404,8 +415,7 @@ class FastLanguageModel(FastLlamaModel):
|
|||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
|
||||
model_name = model_name.lower().removesuffix("-bnb-4bit")
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
|
|
@ -551,8 +561,7 @@ class FastLanguageModel(FastLlamaModel):
|
|||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
|
||||
model_name = model_name.lower().removesuffix("-bnb-4bit")
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
|
|
@ -1019,8 +1028,7 @@ class FastModel(FastBaseModel):
|
|||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
|
||||
model_name = model_name.lower().removesuffix("-bnb-4bit")
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
|
|
@ -1320,8 +1328,7 @@ class FastModel(FastBaseModel):
|
|||
if not ALLOW_PREQUANTIZED_MODELS and model_name.lower().endswith(
|
||||
("-unsloth-bnb-4bit", "-bnb-4bit")
|
||||
):
|
||||
model_name = model_name.lower().removesuffix("-unsloth-bnb-4bit")
|
||||
model_name = model_name.lower().removesuffix("-bnb-4bit")
|
||||
model_name = _strip_unsloth_bnb_4bit_suffix(model_name)
|
||||
# Change -BF16 to all False for 4bit, 8bit etc
|
||||
if model_name.lower().endswith("-bf16"):
|
||||
load_in_4bit = False
|
||||
|
|
@ -1533,14 +1540,72 @@ class FastModel(FastBaseModel):
|
|||
if is_peft:
|
||||
# From https://github.com/huggingface/peft/issues/184
|
||||
# Now add PEFT adapters
|
||||
model = PeftModel.from_pretrained(
|
||||
model,
|
||||
old_model_name,
|
||||
token = token,
|
||||
revision = revision,
|
||||
is_trainable = True,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
|
||||
# Gemma4 ClippableLinear wraps nn.Linear -- PEFT can't inject LoRA
|
||||
# on it directly. Monkey-patch PEFT to target the inner .linear
|
||||
# child instead (same patch as vision.py training path).
|
||||
# See https://github.com/huggingface/peft/issues/3129
|
||||
_clippable_linear_cls = None
|
||||
try:
|
||||
from transformers.models.gemma4.modeling_gemma4 import (
|
||||
Gemma4ClippableLinear as _clippable_linear_cls,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if _clippable_linear_cls is not None:
|
||||
from peft.tuners.lora.model import LoraModel as _LoraModel
|
||||
|
||||
_original_car = _LoraModel._create_and_replace
|
||||
|
||||
def _patched_car(
|
||||
self,
|
||||
peft_config,
|
||||
adapter_name,
|
||||
target,
|
||||
target_name,
|
||||
parent,
|
||||
current_key = None,
|
||||
**kwargs,
|
||||
):
|
||||
if isinstance(target, _clippable_linear_cls):
|
||||
return _original_car(
|
||||
self,
|
||||
peft_config,
|
||||
adapter_name,
|
||||
target.linear,
|
||||
"linear",
|
||||
target,
|
||||
current_key = current_key,
|
||||
**kwargs,
|
||||
)
|
||||
return _original_car(
|
||||
self,
|
||||
peft_config,
|
||||
adapter_name,
|
||||
target,
|
||||
target_name,
|
||||
parent,
|
||||
current_key = current_key,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
_LoraModel._create_and_replace = _patched_car
|
||||
|
||||
try:
|
||||
model = PeftModel.from_pretrained(
|
||||
model,
|
||||
old_model_name,
|
||||
token = token,
|
||||
revision = revision,
|
||||
is_trainable = True,
|
||||
trust_remote_code = trust_remote_code,
|
||||
)
|
||||
finally:
|
||||
# Always restore original PEFT method, even if loading fails
|
||||
if _clippable_linear_cls is not None:
|
||||
_LoraModel._create_and_replace = _original_car
|
||||
|
||||
# Patch it as well!
|
||||
model = FastBaseModel.post_patch_model(
|
||||
model, use_gradient_checkpointing, trust_remote_code = trust_remote_code
|
||||
|
|
|
|||
|
|
@ -162,7 +162,7 @@ def __get_model_name(
|
|||
# Support returning original full -bnb-4bit name if specified specifically
|
||||
# since we'll map it to the dynamic version instead
|
||||
if lower_model_name.endswith("-bnb-4bit"):
|
||||
return lower_model_name
|
||||
return model_name
|
||||
|
||||
new_model_name = FLOAT_TO_INT_MAPPER[lower_model_name]
|
||||
# logger.warning_once(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue