From 5c473fab80e079bb525345b86cb71afd409262c3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Apr 2026 07:05:47 -0700 Subject: [PATCH 01/54] Bump versions --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 44464101f3..77c0034125 100644 --- a/install.ps1 +++ b/install.ps1 @@ -934,7 +934,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.8" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -942,7 +942,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.4.8" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -975,7 +975,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.4.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.4.8" unsloth-zoo } if ($baseInstallExit -eq 0) { $NoTorchReq = Find-NoTorchRuntimeFile if ($NoTorchReq) { @@ -983,7 +983,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.7" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.4.8" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1006,7 +1006,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.4.7" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.4.8" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return diff --git a/install.sh b/install.sh index 07473e441d..6c28b6eda4 100755 --- a/install.sh +++ b/install.sh @@ -1347,7 +1347,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.4.7" unsloth-zoo + "unsloth>=2026.4.8" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1355,7 +1355,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.4.7" unsloth-zoo + "unsloth>=2026.4.8" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -1519,7 +1519,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.4.7" unsloth-zoo + "unsloth>=2026.4.8" unsloth-zoo _NO_TORCH_RT="$(_find_no_torch_runtime)" if [ -n "$_NO_TORCH_RT" ]; then run_install_cmd "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" @@ -1530,7 +1530,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.4.7" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.4.8" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else @@ -1558,7 +1558,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.4.7" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.4.8" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps else From c875dc17e4f608984c5aa2def361c7aa863d9860 Mon Sep 17 00:00:00 2001 From: Konstantin Azizov Date: Fri, 24 Apr 2026 12:53:10 +0200 Subject: [PATCH 02/54] Studio: use (gguf) context length before max seq length (#5111) * fix: use (gguf) context length before max seq length For GGUF models context length is used instead of `maxSeqLength` Fixes #4893 * Studio: split rollback max_seq_length from target load/validate Restore params.maxSeqLength as the value passed to the target model's validateModel and effectiveMaxSeqLength fallback, and introduce a separate rollbackMaxSeqLength used only in the rollback loadModel call. Without the split, switching from a GGUF model with a large native context to a non-GGUF target polluted the new load via the effectiveMaxSeqLength else-branch. Also: - Detect a previous GGUF via isGguf, activeGgufVariant, or a .gguf suffix on the checkpoint, so local/LM Studio paths not present in the models catalog still take the GGUF rollback path. - Use 0 as the GGUF rollback fallback to match the sentinel used at the normal GGUF load site (ggufContextLength fallback to 0); 4096 would reintroduce the original truncation bug when both customContextLength and ggufContextLength are null. - Reuse the captured stateBeforeUnload for modelRequiresTrustRemoteCode and drop the now-unused DEFAULT_INFERENCE_PARAMS value import. * Studio: drop pending customContextLength from GGUF rollback rollbackMaxSeqLength previously preferred customContextLength over ggufContextLength, but customContextLength is a pending, not-yet-applied slider edit: chat-settings-sheet.tsx treats it as the "dirty" marker (ctxDirty = customContextLength !== null) and use-chat-model-runtime.ts clears it to null on every successful load. Rollback exists to restore the previously-loaded model at its confirmed running context, so leaking a pending slider value can load the rollback target at a context the user never validated against VRAM, potentially causing the rollback itself to fail. --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han --- .../chat/hooks/use-chat-model-runtime.ts | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index c9dcd911a2..1281592168 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -419,12 +419,19 @@ export function useChatModelRuntime() { let previousWasUnloaded = false; const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; - const paramsBeforeLoad = useChatRuntimeStore.getState().params; - const trustRemoteCode = paramsBeforeLoad.trustRemoteCode ?? false; - const maxSeqLength = paramsBeforeLoad.maxSeqLength; - const hfToken = useChatRuntimeStore.getState().hfToken || null; + const stateBeforeUnload = useChatRuntimeStore.getState(); + const trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; + const maxSeqLength = stateBeforeUnload.params.maxSeqLength; + const previousIsGguf = + previousModel?.isGguf === true + || previousVariant != null + || (previousCheckpoint?.toLowerCase().endsWith(".gguf") ?? false); + const rollbackMaxSeqLength = previousIsGguf + ? (stateBeforeUnload.ggufContextLength ?? 0) + : maxSeqLength; + const hfToken = stateBeforeUnload.hfToken || null; const previousModelRequiresTrustRemoteCode = - useChatRuntimeStore.getState().modelRequiresTrustRemoteCode; + stateBeforeUnload.modelRequiresTrustRemoteCode; try { // Lightweight pre-flight validation: avoid unloading a working model // if the new identifier is clearly invalid (e.g. bad HF id / path). @@ -542,7 +549,7 @@ export function useChatModelRuntime() { await loadModel({ model_path: previousCheckpoint, hf_token: hfToken, - max_seq_length: maxSeqLength, + max_seq_length: rollbackMaxSeqLength, load_in_4bit: true, is_lora: previousIsLora, gguf_variant: previousVariant, From 06ed94da0d1958c894e68bec863e5ec3d7254c71 Mon Sep 17 00:00:00 2001 From: luo jiyin Date: Fri, 24 Apr 2026 19:51:27 +0800 Subject: [PATCH 03/54] chore: fix typo cleanup across tests and backend strings (#5152) * chore: fix typos in studio/backend/routes/models.py * chore: fix typos in tests/saving/non_peft/test_mistral_non_peft.py * chore: fix typos in tests/saving/non_peft/test_whisper_non_peft.py * chore: fix typos in tests/saving/vision_models/test_index_file_sharded_model.py * chore: fix typos in tests/saving/vision_models/test_push_to_hub_merged.py * chore: fix typos in tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py * chore: fix typos in tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py * chore: fix typos in unsloth/import_fixes.py * Split: keep only 6 file(s) --------- Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- tests/saving/non_peft/test_mistral_non_peft.py | 2 +- tests/saving/non_peft/test_whisper_non_peft.py | 2 +- tests/saving/vision_models/test_index_file_sharded_model.py | 4 ++-- tests/saving/vision_models/test_push_to_hub_merged.py | 4 ++-- .../test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py | 4 ++-- .../test_save_merge_vision_model_ocr_benchmark.py | 4 ++-- 6 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/saving/non_peft/test_mistral_non_peft.py b/tests/saving/non_peft/test_mistral_non_peft.py index e03813367d..b2610360fe 100644 --- a/tests/saving/non_peft/test_mistral_non_peft.py +++ b/tests/saving/non_peft/test_mistral_non_peft.py @@ -27,7 +27,7 @@ model, tokenizer = FastLanguageModel.from_pretrained( print("✅ Base model loaded successfully!") -### Attemtping save merge +### Attempting save merge print(f"\n{'='*80}") diff --git a/tests/saving/non_peft/test_whisper_non_peft.py b/tests/saving/non_peft/test_whisper_non_peft.py index 303d596c85..2d5676ec9c 100644 --- a/tests/saving/non_peft/test_whisper_non_peft.py +++ b/tests/saving/non_peft/test_whisper_non_peft.py @@ -27,7 +27,7 @@ model, tokenizer = FastModel.from_pretrained( print("✅ Base model loaded successfully!") -### Attemtping save merge +### Attempting save merge print(f"\n{'='*80}") diff --git a/tests/saving/vision_models/test_index_file_sharded_model.py b/tests/saving/vision_models/test_index_file_sharded_model.py index f737169841..8d107463e0 100644 --- a/tests/saving/vision_models/test_index_file_sharded_model.py +++ b/tests/saving/vision_models/test_index_file_sharded_model.py @@ -65,7 +65,7 @@ def format_data(sample): print("\n🔄 Formatting dataset for vision training...") system_message = "You are an expert french ocr system." # Convert dataset to OAI messages -# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes +# need to use list comprehension to keep Pil.Image type, .map convert image to bytes train_dataset = [format_data(sample) for sample in train_dataset] eval_dataset = [format_data(sample) for sample in eval_dataset] print("✅ Dataset formatting completed!") @@ -99,7 +99,7 @@ try: finetune_vision_layers = True, # Turn off for just text! finetune_language_layers = True, # Should leave on! finetune_attention_modules = True, # Attention good for GRPO - finetune_mlp_modules = True, # SHould leave on always! + finetune_mlp_modules = True, # Should leave on always! r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 lora_alpha = 32, lora_dropout = 0, # Supports any, but = 0 is optimized diff --git a/tests/saving/vision_models/test_push_to_hub_merged.py b/tests/saving/vision_models/test_push_to_hub_merged.py index 74fa058988..fb2af4b4fe 100644 --- a/tests/saving/vision_models/test_push_to_hub_merged.py +++ b/tests/saving/vision_models/test_push_to_hub_merged.py @@ -66,7 +66,7 @@ def format_data(sample): print("\n🔄 Formatting dataset for vision training...") system_message = "You are an expert french ocr system." # Convert dataset to OAI messages -# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes +# need to use list comprehension to keep Pil.Image type, .map convert image to bytes train_dataset = [format_data(sample) for sample in train_dataset] eval_dataset = [format_data(sample) for sample in eval_dataset] print("✅ Dataset formatting completed!") @@ -100,7 +100,7 @@ try: finetune_vision_layers = True, # Turn off for just text! finetune_language_layers = True, # Should leave on! finetune_attention_modules = True, # Attention good for GRPO - finetune_mlp_modules = True, # SHould leave on always! + finetune_mlp_modules = True, # Should leave on always! r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 lora_alpha = 32, lora_dropout = 0, # Supports any, but = 0 is optimized diff --git a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py index ebe078c73b..2b24bc4a32 100644 --- a/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py +++ b/tests/saving/vision_models/test_save_merge_qwen2.5vl32B_model_ocr_benchmark.py @@ -61,7 +61,7 @@ def format_data(sample): system_message = "You are an expert french ocr system." # Convert dataset to OAI messages -# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes +# need to use list comprehension to keep Pil.Image type, .map convert image to bytes train_dataset = [format_data(sample) for sample in train_dataset] eval_dataset = [format_data(sample) for sample in eval_dataset] @@ -102,7 +102,7 @@ model = FastVisionModel.get_peft_model( finetune_vision_layers = True, # Turn off for just text! finetune_language_layers = True, # Should leave on! finetune_attention_modules = True, # Attention good for GRPO - finetune_mlp_modules = True, # SHould leave on always! + finetune_mlp_modules = True, # Should leave on always! r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 # target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", # "gate_proj", "up_proj", "down_proj",], diff --git a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py index b99785bcb1..16914707c2 100644 --- a/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py +++ b/tests/saving/vision_models/test_save_merge_vision_model_ocr_benchmark.py @@ -61,7 +61,7 @@ def format_data(sample): system_message = "You are an expert french ocr system." # Convert dataset to OAI messages -# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes +# need to use list comprehension to keep Pil.Image type, .map convert image to bytes train_dataset = [format_data(sample) for sample in train_dataset] eval_dataset = [format_data(sample) for sample in eval_dataset] @@ -102,7 +102,7 @@ model = FastVisionModel.get_peft_model( finetune_vision_layers = True, # Turn off for just text! finetune_language_layers = True, # Should leave on! finetune_attention_modules = True, # Attention good for GRPO - finetune_mlp_modules = True, # SHould leave on always! + finetune_mlp_modules = True, # Should leave on always! r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128 # target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", # "gate_proj", "up_proj", "down_proj",], From 0326577b821603878ae9d85b67d52eb09bf7ae46 Mon Sep 17 00:00:00 2001 From: Etherll <61019402+Etherll@users.noreply.github.com> Date: Fri, 24 Apr 2026 15:59:17 +0300 Subject: [PATCH 04/54] fix: guard resolve_model_class fallback against unresolvable transformers AutoModel entries (#5155) * fix: avoid PerceptionEncoder ImportError blocking trust_remote_code model loads * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update config class retrieval in _utils.py Refactor config class retrieval logic to use model mapping. * resolve_model_class: restore _extra_content fallback The previous fallback iterated mapping.items(), which transformers' _LazyAutoMapping defines as _model_mapping entries + _extra_content entries. The PR's per-key loop covers only _model_mapping, so subclasses of configs registered via AutoModel.register(cfg, model) silently resolve to None. Add a safe isinstance pass over _extra_content (no lazy loads, no crash risk) before giving up. * Add tests for resolve_model_class fallback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- tests/test_resolve_model_class.py | 137 ++++++++++++++++++++++++++++++ unsloth/models/_utils.py | 30 +++++-- 2 files changed, 161 insertions(+), 6 deletions(-) create mode 100644 tests/test_resolve_model_class.py diff --git a/tests/test_resolve_model_class.py b/tests/test_resolve_model_class.py new file mode 100644 index 0000000000..5ec2a21a07 --- /dev/null +++ b/tests/test_resolve_model_class.py @@ -0,0 +1,137 @@ +from unsloth.models._utils import resolve_model_class + + +class _AutoModelLike: + def __init__(self, mapping): + self._model_mapping = mapping + + +class _FakeLazyMapping: + def __init__(self, entries, extra_content = None, broken_keys = ()): + self._entries = dict(entries) + self._config_mapping = {k: f"Cfg_{k}" for k in self._entries} + self._model_mapping = {k: f"Mdl_{k}" for k in self._entries} + self._extra_content = dict(extra_content or {}) + self._broken_keys = set(broken_keys) + + def __getitem__(self, key): + if key in self._extra_content: + return self._extra_content[key] + for k, (cfg_cls, mdl_cls) in self._entries.items(): + if k in self._broken_keys: + raise ValueError(f"broken entry {k}") + if cfg_cls is key: + return mdl_cls + raise KeyError(key) + + def _load_attr_from_module(self, key, attr): + if key in self._broken_keys: + raise ValueError(f"broken entry {key}") + cfg_cls, mdl_cls = self._entries[key] + if attr == self._config_mapping[key]: + return cfg_cls + if attr == self._model_mapping[key]: + return mdl_cls + raise KeyError(attr) + + +class CfgA: + pass + + +class CfgB: + pass + + +class CfgBChild(CfgB): + pass + + +class ModelA: + pass + + +class ModelB: + pass + + +class RegBase: + pass + + +class RegChild(RegBase): + pass + + +class RegModel: + pass + + +class UnknownCfg: + pass + + +def test_fast_path_exact_match(): + m = _FakeLazyMapping({"a": (CfgA, ModelA), "b": (CfgB, ModelB)}) + am = _AutoModelLike(m) + assert resolve_model_class(am, CfgA()) is ModelA + + +def test_fallback_subclass_match_via_lazy_mapping(): + m = _FakeLazyMapping({"a": (CfgA, ModelA), "b": (CfgB, ModelB)}) + am = _AutoModelLike(m) + assert resolve_model_class(am, CfgBChild()) is ModelB + + +def test_broken_lazy_entry_does_not_crash(): + m = _FakeLazyMapping( + {"broken": (CfgA, ModelA), "b": (CfgB, ModelB)}, + broken_keys = ("broken",), + ) + am = _AutoModelLike(m) + assert resolve_model_class(am, CfgBChild()) is ModelB + + +def test_unknown_config_returns_none(): + m = _FakeLazyMapping({"a": (CfgA, ModelA)}) + am = _AutoModelLike(m) + assert resolve_model_class(am, UnknownCfg()) is None + + +def test_extra_content_subclass_fallback(): + m = _FakeLazyMapping( + {"a": (CfgA, ModelA)}, + extra_content = {RegBase: RegModel}, + ) + am = _AutoModelLike(m) + assert resolve_model_class(am, RegChild()) is RegModel + + +def test_extra_content_exact_match_fast_path(): + m = _FakeLazyMapping( + {"a": (CfgA, ModelA)}, + extra_content = {RegBase: RegModel}, + ) + am = _AutoModelLike(m) + assert resolve_model_class(am, RegBase()) is RegModel + + +def test_broken_entry_with_extra_content_subclass(): + m = _FakeLazyMapping( + {"broken": (CfgA, ModelA)}, + extra_content = {RegBase: RegModel}, + broken_keys = ("broken",), + ) + am = _AutoModelLike(m) + assert resolve_model_class(am, RegChild()) is RegModel + + +def test_plain_dict_mapping_is_not_required(): + am = _AutoModelLike({CfgA: ModelA}) + assert resolve_model_class(am, CfgA()) is ModelA + + +def test_tuple_result_unwrapped(): + m = _FakeLazyMapping({"a": (CfgA, (ModelA, "extra"))}) + am = _AutoModelLike(m) + assert resolve_model_class(am, CfgA()) is ModelA diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 8b8dec3d3e..e87c2a1a83 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -415,13 +415,31 @@ def resolve_model_class(auto_model, config): try: result = mapping[config.__class__] except Exception: - for config_class, model_class in mapping.items(): - if isinstance(config, config_class): - result = model_class - break - else: + result = None + for key in list(getattr(mapping, "_model_mapping", {})): + try: + config_class = mapping._load_attr_from_module( + key, mapping._config_mapping[key] + ) + if isinstance(config, config_class): + result = mapping._load_attr_from_module( + key, mapping._model_mapping[key] + ) + break + except Exception: + continue + if result is None: + for extra_cls, extra_model in getattr( + mapping, "_extra_content", {} + ).items(): + try: + if isinstance(config, extra_cls): + result = extra_model + break + except Exception: + continue + if result is None: return None - return result[0] if isinstance(result, (list, tuple)) else result From c2dc2eb1b128747e8f522d271de16586a07a0837 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Apr 2026 09:05:31 -0700 Subject: [PATCH 05/54] Studio: kill in-flight llama-server before spawning a new one (#5171) * Studio: kill in-flight llama-server before spawning a new one Two rapid Apply clicks in the chat settings panel can race two load_model calls. Both pass the Phase 1 _kill_process (because neither has stored its Popen handle yet), both download / read metadata, and both reach Phase 3 and spawn a server. Only the last reference is tracked in self._process. The first server becomes an orphan that holds the model in RAM until the kernel OOM kicks in. Addresses #5161. Two complementary changes in studio/backend/core/inference/llama_cpp.py: 1. At load_model entry, set the existing _cancel_event so any in-flight load aborts at its next checkpoint, then bind a fresh Event for the new load. Subsequent _kill_process and download phases pick up the new event. 2. Inside the Phase 3 lock, immediately before subprocess.Popen, run a defensive _kill_process that removes any orphan handle a racing load might have stored after the first kill ran. Speculative decoding cleanup (also touched while in this code path): The chat UI used to send "ngram-mod" as the wire value when the speculative dropdown was On, and the backend mapped that to the 4-flag combo --spec-type ngram-mod --spec-ngram-size-n 24 --draft-min 48 --draft-max 64. Switch the wire value to "default" and have the backend pass the single llama-server flag --spec-default. That flag expands to the exact same params (see common/arg.cpp:3905-3914 in llama.cpp). Default-on for non-vision models is preserved. Verified: - Unit test: planted Popen orphan handle is terminated before self._process is overwritten. - All ten spec-cmd mappings produce the expected llama-server args ("default" -> --spec-default, "off" / null -> no flag, vision -> always disabled, manual "ngram-mod" / "ngram-simple" still work). * Address review feedback on PR #5171 The previous attempt at cancelling in-flight loads via ``self._cancel_event.set()`` followed by ``self._cancel_event = threading.Event()`` was broken in two ways (flagged independently by Gemini and Codex): 1. Sub-methods like _download_gguf consult ``self._cancel_event`` on every check. After the rebind, the in-flight thread reads the FRESH unset Event, not the one we just set, so cancellation never propagates. 2. Worse, if unload_model() lands between ``set()`` and the rebind, unload's signal hits the OLD event and is then immediately discarded when load_model swaps in a fresh Event. The user's stop-request silently no-ops. Revert to the original ``self._cancel_event.clear()``. The Phase 3 defensive ``_kill_process()`` introduced in this branch still closes the orphan-process race that #5161 reports: even if two concurrent loads both pass Phase 1 with self._process == None, the loser's Phase 3 kill terminates the winner's Popen handle before overwriting it, so we end up with exactly one llama-server process. * Studio: coerce legacy speculative-type values for the simplified dropdown The Speculative Decoding control was simplified to On (default) / Off, but the backend still accepts and reports the older manual modes (ngram-mod, ngram-simple). When a load response or status refresh comes back with one of those values -- whether from an external API caller, a model loaded before this PR landed, or a not-yet-upgraded backend -- the controlled Select renders with an empty trigger because the value is not in the SelectItem list. Add a tiny normaliser at both entry points (status refresh + post-load) so legacy manual modes coerce to "default". The user sees "On" instead of a blank dropdown, and reapplying lets llama.cpp pick its own preferred strategy via --spec-default. Reviewer-flagged finding on PR #5171. --- studio/backend/core/inference/llama_cpp.py | 26 +++++++++++++++---- .../src/features/chat/chat-settings-sheet.tsx | 2 +- .../chat/hooks/use-chat-model-runtime.ts | 17 ++++++++++-- .../chat/stores/chat-runtime-store.ts | 4 +-- 4 files changed, 39 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c320f03b2c..800d7fdd8c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1574,11 +1574,21 @@ class LlamaCppBackend: # 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 + # ``"default"`` -> let llama-server pick a sensible spec + # config via ``--spec-default``. Explicit type names are + # passed through with the manual draft tuning we've shipped + # historically so power users keep their overrides. _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": + normalized_spec = ( + speculative_type.lower().strip() if speculative_type else None + ) + if normalized_spec and normalized_spec != "off" and not is_vision: + if normalized_spec == "default": + cmd.append("--spec-default") + self._speculative_type = "default" + elif normalized_spec in _valid_spec_types: + cmd.extend(["--spec-type", normalized_spec]) + if normalized_spec == "ngram-mod": cmd.extend( [ "--spec-ngram-size-n", @@ -1589,7 +1599,7 @@ class LlamaCppBackend: "64", ] ) - self._speculative_type = speculative_type + self._speculative_type = normalized_spec else: self._speculative_type = None else: @@ -1750,6 +1760,12 @@ class LlamaCppBackend: if gpu_indices is not None: env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in gpu_indices) + # Defensive kill: if a concurrent load slipped past Phase 1 + # (because its `self._process` was None at the time) and + # already stored a Popen handle here, drop that orphan + # before we overwrite the reference. See issue #5161. + self._kill_process() + self._stdout_lines = [] self._process = subprocess.Popen( cmd, diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index e4574dab74..fc5f097969 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1037,7 +1037,7 @@ export function ChatSettingsPanel({ - On + On Off diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 1281592168..e90cab73fa 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -24,6 +24,19 @@ import type { InferenceParams, } from "../types/runtime"; +// The simplified Speculative Decoding control surfaces "default" (which +// maps to llama.cpp's --spec-default) and "off". A backend status / load +// response can still report the older manual modes (ngram-mod, +// ngram-simple) when a model is loaded via the API or carried over from an +// older Studio version. The Select would render an empty trigger for those +// values, so coerce them to "default" -- llama.cpp's own --spec-default +// picks an equivalent strategy and keeps the dropdown coherent. +function normalizeSpeculativeType(v: string | null | undefined): string | null { + if (v == null) return null; + if (v === "default" || v === "off") return v; + return "default"; +} + type SelectedModelInput = { id: string; isLora?: boolean; @@ -279,7 +292,7 @@ export function useChatModelRuntime() { const ggufNativeContextLength = statusRes.is_gguf ? (statusRes.native_context_length ?? null) : null; - const currentSpecType = statusRes.speculative_type ?? null; + const currentSpecType = normalizeSpeculativeType(statusRes.speculative_type); useChatRuntimeStore.setState({ supportsReasoning, reasoningAlwaysOn, @@ -492,7 +505,7 @@ export function useChatModelRuntime() { } } const loadedKv = loadResponse.cache_type_kv ?? null; - const loadedSpec = loadResponse.speculative_type ?? null; + const loadedSpec = normalizeSpeculativeType(loadResponse.speculative_type); const nativeCtx = loadResponse.is_gguf ? (loadResponse.context_length ?? 131072) : null; diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 2a1249aea0..ce69d8f6dc 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -270,7 +270,7 @@ export const useChatRuntimeStore = create((set) => ({ toolCallTimeout: loadInt(TOOL_CALL_TIMEOUT_KEY, 5), kvCacheDtype: null, loadedKvCacheDtype: null, - speculativeType: "ngram-mod", + speculativeType: "default", loadedSpeculativeType: null, customContextLength: null, defaultChatTemplate: null, @@ -365,7 +365,7 @@ export const useChatRuntimeStore = create((set) => ({ toolStatus: null, kvCacheDtype: null, loadedKvCacheDtype: null, - speculativeType: "ngram-mod", + speculativeType: "default", loadedSpeculativeType: null, customContextLength: null, defaultChatTemplate: null, From ae9de7f2df80b592af02d0d610627a4c21663b35 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Apr 2026 09:06:01 -0700 Subject: [PATCH 06/54] Studio: stop currency escape from breaking inline LaTeX (#5170) * Studio: stop currency escape from breaking inline LaTeX The currency-escape preprocessor in studio/frontend/src/lib/latex.ts matched the opening dollar of any $...$ span and inserted a backslash. The result was that text like "$30^\circ$" or "**$90 - x$**" rendered as raw characters with stray dollar signs. Fixes #5164. Add two helpers in front of the escape: - hasInlineMathCloser looks for an unescaped, non-doubled closing dollar within the same line. Bold-wrapped spans (**$X$**) are always treated as math since LLMs use that form for bold math. - looksLikeMathBody filters multi-token bodies that look like prose between two currency tokens ($5 to $10, $5, $10). Verified against 111 inputs: the issue body, common LaTeX patterns (Greek vars, fractions, integrals, vectors, exponents), prose currency in lists and sentences, code blocks, and headings. All pass. * Address review feedback on PR #5170 - Drop ^ and _ from MATH_OP_RE since LATEX_CHAR_RE already short- circuits on those before MATH_OP_RE is consulted (Gemini comment). - Treat compact currency ranges like $5-$10 and $5/$10 as currency rather than math. The body between the first two dollars in those forms is "5-" or "5/", a single non-whitespace token that previously hit the math shortcut. Extend TRAIL_PUNCT_RE to strip - and / so the trimmed body comes back as pure currency. (Codex comment.) - Honour __underscore-bold__ around math the same way as **-bold**. Markdown allows both delimiters and LLMs do reach for the underscore form. (Gemini comment.) Verified against the existing 18 cases plus 5 new ones for the range, slash, and underscore-bold scenarios. All pass. * Studio: fix numeric inline math + currency-as-closer in LaTeX preprocess Two reviewer-flagged real-world misses in the inline-math heuristic. 1) Numeric-only operator forms like $2 + 2$, $100 < 200$, $1,000 - 500$ were getting their leading $ escaped, so the renderer never saw them as math. The body has a math op but no lone-letter variable, so the old looksLikeMathBody required the lone-letter clause and rejected purely numeric expressions. Add SIMPLE_MATH_RE to recognise number- or-letter operands joined by math operators. 2) Prose like "Starts at $5 + a $10 add-on" was being treated as one math span "5 + a " with the second currency token mistaken for the closer. The body satisfied the math-op + lone-letter check, so the span got accepted and the renderer ate "10 add-on". In hasInlineMathCloser, reject any candidate $ whose next character is a digit -- that's almost always another currency token starting, not the closer of a real math span (math doesn't follow $ with a bare digit). Verified via temp/pr_simulation/sim_5170_latex.mjs: 25/25 cases pass, including the 6 reviewer numeric-math cases, 3 currency-as-closer cases, and 16 regression checks against the originally shipped behavior. --- studio/frontend/src/lib/latex.ts | 117 +++++++++++++++++++++++++++++++ 1 file changed, 117 insertions(+) diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 954d0f7bc6..0092d51be0 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -74,12 +74,126 @@ function isInCodeBlock( return false; } +/** + * A token (no whitespace) that looks purely like currency, e.g. `5`, + * `1,000`, `5.99`, `100K`, `3.5M`. + */ +const CURRENCY_BODY_RE = /^\d+(?:,\d{3})*(?:\.\d+)?[KMBkmb]?$/; + +/** Body characters that almost always indicate real LaTeX. */ +const LATEX_CHAR_RE = /[\\^_{}]/; + +/** + * Operators that strongly suggest a math expression. We deliberately + * leave out `^` and `_` because `LATEX_CHAR_RE` already short-circuits + * on those before we ever consult this regex. + */ +const MATH_OP_RE = /[=+\-<>/*]/; + +/** + * Trailing chars stripped before the currency check. Includes prose + * punctuation plus the connectors `-` and `/` that appear in compact + * currency ranges like `$5-$10` or `$5/$10`; without them the body + * `5-` or `5/` would slip through the single-token math shortcut. + */ +const TRAIL_PUNCT_RE = /[.,;:!?\-/]+$/; + +/** + * A standalone single letter (variable name) inside the body. We require + * that the letter is not part of a longer word so that prose like + * "5 to attend" doesn't get misread as a math expression with the + * variable `t`. + */ +const LONE_LETTER_RE = /(?/*]\s*(?:\d+(?:,\d{3})*(?:\.\d+)?|[a-zA-Z]))+$/; + +/** + * Return true if the substring between two `$` delimiters looks like a + * LaTeX expression rather than a span of prose between two currency + * tokens. + * + * Rule of thumb: + * - `$30^\circ$` -> math (LaTeX chars) + * - `$x$` -> math (single non-currency token) + * - `$90 - x$` -> math (math op + lone variable) + * - `$5 to $10` -> NOT math (multi-token prose, no math op) + * - `$5, $10` -> NOT math (currency-like token + trailing punct) + * - `$1,000$` -> NOT math (single currency-like token) + */ +function looksLikeMathBody(body: string): boolean { + if (LATEX_CHAR_RE.test(body)) return true; + const trimmed = body.trim().replace(TRAIL_PUNCT_RE, ""); + if (!trimmed) return false; + if (CURRENCY_BODY_RE.test(trimmed)) return false; + // Numeric-only operator forms: `2 + 2`, `100 < 200`, `1,000 - 500`. + // Recognised without requiring a lone-variable letter. + if (SIMPLE_MATH_RE.test(trimmed)) return true; + if (!/\s/.test(trimmed)) return true; + if (!MATH_OP_RE.test(trimmed)) return false; + return LONE_LETTER_RE.test(trimmed); +} + +/** + * Return true if the `$` at `offset` opens a balanced inline math span + * (`$...$`) on the same line. The closer must be unescaped, must not be + * part of `$$`, and must lie inside a 200-character window. The body + * must look like LaTeX so we don't pair two currency tokens that share + * a line (e.g. "$5 to $10"). Bold-wrapped spans (`**$X$**` and the + * underscore equivalent `__$X$__`) are always treated as math because + * LLMs reach for that pattern when they want "bold math" and the + * heuristic would otherwise reject prose-shaped bodies like "90 - x". + */ +function hasInlineMathCloser(content: string, offset: number): boolean { + const MAX_SPAN = 200; + const limit = Math.min(content.length, offset + 1 + MAX_SPAN); + for (let i = offset + 1; i < limit; i++) { + const c = content[i]; + if (c === "\n") return false; + if (c !== "$") continue; + if (content[i - 1] === "\\") continue; + if (content[i + 1] === "$") { + i++; + continue; + } + // A `$` immediately followed by a digit is far more likely the start + // of another currency token than the closer for the current span. Keep + // scanning so prose like `Starts at $5 + a $10 add-on` doesn't pair + // the two currency markers as a math span. + if (/\d/.test(content[i + 1] ?? "")) { + continue; + } + if (offset >= 2) { + const op = content[offset - 1]; + if ( + (op === "*" || op === "_") && + content[offset - 2] === op && + content[i + 1] === op && + content[i + 2] === op + ) { + return true; + } + } + return looksLikeMathBody(content.slice(offset + 1, i)); + } + return false; +} + /** * Preprocess a markdown string to escape currency dollar signs so they are not * parsed as LaTeX math delimiters. * * - `$5` alone becomes `\$5` (currency, not math) * - `$\alpha$` is untouched (real LaTeX) + * - `$30^\circ$` is untouched (LaTeX whose body starts with a digit) + * - `**$30^\circ$**` is untouched (LaTeX wrapped in bold) * - `$$E = mc^2$$` is untouched (display math) * - Currency inside code blocks/spans is untouched */ @@ -92,6 +206,9 @@ export function preprocessLaTeX(content: string): string { if (isInCodeBlock(offset, codeRegions)) { return match; } + if (hasInlineMathCloser(content, offset)) { + return match; + } return "\\" + match; }); } From 8264e80dd952df0fc05c85c6f9671ab7d9991638 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Apr 2026 10:00:42 -0700 Subject: [PATCH 07/54] Studio: probe AMD GPUs in llama-server VRAM detection (#5172) * Studio: probe AMD GPUs in llama-server VRAM detection _get_gpu_free_memory in studio/backend/core/inference/llama_cpp.py only queried nvidia-smi. On AMD ROCm hosts that returns nothing, so the GPU list is empty, the auto-fit logic falls into the no-gpus branch, and llama-server gets --fit on with no -ngl to anchor it. The model loads on CPU even though the GPU is detected elsewhere in Studio. Addresses #5106. Add a torch-based fallback that runs after nvidia-smi fails or returns empty: import torch if torch.cuda.is_available() and hasattr(torch.cuda, "mem_get_info"): for ordinal in range(torch.cuda.device_count()): free, _total = torch.cuda.mem_get_info(ordinal) gpus.append((ordinal, free // (1024 * 1024))) Works on AMD because the ROCm torch wheels Studio installs reuse the entire torch.cuda.* namespace via HIP. Also rescues NVIDIA hosts where nvidia-smi is missing from PATH (a secondary cause of the bug on Windows). Matches the convention studio/backend/utils/hardware/hardware.py:412 already uses for the same fallback purpose. Verified locally: nvidia-smi path returns the expected GPU and free MiB; torch fallback returns valid VRAM when nvidia-smi is forced to fail. Note: PR #4874 is a draft taking a different approach (parsing vulkaninfo); the two are complementary. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review feedback on PR #5172 torch.cuda.device_count() enumerates GPUs RELATIVE to the current CUDA_VISIBLE_DEVICES (or HIP_VISIBLE_DEVICES on ROCm). Returning those visible ordinals directly lets _select_gpus rewrite CUDA_VISIBLE_DEVICES with the wrong physical IDs: a process started with CUDA_VISIBLE_DEVICES=2,3 would get its child llama-server relaunched with CUDA_VISIBLE_DEVICES=0,1, targeting the wrong GPUs and violating any scheduler pinning. Translate visible ordinals back through the active CVD/HIP/ROCR mask before returning. Falls through to bare ordinal when no mask is set. Also drop the redundant int() cast on // -- bytes // 2**20 already returns int. Verified: with CUDA_VISIBLE_DEVICES=6 and nvidia-smi forced to fail, the torch fallback now returns (6, free_mib) instead of (0, free_mib). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix ROCm visibility precedence + narrow ROCm child env Two reviewer-flagged correctness bugs in the AMD GPU probe path. 1) ROCm visibility precedence was reversed. torch.cuda enumerates GPUs relative to HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES on ROCm builds, but the probe's env-var lookup checked CUDA_VISIBLE_DEVICES first. With CUDA_VISIBLE_DEVICES=0,1 and HIP_VISIBLE_DEVICES=6,7 the probe returned [(0, ...), (1, ...)] when torch's view was actually [(6, ...), (7, ...)]. The wrong physical IDs flowed downstream into CUDA_VISIBLE_DEVICES for the llama-server subprocess, pinning it to GPUs 0,1 instead of 6,7. Fix: branch on torch.version.hip. On ROCm, prefer HIP > ROCR > CUDA (matches torch's own ordering). On NVIDIA, use CUDA only -- ignoring any HIP/ROCR vars the parent happens to have set. 2) Child env narrowing only set CUDA_VISIBLE_DEVICES. On ROCm, llama-server honors HIP/ROCR; if the parent shell exported HIP_VISIBLE_DEVICES=4,5 and the selector picked just GPU 4, the child still saw both because we never narrowed HIP/ROCR. Now we set all three on ROCm so the AMD subprocess actually sees the planned subset. Both branches verified via temp/pr_simulation/sim_5172_rocm_precedence.py (7/7 cases pass), including the reviewer's verbatim R5 case (CVD=0,1 + HIP/ROCR=6,7). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: sort GPU probe result + honor explicitly empty ROCm masks Two reviewer-flagged correctness nits on top of eff55fb8. 1) Gemini medium: the torch fallback returned an unsorted list when the visibility mask was non-sequential (e.g. CUDA_VISIBLE_DEVICES=5,2,9), diverging from the docstring guarantee and the nvidia-smi path. Now sorted by physical id. 2) Codex P2: an explicitly empty HIP_VISIBLE_DEVICES="" should mean "no GPUs" per the codebase convention in utils/hardware/hardware.py::_get_parent_visible_gpu_spec. The previous `or` chain treated empty string as falsy and silently fell through to ROCR / CUDA, producing wrong physical IDs. Switch to `is not None` checks to match. Verified via sim_5172_rocm_precedence.py (9/9 cases pass) including the two new R8 (sort) and R9 (empty-HIP honored) cases. * Studio: align nvidia-smi probe with torch fallback (sort + robust CVD) Two follow-up Gemini-medium nits on PR #5172. 1) Fragile CVD parsing on the nvidia-smi path: `cvd.split(",")` would raise ValueError on a trailing comma like "0,1," because the empty trailing token is not skipped. The torch fallback already filters empty tokens via `if x.strip()`; mirror that here. 2) Missing sort guarantee on the nvidia-smi path: the docstring promises sort-by-id, the torch fallback now sorts, but the nvidia-smi path relied on driver enumeration order. Add an explicit sort. Both changes match what shipped in 6b1cccd6 for the torch fallback, so the two probe paths now have identical CVD parsing + ordering semantics. * Studio: drop cvd.strip() truthiness so empty CVD filters all GPUs Reviewer-flagged correctness bug. The previous `if cvd is not None and cvd.strip():` guard treated `CUDA_VISIBLE_DEVICES=""` as if the variable were unset, leaving `allowed=None` (and `physical_ids=None` on the torch path). On the nvidia-smi path that mattered: nvidia-smi ignores CVD entirely, so the probe's `allowed` filter is the only thing that respects the parent's "no GPUs" intent. Pre-fix the probe returned every physical GPU when the parent had explicitly hidden them. Drop the `.strip()` truthiness check on both paths. The downstream `if x.strip()` token filter still keeps trailing-comma masks like "0,1," safe, and an empty mask now produces an empty allowed/physical set as expected (matching utils/hardware/hardware.py convention). Verified via sim_5172_rocm_precedence.py R10 + R11 (now 11/11 cases pass): nvidia-smi path with `CUDA_VISIBLE_DEVICES=""` now returns [] instead of leaking the hidden GPUs. * Studio: log ROCm env-var failures instead of silently swallowing Reviewer-flagged defensive logging gap. The bare `except Exception: pass` around the HIP/ROCR env-var assignment would mask anything from a missing torch import to an unexpected version object shape. Log at debug so a failed AMD child-env narrowing is at least traceable. Behavior is unchanged: torch missing or version probe failing still leaves the child with only CUDA_VISIBLE_DEVICES set. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 147 +++++++++++++++++---- 1 file changed, 118 insertions(+), 29 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 800d7fdd8c..35ef90ee37 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -549,14 +549,24 @@ class LlamaCppBackend: @staticmethod def _get_gpu_free_memory() -> list[tuple[int, int]]: - """Query free memory per GPU via nvidia-smi. + """Query free memory per GPU. - Returns list of (gpu_index, free_mib) sorted by index. - Respects CUDA_VISIBLE_DEVICES if set. - Returns empty list if nvidia-smi is not available. + Order: + 1. ``nvidia-smi`` (NVIDIA CUDA hosts) -- respects + ``CUDA_VISIBLE_DEVICES``. + 2. ``torch.cuda.mem_get_info`` -- universal fallback that + works on AMD ROCm too because the HIP runtime + reuses the entire ``torch.cuda.*`` namespace. Covers the + AMD case for issue #5106 (nvidia-smi-only probe silently + returned [] on AMD hosts) and also rescues NVIDIA hosts + where ``nvidia-smi`` is missing from PATH. + + Returns list of (gpu_index, free_mib) sorted by index. Empty + list if no supported GPU is reachable. """ import os + # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( [ @@ -569,30 +579,95 @@ class LlamaCppBackend: timeout = 10, **_windows_hidden_subprocess_kwargs(), ) - if result.returncode != 0: - return [] - - # Parse which GPUs are allowed by existing CUDA_VISIBLE_DEVICES - allowed = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None and cvd.strip(): - try: - allowed = set(int(x.strip()) for x in cvd.split(",")) - except ValueError: - pass # Non-numeric (e.g., "GPU-uuid"), ignore filter - - gpus = [] - for line in result.stdout.strip().splitlines(): - parts = line.split(",") - if len(parts) == 2: - idx = int(parts[0].strip()) - free_mib = int(parts[1].strip()) - if allowed is not None and idx not in allowed: - continue - gpus.append((idx, free_mib)) - return gpus + if result.returncode == 0: + allowed: Optional[set[int]] = None + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is not None: + try: + # `if x.strip()` filters trailing-comma masks like + # "0,1," which would otherwise raise ValueError on + # an empty token. An explicitly empty mask (CVD="") + # yields an empty `allowed` set so all GPUs are + # filtered out, matching the codebase convention. + allowed = set( + int(x.strip()) for x in cvd.split(",") if x.strip() + ) + except ValueError: + pass + gpus: list[tuple[int, int]] = [] + for line in result.stdout.strip().splitlines(): + parts = line.split(",") + if len(parts) == 2: + idx = int(parts[0].strip()) + free_mib = int(parts[1].strip()) + if allowed is not None and idx not in allowed: + continue + gpus.append((idx, free_mib)) + # Match the docstring's sort-by-id guarantee. nvidia-smi + # almost always returns sorted output, but driver order + # is not formally guaranteed. + gpus.sort(key = lambda g: g[0]) + if gpus: + return gpus except Exception as e: - logger.debug(f"Failed to query GPU free memory via nvidia-smi: {e}") + logger.debug(f"nvidia-smi probe failed: {e}") + + # ── Torch fallback (covers AMD ROCm and missing nvidia-smi) ── + try: + import torch + + if not hasattr(torch, "cuda") or not torch.cuda.is_available(): + return [] + if not hasattr(torch.cuda, "mem_get_info"): + return [] + # torch.cuda enumerates GPUs RELATIVE to the visibility mask. + # On NVIDIA builds the mask is CUDA_VISIBLE_DEVICES; on AMD + # ROCm builds it is HIP_VISIBLE_DEVICES (or ROCR_VISIBLE_DEVICES + # if HIP is unset). Downstream we feed these IDs back into the + # llama-server subprocess as CVD, so we must translate visible + # ordinals back to physical indices first; otherwise launching + # with ``CUDA_VISIBLE_DEVICES=2,3`` would get rewritten to + # ``CUDA_VISIBLE_DEVICES=0,1`` and target the wrong GPUs. + physical_ids: Optional[list[int]] = None + # Match the codebase convention in + # ``utils/hardware/hardware.py::_get_parent_visible_gpu_spec``: + # treat an explicitly empty mask (``HIP_VISIBLE_DEVICES=""``) + # as "set to no GPUs" rather than falling through to the next + # var. ``or`` would coerce empty string to falsy and silently + # promote the wrong source. + if getattr(torch.version, "hip", None) is not None: + hip_v = os.environ.get("HIP_VISIBLE_DEVICES") + rocr_v = os.environ.get("ROCR_VISIBLE_DEVICES") + cvd = ( + hip_v + if hip_v is not None + else rocr_v + if rocr_v is not None + else os.environ.get("CUDA_VISIBLE_DEVICES") + ) + else: + cvd = os.environ.get("CUDA_VISIBLE_DEVICES") + if cvd is not None: + try: + # Empty mask (CVD="") yields an empty list so the + # below loop produces no GPUs, consistent with the + # nvidia-smi path and utils/hardware/hardware.py. + physical_ids = [int(x.strip()) for x in cvd.split(",") if x.strip()] + except ValueError: + physical_ids = None + gpus = [] + for ordinal in range(torch.cuda.device_count()): + free_bytes, _total_bytes = torch.cuda.mem_get_info(ordinal) + idx = ( + physical_ids[ordinal] + if physical_ids is not None and ordinal < len(physical_ids) + else ordinal + ) + gpus.append((idx, free_bytes // (1024 * 1024))) + # Match the nvidia-smi path's docstring guarantee of sorted-by-id. + return sorted(gpus, key = lambda g: g[0]) + except Exception as e: + logger.debug(f"torch GPU probe failed: {e}") return [] @staticmethod @@ -1756,9 +1831,23 @@ class LlamaCppBackend: f"{new_ld}:{existing_ld}" if existing_ld else new_ld ) - # Pin to selected GPU(s) via CUDA_VISIBLE_DEVICES + # Pin to selected GPU(s). On ROCm, llama-server (and any torch + # in the subprocess) honors HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES; + # narrowing only CUDA_VISIBLE_DEVICES leaves an AMD child seeing + # the full HIP/ROCR set the parent inherited. if gpu_indices is not None: - env["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in gpu_indices) + pinned = ",".join(str(i) for i in gpu_indices) + env["CUDA_VISIBLE_DEVICES"] = pinned + try: + import torch as _torch + + if getattr(_torch.version, "hip", None) is not None: + env["HIP_VISIBLE_DEVICES"] = pinned + env["ROCR_VISIBLE_DEVICES"] = pinned + except Exception as e: + logger.debug( + "Failed to set ROCm visibility env vars for child: %s", e + ) # Defensive kill: if a concurrent load slipped past Phase 1 # (because its `self._process` was None at the time) and From eb8b0dee2e7b06727360add18e2759ab27849f93 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Apr 2026 10:09:25 -0700 Subject: [PATCH 08/54] Studio: make stop button actually stop generation (#5069) * Studio: make stop button actually stop generation The UI stop button routes through assistant-ui's cancelRun, which aborts the frontend fetch. Four issues combined to let llama-server keep decoding long after the user clicked stop: 1. request.is_disconnected() does not fire reliably behind proxies (e.g. Colab) that don't propagate fetch aborts. 2. llama-server defaults n_predict to n_ctx when max_tokens is not sent, so a cancelled request keeps producing tokens up to 262144. 3. The httpx.Client pool keeps TCP keep-alive, so even a cleanly closed stream reuses the same connection and llama-server's liveness poll never sees a disconnect. 4. No explicit backend route to cancel - every cancel path relied on is_disconnected. Changes: - Add POST /api/inference/cancel keyed by session_id/completion_id, with a registry populated for the lifetime of each streaming response. - Have the frontend (chat-adapter.ts) POST /inference/cancel on AbortController abort, alongside the existing fetch teardown. - Send max_tokens=4096 + t_max_predict_ms=120000 as defaults on every outbound chat completion to llama-server; honoured by user overrides. - Disable httpx keep-alive on the streaming client so connection close reaches llama-server and its 1s liveness check fires. No behaviour changes for non-streaming paths or for existing callers that already pass max_tokens/session_id. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: harden stop-button cancel path and scope cancel route - Require at least one identifier for /api/inference/cancel so a missing thread id cannot silently cancel every in-flight generation. - Scope /cancel to a dedicated studio_router so it is not exposed under the /v1 OpenAI-compat prefix as a surprise endpoint. - Store a set of cancel events per key in _CANCEL_REGISTRY so concurrent requests on the same session_id do not overwrite each other, and deduplicate in _cancel_by_keys so the cancelled count reflects unique requests. - Always send session_id with chat completions (not only when tools are enabled) so non-tool GGUF streams register under it and are reachable from /cancel. - Register the non-GGUF stream_chunks path in the cancel registry too, so transformers-based stop-button works behind proxies that swallow fetch aborts. - Only apply the 2-minute t_max_predict_ms wall-clock cap when the caller did not pass max_tokens, so legitimate long generations on slow CPU/macOS/Windows supported installs are not silently truncated. - Remove the abort listener on normal stream completion so reused AbortSignals cannot fire a spurious cancel POST after the fact. * studio: close cancel-race and stale-cancel gaps in stop path - Register the cancel tracker before returning StreamingResponse so a stop POST that arrives during prefill / warmup / proxy buffering finds an entry in _CANCEL_REGISTRY. Cleanup now runs via a Starlette BackgroundTask instead of a finally inside the async generator body. - Add a per-run cancel_id on the frontend (crypto.randomUUID) and in ChatCompletionRequest so /api/inference/cancel matches one specific generation. Removes the stale-cancel bug where pressing stop then starting a new run in the same thread would cancel the retry. - Apply t_max_predict_ms unconditionally in all three llama-server payload builders (previously gated on max_tokens=None, which made it dead code for UI callers that always send params.maxTokens). Raise the default to 10 minutes so slow CPU / macOS / Windows installs are not cut off mid-generation. - Make _cancel_by_keys refuse empty input (return 0) so a future internal caller can not accidentally mass-cancel every in-flight request. - Accept cancel_id (primary), session_id, and completion_id on the /api/inference/cancel route. Unify the three streaming sites on the same _cancel_keys / _tracker variable names. - Annotate _CANCEL_REGISTRY as dict[str, set[threading.Event]]. * Add review tests for PR #5069 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: harden stop-button cancel semantics and wall-clock cap - Make /inference/cancel match cancel_id EXCLUSIVELY when supplied. Previously the handler iterated ('cancel_id','session_id','completion_id') and unioned matches, so a stale cancel POST carrying {cancel_id:old, session_id:thr} would still cancel a later run on the same thread via the shared session_id. cancel_id is now a per-run exclusive key; session_id / completion_id are only used as fallbacks when cancel_id is absent. - Close the early-cancel race. If /inference/cancel lands before the streaming handler reaches _TrackedCancel.__enter__() (stop clicked during prefill / warmup / proxy buffering), the cancel was silently dropped. Stash unmatched cancel_ids in _PENDING_CANCELS with a 30 s TTL; _TrackedCancel.__enter__() now replays any matching pending cancel by set()-ing the event immediately after registration. - Make t_max_predict_ms = _DEFAULT_T_MAX_PREDICT_MS conditional on max_tokens is None at all three llama-server payload sites. The cap is a safety net for callers who leave max_tokens unset (otherwise llama-server defaults n_predict to n_ctx, up to 262144). Callers who set an explicit max_tokens are already self-limiting and must not be silently truncated at 10 minutes on slow CPU / macOS / Windows legitimate long generations. - Guard each StreamingResponse return with try/except BaseException so _tracker.__exit__ runs even if StreamingResponse construction or any preceding statement raises between _tracker.__enter__() and the BackgroundTask attachment. Prevents a registry leak on that narrow window. * studio: close TOCTOU race and restore wall-clock backstop on UI path - Close TOCTOU race in the pending-cancel mechanism. The previous fix split cancel_inference's (cancel_by_keys + remember_pending_cancel) and _TrackedCancel.__enter__'s (register + consume_pending) into four separate lock acquisitions. Under contention a cancel POST could acquire-then-release the lock, find the registry empty, and stash ONLY AFTER __enter__ had already registered and consumed an empty pending map -- silently dropping the cancel. Both call sites now do their work inside a single _CANCEL_LOCK critical section, via the new atomic helper _cancel_by_cancel_id_or_stash() and an inlined consume-pending step in __enter__. Reproduced the race under forced interleaving pre-fix; 0/2000 drops post-fix under parallel stress. - Apply t_max_predict_ms UNCONDITIONALLY at all three llama-server payload sites. The previous iteration gated the cap on `max_tokens is None`, which turned out to be dead code on the primary Studio UI path: chat-adapter.ts sets maxTokens=loadResp.context_length after every model load, so every chat request carries an explicit max_tokens and the wall-clock safety net never fired. The cap's original purpose is to bound stuck decodes regardless of the token budget; it must always apply. - Raise _DEFAULT_T_MAX_PREDICT_MS from 10 minutes to 1 hour. 10 minutes was too aggressive for legitimate slow-CPU chat responses (a 4096-token reply at 2 tok/s takes ~34 min); 1 hour accommodates that and still catches genuine zombie decodes. - Prune _PENDING_CANCELS inside _cancel_by_keys as well, so stashed entries expire proportionally to overall cancel traffic rather than only to cancel_id-specific POSTs. * studio: trim verbose comments and docstrings in cancel path * studio/llama_cpp: drop upstream PR hashes from benchmark comment * Add review tests for Studio stop button * Consolidate review tests for Studio stop button * Align cancel-route test with exclusive cancel_id semantics * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: move cancel cleanup to generator finally; drop dead helper - Move _tracker.__exit__ from Starlette BackgroundTask into each streaming generator's finally block. Starlette skips the background callback when stream_response raises (OSError / ClientDisconnect), which leaked _CANCEL_REGISTRY entries on abrupt disconnect. - Check cancel_event.is_set() at the top of each GGUF while loop so a pending-replay cancel falls through to final_chunk + [DONE] instead of propagating GeneratorExit out of _stream_with_retry. - Remove unused _remember_pending_cancel; _cancel_by_cancel_id_or_stash superseded it. * Add review tests for Studio stop-button * studio: wire audio-input stream into cancel registry - Register cancel_event with _TrackedCancel on the audio-input streaming path so POST /api/inference/cancel can stop whisper / audio-input GGUF runs. Previously the registry stayed empty on this branch, so the stop button returned {"cancelled":0} and the decode ran to completion. - Apply the same finally-based cleanup and pre-iteration cancel-event check used on the other three streaming paths. - Update the _CANCEL_REGISTRY block comment to list cancel_id as the primary key (was stale "session_id preferred"). * Consolidate review tests for Studio stop-button cancel flow - Merge the 6 behavioral tests from test_stream_cleanup_on_disconnect.py (finally cleanup on normal/exception/aclose, pre-set cancel_event pattern, and its regressions) into test_stream_cancel_registration_timing.py, which is the PR's existing file covering the same area. - Extend structural invariants to include audio_input_stream alongside the three GGUF / Unsloth streaming generators: no _tracker.__enter__ inside the async gen body, cleanup via try/finally, no background= on StreamingResponse. - Delete test_stream_cleanup_on_disconnect.py (now empty). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: make cancel-via-POST interrupt Unsloth and audio-input streams Close two remaining gaps in the stop-button cancellation wiring: - stream_chunks (Unsloth path): add a top-of-loop cancel_event check and call backend.reset_generation_state() so cancel POSTs flush GPU state and close the SSE cleanly instead of relying on request.is_disconnected (which does not fire through proxies like Colab's). - audio_input_stream: run the synchronous audio_input_generate() via asyncio.to_thread so blocking whisper chunks do not freeze the event loop, matching the pattern already used by the GGUF streaming paths. * Add review tests for Studio stop-button cancel flow * Consolidate review tests for Studio stop-button cancel flow - Delete standalone test_cancel_registry.py at repo root: tests duplicated test_cancel_atomicity.py / test_cancel_id_wiring.py and re-implemented registry primitives inline (scaffolding). - Extend tests/studio/test_stream_cancel_registration_timing.py with regression guards for the iter-1 cancel-loop fixes: structural: each streaming generator checks cancel_event in its loop; audio_input_stream offloads next() via asyncio.to_thread; stream_chunks cancel branch calls reset_generation_state(). runtime: Unsloth loop breaks on external cancel and resets state; audio loop stays responsive under blocking next(); both loops emit zero tokens on pre-set cancel (replay path). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: extend stop-path to passthrough streams; tighten wall-clock cap - Lower _DEFAULT_T_MAX_PREDICT_MS from 1 hour to 10 minutes so the wall-clock backstop actually bounds runaway decodes when cancel signaling fails. - Wire _TrackedCancel and cancel_event.is_set() into _openai_passthrough_stream and _anthropic_passthrough_stream and disable httpx keepalive so stop requests from /v1 and /v1/messages tool-calling clients reach llama-server. - Apply t_max_predict_ms to the tool-passthrough request body so the backstop covers passthrough paths as well. - Symmetric pre-registration stash for session_id/completion_id cancels (_cancel_by_keys_or_stash) so early cancels by those keys replay on later registration like cancel_id. - Drop dead except BaseException guards around StreamingResponse() at four streaming sites; cleanup lives in the generator's finally. * studio: harden cancel registry against ghost-cancel and leak paths - Revert the session_id/completion_id stash in the fallback cancel helper. session_id is thread-scoped and reused across runs, so stashing it on an unmatched POST would fire cancel_event for the user's next unrelated request via _TrackedCancel.__enter__. cancel_id remains the only per-run unique key that gets stashed. - Default max_tokens to _DEFAULT_MAX_TOKENS in the tool-passthrough body. Mirror the direct GGUF path so OpenAI/Anthropic passthrough callers who omit max_tokens get the same zombie-decode cap instead of relying on the wall-clock backstop alone. - Wrap _openai_passthrough_stream setup with an outer try/except BaseException. The inner except httpx.RequestError does not catch asyncio.CancelledError at await client.send, which would otherwise leave _tracker registered in _CANCEL_REGISTRY indefinitely. - Frontend stop POST uses plain fetch + manual Authorization header instead of authFetch. A 401 on the cancel POST no longer refreshes tokens or redirects the user to the login page mid-stop. * Add review tests for Studio stop-button cancel flow * studio: trim comments on stop-button review changes Collapse multi-paragraph rationale blocks on the cancel registry, _openai_passthrough_stream, and the frontend onAbortCancel handler into one-line explanations of why the non-obvious behaviour exists. Drop authFetch import that became unused when the cancel POST switched to plain fetch. * Consolidate review tests for Studio stop-button cancel flow Move review-added tests out of test_cancel_dispatch_edges.py into the existing PR test files that already cover the same areas: - backend registry fan-out / exclusivity / idempotency / falsy-keys edge cases moved into tests/studio/test_cancel_atomicity.py - frontend plain-fetch (not authFetch) + manual Authorization header moved into tests/studio/test_cancel_id_wiring.py Delete the now-empty test_cancel_dispatch_edges.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: stop default-capping responses at 4096 tokens (follow-up to #5069) (#5174) * Studio: stop default-capping responses at 4096 tokens Follow-up to #5069. The 4096 default introduced for runaway-decode defense silently truncates any caller that omits max_tokens. The Studio chat UI sets params.maxTokens = loadResp.context_length after a GGUF load, so it's fine, but every other consumer is not: - OpenAI-API direct callers (/v1/chat/completions, /v1/responses, /v1/messages, /v1/completions) where the OpenAI default is effectively unlimited per response. langchain, llama-index, raw curl, and the openai SDK all rely on that. - Reasoning models. Qwen3 / gpt-oss reasoning traces routinely exceed 4096 tokens before the model emits a single visible content token. The user sees the trace cut off mid-thought. - Long-form generation ("write a chapter", "produce a full SVG"). Reproduced on this branch: gemma-4-E2B-it-GGUF Q8_0, prompt asking for a 10000-word story, no max_tokens in the request: finish_reason: stop (misleading -- should be 'length') content_chars: 19772 content_tail: ...'a comforting, yet immense, pressure.\n\n*"' Body ended mid-sentence on a stray opening quote, right at the 4096 token mark. After this patch the same request returns 38357 chars ending with '...held in a perfect, dynamic equilibrium.' -- a natural stop, not a truncation. Implementation: rename the constant to _DEFAULT_MAX_TOKENS_FLOOR and set it to 32768. Each call site now uses the model's effective context length when known, falling back to the floor: default_cap = self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR The 10-minute t_max_predict_ms wall-clock backstop from #5069 is preserved as the second line of defense. Plumbed _build_passthrough_payload + _build_openai_passthrough_body through the routes layer so the Anthropic and OpenAI passthrough paths also respect the model's context length. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Studio: cancel passthrough streams during llama-server prefill + route through apiUrl for Tauri Three reviewer-flagged correctness gaps in the stop-button mechanism. 1) `_openai_passthrough_stream` could not honor cancel during prefill. The cancel check ran inside the `async for raw_line in lines_iter` body, so a cancel POST that arrived before llama-server emitted the first SSE line was unobservable until prefill completed. With a long prompt under proxy/Colab conditions -- the exact target scenario for this PR -- that left the model decoding for a long time after the user clicked Stop. Add an asyncio watcher task that closes `resp` as soon as `cancel_event` is set, raising in `aiter_lines` so the generator can exit. The watcher polls a threading.Event because the cancel registry is keyed by threading.Event for the synchronous /cancel handler. 2) `_anthropic_passthrough_stream` had the same blocking-prefill pattern. Same fix. 3) The frontend's stop-button cancel POST used a bare relative `fetch("/api/inference/cancel", ...)`, which targets the webview origin in Tauri production builds (where the backend is at `http://127.0.0.1:8888`). Route through the existing `apiUrl()` helper from `lib/api-base.ts` to match every other Studio call. Browser/dev builds get the empty base, so behavior is unchanged there. Verified via temp/pr_simulation/sim_5069_prefill_cancel.py: cancel during prefill terminates within ~250ms on both passthrough paths (was 145s+ on the Anthropic path before this change), and the standard non-passthrough chat path still cancels with no regression. * Studio: log cancel-body parse errors instead of silently swallowing Reviewer-flagged defensive logging gap. The bare `except Exception: pass` in `cancel_inference` would mask malformed payloads that hint at a buggy client or a transport issue. Log at debug so future investigation isn't left guessing whether `body={}` came from a missing body or a parse failure. Behavior is unchanged: an unparseable body still falls through to the empty-dict path and the cancel call returns `{"cancelled": 0}`. * Studio: Anthropic passthrough cancel parity with OpenAI passthrough Two reviewer-flagged consistency gaps in the cancel surface for /v1/messages. 1) Anthropic passthrough did not register cancel_id, so a per-run cancel POST (the cleanest Studio-style cancel path) silently missed when the route hit `_anthropic_passthrough_stream`. The OpenAI passthrough has registered (cancel_id, session_id, completion_id) since this PR was first opened; mirror that here. Also add `cancel_id` to `AnthropicMessagesRequest` so the route handler can plumb it through. 2) The cancel handler's fallback key list checked only completion_id and session_id, never message_id. Anthropic clients that send their native `id` (returned in the SSE message_start event) for cancel had no way to hit the registry. Add message_id to the fallback list. Verified via temp/pr_simulation/sim_5069_prefill_cancel.py: P2 now cancels by cancel_id in 137ms (was hanging pre-fix), and the new P2b case cancels by message_id in 77ms. P1 (OpenAI) and P3 (standard chat) still pass with no regression. --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/core/inference/llama_cpp.py | 57 +- studio/backend/main.py | 4 + studio/backend/models/inference.py | 5 + studio/backend/routes/__init__.py | 2 + studio/backend/routes/inference.py | 495 +++++++++--- .../src/features/chat/api/chat-adapter.ts | 41 +- .../frontend/src/features/chat/types/api.ts | 1 + tests/studio/test_cancel_atomicity.py | 289 +++++++ tests/studio/test_cancel_id_wiring.py | 169 +++++ tests/studio/test_llama_cpp_wall_clock_cap.py | 123 +++ .../test_stream_cancel_registration_timing.py | 718 ++++++++++++++++++ 11 files changed, 1785 insertions(+), 119 deletions(-) create mode 100644 tests/studio/test_cancel_atomicity.py create mode 100644 tests/studio/test_cancel_id_wiring.py create mode 100644 tests/studio/test_llama_cpp_wall_clock_cap.py create mode 100644 tests/studio/test_stream_cancel_registration_timing.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 35ef90ee37..81ff9ae4e5 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -53,6 +53,21 @@ _INTENT_SIGNAL = re.compile( r")" ) _MAX_REPROMPTS = 3 + +# Without max_tokens, llama-server defaults to n_predict = n_ctx (up to +# 262144 for Qwen3.5), producing many-minute zombie decodes when cancel +# fails. t_max_predict_ms is a wall-clock backstop applied unconditionally, +# but the llama.cpp README notes it ONLY fires after a newline has been +# generated -- a model stuck in a long unbroken non-newline sequence is +# unbounded by it. So we still want a token cap as the front-line limiter. +# +# The cap is the model's effective context length when we know it, +# falling back to a generous floor when metadata is unavailable. 4096 was +# too low: Qwen3 / gpt-oss reasoning traces routinely exceed it, and any +# OpenAI-API caller that omits max_tokens (langchain, llama-index, raw +# curl) sees responses silently truncated mid-sentence. +_DEFAULT_MAX_TOKENS_FLOOR = 32768 +_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min _REPROMPT_MAX_CHARS = 2000 # ── Pre-compiled patterns for GGUF shard detection ─────────── @@ -1636,7 +1651,7 @@ class LlamaCppBackend: # existing text (code refactoring, summarization, reasoning). # For general chat with low repetition, overhead is ~5 ms. # - # Benchmarks from llama.cpp PRs #18471, #19164: + # Benchmarks from upstream llama.cpp speculative-decoding PRs: # 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 @@ -2549,8 +2564,15 @@ class LlamaCppBackend: ) if _reasoning_kw is not None: payload["chat_template_kwargs"] = _reasoning_kw - if max_tokens is not None: - payload["max_tokens"] = max_tokens + # Default cap to the model's effective context length when known, + # otherwise the conservative floor. The wall-clock backstop below + # keeps a stuck model from running indefinitely either way. + payload["max_tokens"] = ( + max_tokens + if max_tokens is not None + else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) + ) + payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop payload["stream_options"] = {"include_usage": True} @@ -2570,7 +2592,9 @@ class LlamaCppBackend: _auth_headers = ( {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None ) - with httpx.Client(timeout = stream_timeout) as client: + with httpx.Client( + timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) + ) as client: with self._stream_with_retry( client, url, @@ -2769,8 +2793,12 @@ class LlamaCppBackend: ) if _reasoning_kw is not None: payload["chat_template_kwargs"] = _reasoning_kw - if max_tokens is not None: - payload["max_tokens"] = max_tokens + payload["max_tokens"] = ( + max_tokens + if max_tokens is not None + else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) + ) + payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: payload["stop"] = stop @@ -2809,7 +2837,10 @@ class LlamaCppBackend: write = 10, pool = 10, ) - with httpx.Client(timeout = stream_timeout) as client: + with httpx.Client( + timeout = stream_timeout, + limits = httpx.Limits(max_keepalive_connections = 0), + ) as client: with self._stream_with_retry( client, url, @@ -3422,8 +3453,12 @@ class LlamaCppBackend: ) if _reasoning_kw is not None: stream_payload["chat_template_kwargs"] = _reasoning_kw - if max_tokens is not None: - stream_payload["max_tokens"] = max_tokens + stream_payload["max_tokens"] = ( + max_tokens + if max_tokens is not None + else (self._effective_context_length or _DEFAULT_MAX_TOKENS_FLOOR) + ) + stream_payload["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: stream_payload["stop"] = stop stream_payload["stream_options"] = {"include_usage": True} @@ -3442,7 +3477,9 @@ class LlamaCppBackend: _auth_headers = ( {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None ) - with httpx.Client(timeout = stream_timeout) as client: + with httpx.Client( + timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0) + ) as client: with self._stream_with_retry( client, url, diff --git a/studio/backend/main.py b/studio/backend/main.py index 05adcaa2ea..9212404b30 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -62,6 +62,7 @@ from routes import ( datasets_router, export_router, inference_router, + inference_studio_router, models_router, training_history_router, training_router, @@ -207,6 +208,9 @@ app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"]) app.include_router(training_router, prefix = "/api/train", tags = ["training"]) app.include_router(models_router, prefix = "/api/models", tags = ["models"]) app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"]) +# Studio-only inference endpoints (cancel, etc.) are intentionally NOT +# exposed on the /v1 OpenAI-compat prefix below. +app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"]) # OpenAI-compatible endpoints: mount the same inference router at /v1 # so external tools (Open WebUI, SillyTavern, etc.) can use the diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index e5b037755d..bf0177efbf 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -531,6 +531,10 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.", ) + cancel_id: Optional[str] = Field( + None, + description = "[x-unsloth] Per-request cancellation token. Frontend sends a fresh UUID per run so /inference/cancel matches one specific generation.", + ) # ── Streaming response chunks ──────────────────────────────────── @@ -992,6 +996,7 @@ class AnthropicMessagesRequest(BaseModel): enable_tools: Optional[bool] = None enabled_tools: Optional[list[str]] = None session_id: Optional[str] = None + cancel_id: Optional[str] = None model_config = {"extra": "allow"} diff --git a/studio/backend/routes/__init__.py b/studio/backend/routes/__init__.py index e79f6553f9..cf4586281b 100644 --- a/studio/backend/routes/__init__.py +++ b/studio/backend/routes/__init__.py @@ -8,6 +8,7 @@ API Routes from routes.training import router as training_router from routes.models import router as models_router from routes.inference import router as inference_router +from routes.inference import studio_router as inference_studio_router from routes.datasets import router as datasets_router from routes.auth import router as auth_router from routes.data_recipe import router as data_recipe_router @@ -18,6 +19,7 @@ __all__ = [ "training_router", "models_router", "inference_router", + "inference_studio_router", "datasets_router", "auth_router", "data_recipe_router", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ed331a5660..cf3b37a2fd 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -113,7 +113,12 @@ if str(backend_path) not in sys.path: # Import backend functions try: from core.inference import get_inference_backend - from core.inference.llama_cpp import LlamaCppBackend, detect_reasoning_flags + from core.inference.llama_cpp import ( + LlamaCppBackend, + _DEFAULT_MAX_TOKENS_FLOOR, + _DEFAULT_T_MAX_PREDICT_MS, + detect_reasoning_flags, + ) from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import load_model_defaults @@ -122,7 +127,12 @@ except ImportError: if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) from core.inference import get_inference_backend - from core.inference.llama_cpp import LlamaCppBackend, detect_reasoning_flags + from core.inference.llama_cpp import ( + LlamaCppBackend, + _DEFAULT_MAX_TOKENS_FLOOR, + _DEFAULT_T_MAX_PREDICT_MS, + detect_reasoning_flags, + ) from utils.models import ModelConfig from utils.inference import load_inference_config from utils.models.model_config import load_model_defaults @@ -185,6 +195,126 @@ import numpy as np from datetime import date as _date router = APIRouter() +# Studio-only router (not mounted on /v1 OpenAI-compat). +studio_router = APIRouter() + + +# Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts +# so is_disconnected() never fires. POST /inference/cancel looks up +# in-flight cancel_events here by cancel_id (per-run) or session_id / +# completion_id (fallbacks). +_CANCEL_REGISTRY: dict[str, set[threading.Event]] = {} +_CANCEL_LOCK = threading.Lock() + +# Cancel POSTs that arrive before registration are stashed; the next +# matching __enter__ replays set() within the TTL. +_PENDING_CANCELS: dict[str, float] = {} +_PENDING_CANCEL_TTL_S = 30.0 + + +def _prune_pending(now: float) -> None: + for k in [ + k for k, ts in _PENDING_CANCELS.items() if now - ts > _PENDING_CANCEL_TTL_S + ]: + _PENDING_CANCELS.pop(k, None) + + +class _TrackedCancel: + """Register cancel_event in _CANCEL_REGISTRY for the block's duration.""" + + def __init__(self, event: threading.Event, *keys): + self.event = event + self.keys = tuple(k for k in keys if k) + + def __enter__(self): + # Register + consume-pending must be one critical section to close + # the TOCTOU race against a concurrent cancel POST. + should_cancel = False + with _CANCEL_LOCK: + for k in self.keys: + _CANCEL_REGISTRY.setdefault(k, set()).add(self.event) + now = time.monotonic() + _prune_pending(now) + for k in self.keys: + if k and _PENDING_CANCELS.pop(k, None) is not None: + should_cancel = True + if should_cancel: + self.event.set() + return self.event + + def __exit__(self, *exc): + with _CANCEL_LOCK: + for k in self.keys: + bucket = _CANCEL_REGISTRY.get(k) + if bucket is None: + continue + bucket.discard(self.event) + if not bucket: + _CANCEL_REGISTRY.pop(k, None) + return False + + +def _cancel_by_keys(keys) -> int: + """Set cancel_event for matching registry entries; no stash. + session_id/completion_id are shared across runs on the same thread, + so stashing them would ghost-cancel the user's next request. Only + cancel_id is per-run unique (see _cancel_by_cancel_id_or_stash).""" + if not keys: + return 0 + events: set[threading.Event] = set() + with _CANCEL_LOCK: + _prune_pending(time.monotonic()) + for k in keys: + bucket = _CANCEL_REGISTRY.get(k) + if bucket: + events.update(bucket) + for ev in events: + ev.set() + return len(events) + + +def _cancel_by_cancel_id_or_stash(cancel_id: str) -> int: + """Atomic lookup-or-stash; pairs with _TrackedCancel.__enter__ to + close the TOCTOU race.""" + now = time.monotonic() + events: set[threading.Event] = set() + with _CANCEL_LOCK: + _prune_pending(now) + bucket = _CANCEL_REGISTRY.get(cancel_id) + if bucket: + events.update(bucket) + else: + _PENDING_CANCELS[cancel_id] = now + for ev in events: + ev.set() + return len(events) + + +async def _await_cancel_then_close(cancel_event, resp) -> None: + """Watch a threading.Event from asyncio and close ``resp`` when it fires. + + Used by the passthrough streamers so a /cancel POST can interrupt + while the async iterator is blocked waiting for llama-server prefill. + Without this watcher the in-loop ``cancel_event.is_set()`` check is + unreachable until the first SSE chunk arrives, which is exactly the + proxy/Colab scenario the cancel POST exists to handle. + + Polls a threading.Event because the cancel registry is keyed by + threading.Event so the synchronous /cancel handler can call .set(). + 50ms cadence adds at most that much latency to a prefill cancel; the + common-case streaming cancel path still observes the event in the + iterator's first iteration after the next chunk. + """ + try: + while not cancel_event.is_set(): + await asyncio.sleep(0.05) + try: + await resp.aclose() + except Exception: + pass + except asyncio.CancelledError: + return + # Appended to tool-use nudge to discourage plan-without-action _TOOL_ACTION_NUDGE = ( @@ -706,6 +836,48 @@ async def unload_model( raise HTTPException(status_code = 500, detail = f"Failed to unload model: {str(e)}") +@studio_router.post("/cancel") +async def cancel_inference( + request: Request, + current_subject: str = Depends(get_current_subject), +): + """Cancel in-flight inference requests. + + Body (JSON, at least one key required): + cancel_id - preferred: per-run UUID, matched exclusively. + session_id - fallback when cancel_id is absent. + completion_id - fallback when cancel_id is absent. + + A cancel_id arriving before its stream registers is stashed briefly + and replayed on registration. Returns {"cancelled": N}. + """ + try: + body = await request.json() + if not isinstance(body, dict): + body = {} + except Exception as e: + logger.debug("Failed to parse cancel request body: %s", e) + body = {} + + cancel_id = body.get("cancel_id") + if isinstance(cancel_id, str) and cancel_id: + return {"cancelled": _cancel_by_cancel_id_or_stash(cancel_id)} + + keys = [] + # `message_id` is the Anthropic passthrough's per-run identifier -- + # included so /v1/messages clients can cancel by their native id. + for k in ("completion_id", "session_id", "message_id"): + v = body.get(k) + if isinstance(v, str) and v: + keys.append(v) + + if not keys: + return {"cancelled": 0} + + n = _cancel_by_keys(keys) + return {"cancelled": n} + + @router.post("/generate/stream") async def generate_stream( request: GenerateRequest, @@ -1169,6 +1341,9 @@ async def openai_chat_completions( ) if payload.stream: + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() async def audio_input_stream(): try: @@ -1185,10 +1360,17 @@ async def openai_chat_completions( ) yield f"data: {first_chunk.model_dump_json(exclude_none = True)}\n\n" - for chunk_text in audio_input_generate(): + gen = audio_input_generate() + _DONE = object() + while True: + if cancel_event.is_set(): + break if await request.is_disconnected(): cancel_event.set() return + chunk_text = await asyncio.to_thread(next, gen, _DONE) + if chunk_text is _DONE: + break if chunk_text: chunk = ChatCompletionChunk( id = completion_id, @@ -1221,6 +1403,8 @@ async def openai_chat_completions( f"Error during audio input streaming: {e}", exc_info = True ) yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" + finally: + _tracker.__exit__(None, None, None) return StreamingResponse( audio_input_stream(), @@ -1466,6 +1650,10 @@ async def openai_chat_completions( _tool_sentinel = object() + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + async def gguf_tool_stream(): try: first_chunk = ChatCompletionChunk( @@ -1488,6 +1676,8 @@ async def openai_chat_completions( _stream_usage = None _stream_timings = None while True: + if cancel_event.is_set(): + break if await request.is_disconnected(): cancel_event.set() return @@ -1595,6 +1785,8 @@ async def openai_chat_completions( }, } yield f"data: {json.dumps(error_chunk)}\n\n" + finally: + _tracker.__exit__(None, None, None) return StreamingResponse( gguf_tool_stream(), @@ -1628,6 +1820,9 @@ async def openai_chat_completions( _gguf_sentinel = object() if payload.stream: + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() async def gguf_stream_chunks(): try: @@ -1652,6 +1847,8 @@ async def openai_chat_completions( _stream_usage = None _stream_timings = None while True: + if cancel_event.is_set(): + break if await request.is_disconnected(): cancel_event.set() return @@ -1735,6 +1932,8 @@ async def openai_chat_completions( }, } yield f"data: {json.dumps(error_chunk)}\n\n" + finally: + _tracker.__exit__(None, None, None) return StreamingResponse( gguf_stream_chunks(), @@ -1834,6 +2033,9 @@ async def openai_chat_completions( # ── Streaming response ──────────────────────────────────────── if payload.stream: + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() async def stream_chunks(): try: @@ -1861,6 +2063,9 @@ async def openai_chat_completions( loop = asyncio.get_event_loop() gen = generate() while True: + if cancel_event.is_set(): + backend.reset_generation_state() + break # next(gen, _DONE) returns _DONE instead of raising # StopIteration — StopIteration cannot propagate # through asyncio futures (Python limitation). @@ -1916,6 +2121,8 @@ async def openai_chat_completions( }, } yield f"data: {json.dumps(error_chunk)}\n\n" + finally: + _tracker.__exit__(None, None, None) return StreamingResponse( stream_chunks(), @@ -2596,7 +2803,9 @@ async def _responses_stream( ), ) - body = _build_openai_passthrough_body(chat_req) + body = _build_openai_passthrough_body( + chat_req, backend_ctx = llama_backend.context_length + ) target_url = f"{llama_backend.base_url}/v1/chat/completions" async def event_generator(): @@ -3081,6 +3290,8 @@ async def anthropic_messages( repetition_penalty = repetition_penalty, presence_penalty = presence_penalty, tool_choice = openai_tool_choice, + session_id = payload.session_id, + cancel_id = payload.cancel_id, ) return await _anthropic_passthrough_non_streaming( llama_backend, @@ -3441,6 +3652,7 @@ def _build_passthrough_payload( repetition_penalty = None, presence_penalty = None, tool_choice = "auto", + backend_ctx = None, ): body = { "messages": openai_messages, @@ -3453,8 +3665,12 @@ def _build_passthrough_payload( } if stream: body["stream_options"] = {"include_usage": True} - if max_tokens is not None: - body["max_tokens"] = max_tokens + body["max_tokens"] = ( + max_tokens + if max_tokens is not None + else (backend_ctx or _DEFAULT_MAX_TOKENS_FLOOR) + ) + body["t_max_predict_ms"] = _DEFAULT_T_MAX_PREDICT_MS if stop: body["stop"] = stop if min_p is not None: @@ -3484,6 +3700,8 @@ async def _anthropic_passthrough_stream( repetition_penalty = None, presence_penalty = None, tool_choice = "auto", + session_id = None, + cancel_id = None, ): """Streaming client-side pass-through: forward tools to llama-server and translate its streaming response to Anthropic SSE without executing anything.""" @@ -3501,8 +3719,14 @@ async def _anthropic_passthrough_stream( repetition_penalty = repetition_penalty, presence_penalty = presence_penalty, tool_choice = tool_choice, + backend_ctx = llama_backend.context_length, ) + # cancel_id mirrors the OpenAI passthrough so a per-run cancel POST + # works without the caller having to know the local message_id. + _tracker = _TrackedCancel(cancel_event, cancel_id, session_id, message_id) + _tracker.__enter__() + async def _stream(): emitter = AnthropicPassthroughEmitter() for line in emitter.start(message_id, model_name): @@ -3535,15 +3759,28 @@ async def _anthropic_passthrough_stream( # has anything orphaned to finalize. Each aclose is wrapped in # `try: ... except Exception: pass` so anyio cleanup noise from # nested aclose paths can't bubble out. - client = httpx.AsyncClient(timeout = 600) + client = httpx.AsyncClient( + timeout = 600, + limits = httpx.Limits(max_keepalive_connections = 0), + ) resp = None lines_iter = None + cancel_watcher = None try: req = client.build_request("POST", target_url, json = body) resp = await client.send(req, stream = True) + # See _openai_passthrough_stream for rationale: aiter_lines() + # blocks during llama-server prefill, so the in-loop cancel + # check is unreachable until the first SSE chunk arrives. + # The watcher closes `resp` on cancel, raising in aiter_lines. + cancel_watcher = asyncio.create_task( + _await_cancel_then_close(cancel_event, resp) + ) lines_iter = resp.aiter_lines() async for raw_line in lines_iter: + if cancel_event.is_set(): + break if await request.is_disconnected(): cancel_event.set() break @@ -3558,9 +3795,18 @@ async def _anthropic_passthrough_stream( continue for line in emitter.feed_chunk(chunk): yield line + except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): + if not cancel_event.is_set(): + raise except Exception as e: logger.error("anthropic_messages passthrough stream error: %s", e) finally: + if cancel_watcher is not None: + cancel_watcher.cancel() + try: + await cancel_watcher + except (asyncio.CancelledError, Exception): + pass if lines_iter is not None: try: await lines_iter.aclose() @@ -3575,6 +3821,7 @@ async def _anthropic_passthrough_stream( await client.aclose() except Exception: pass + _tracker.__exit__(None, None, None) for line in emitter.finish(): yield line @@ -3621,6 +3868,7 @@ async def _anthropic_passthrough_non_streaming( repetition_penalty = repetition_penalty, presence_penalty = presence_penalty, tool_choice = tool_choice, + backend_ctx = llama_backend.context_length, ) async with httpx.AsyncClient() as client: @@ -3742,7 +3990,7 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages -def _build_openai_passthrough_body(payload) -> dict: +def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: """Assemble the llama-server request body from a ChatCompletionRequest. Only explicitly-known OpenAI / llama-server fields are forwarded so that @@ -3764,6 +4012,7 @@ def _build_openai_passthrough_body(payload) -> dict: repetition_penalty = payload.repetition_penalty, presence_penalty = payload.presence_penalty, tool_choice = tool_choice, + backend_ctx = backend_ctx, ) @@ -3784,103 +4033,56 @@ async def _openai_passthrough_stream( observes a standard OpenAI response. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body(payload) + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length + ) - # Dispatch the upstream request BEFORE returning StreamingResponse so - # transport errors and non-200 upstream statuses surface as real HTTP - # errors to the client. OpenAI SDKs rely on status codes to raise - # ``APIError``/``BadRequestError``/...; burying the failure inside a - # 200 SSE ``error`` frame silently breaks their error handling. - client = httpx.AsyncClient(timeout = 600) - resp = None + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + + # Outer guard: asyncio.CancelledError at `await client.send(...)` is + # a BaseException that bypasses `except httpx.RequestError`; without + # this the tracker leaks. The generator's finally only runs once + # iteration starts. try: - req = client.build_request("POST", target_url, json = body) - resp = await client.send(req, stream = True) - except httpx.RequestError as e: - # llama-server subprocess crashed / still starting / unreachable. - logger.error("openai passthrough stream: upstream unreachable: %s", e) - if resp is not None: - try: - await resp.aclose() - except Exception: - pass - try: - await client.aclose() - except Exception: - pass - raise HTTPException( - status_code = 502, - detail = _friendly_error(e), + # Dispatch BEFORE returning StreamingResponse so transport errors + # and non-200 upstream statuses surface as real HTTP errors -- + # OpenAI SDKs rely on status codes to raise APIError/BadRequestError. + client = httpx.AsyncClient( + timeout = 600, + limits = httpx.Limits(max_keepalive_connections = 0), ) - - if resp.status_code != 200: - err_bytes = await resp.aread() - err_text = err_bytes.decode("utf-8", errors = "replace") - logger.error( - "openai passthrough upstream error: status=%s body=%s", - resp.status_code, - err_text[:500], - ) - upstream_status = resp.status_code + resp = None try: - await resp.aclose() - except Exception: - pass - try: - await client.aclose() - except Exception: - pass - raise HTTPException( - status_code = upstream_status, - detail = f"llama-server error: {err_text[:500]}", - ) - - async def _stream(): - # Same httpx lifecycle pattern as _anthropic_passthrough_stream: - # avoid `async with` on the client/response AND explicitly save - # resp.aiter_lines() so we can close it ourselves in the finally - # block. See the long comment there for the full rationale on - # why the anonymous `async for raw_line in resp.aiter_lines():` - # pattern leaks an unclosed async generator that Python's - # asyncgen GC hook then finalizes in a different asyncio task, - # producing "Exception ignored in:" / "async generator ignored - # GeneratorExit" / anyio cancel-scope traces on Python 3.13 + - # httpcore 1.0.x. - lines_iter = None - try: - lines_iter = resp.aiter_lines() - async for raw_line in lines_iter: - if await request.is_disconnected(): - cancel_event.set() - break - if not raw_line: - continue - if not raw_line.startswith("data: "): - continue - # Relay the llama-server SSE chunk verbatim so the client - # sees its native `id`, `finish_reason`, `delta.tool_calls`, - # and final `usage` unchanged. - yield raw_line + "\n\n" - if raw_line[6:].strip() == "[DONE]": - break - except Exception as e: - # Mid-stream failures still have to be reported inside the SSE - # body because the 200 response headers have already been - # committed by the time the first chunk flushes. - logger.error("openai passthrough stream error: %s", e) - err = { - "error": { - "message": _friendly_error(e), - "type": "server_error", - }, - } - yield f"data: {json.dumps(err)}\n\n" - finally: - if lines_iter is not None: + req = client.build_request("POST", target_url, json = body) + resp = await client.send(req, stream = True) + except httpx.RequestError as e: + # llama-server subprocess crashed / still starting / unreachable. + logger.error("openai passthrough stream: upstream unreachable: %s", e) + if resp is not None: try: - await lines_iter.aclose() + await resp.aclose() except Exception: pass + try: + await client.aclose() + except Exception: + pass + raise HTTPException( + status_code = 502, + detail = _friendly_error(e), + ) + + if resp.status_code != 200: + err_bytes = await resp.aread() + err_text = err_bytes.decode("utf-8", errors = "replace") + logger.error( + "openai passthrough upstream error: status=%s body=%s", + resp.status_code, + err_text[:500], + ) + upstream_status = resp.status_code try: await resp.aclose() except Exception: @@ -3889,16 +4091,91 @@ async def _openai_passthrough_stream( await client.aclose() except Exception: pass + raise HTTPException( + status_code = upstream_status, + detail = f"llama-server error: {err_text[:500]}", + ) - return StreamingResponse( - _stream(), - media_type = "text/event-stream", - headers = { - "Cache-Control": "no-cache", - "Connection": "keep-alive", - "X-Accel-Buffering": "no", - }, - ) + async def _stream(): + # Same httpx lifecycle pattern as _anthropic_passthrough_stream: + # save resp.aiter_lines() so the finally block can aclose() it + # on our task. See that function for full rationale. + lines_iter = None + # During llama-server prefill, `aiter_lines()` blocks until the + # first SSE chunk arrives. The in-loop `cancel_event` check + # cannot fire until then, which is the exact proxy/Colab + # scenario the cancel POST is meant to recover from. Run a + # tiny watcher that closes `resp` as soon as cancel fires, + # unblocking the iterator with a RemoteProtocolError caught + # in the except clause below. + cancel_watcher = asyncio.create_task( + _await_cancel_then_close(cancel_event, resp) + ) + try: + lines_iter = resp.aiter_lines() + async for raw_line in lines_iter: + if cancel_event.is_set(): + break + if await request.is_disconnected(): + cancel_event.set() + break + if not raw_line: + continue + if not raw_line.startswith("data: "): + continue + # Relay verbatim to preserve llama-server's native id, + # finish_reason, delta.tool_calls, and usage chunks. + yield raw_line + "\n\n" + if raw_line[6:].strip() == "[DONE]": + break + except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): + # Watcher closed resp on cancel. Emit nothing extra; the + # client either initiated the cancel or already disconnected. + if not cancel_event.is_set(): + raise + except Exception as e: + # 200 headers are already flushed; errors must be in the SSE body. + logger.error("openai passthrough stream error: %s", e) + err = { + "error": { + "message": _friendly_error(e), + "type": "server_error", + }, + } + yield f"data: {json.dumps(err)}\n\n" + finally: + cancel_watcher.cancel() + try: + await cancel_watcher + except (asyncio.CancelledError, Exception): + pass + if lines_iter is not None: + try: + await lines_iter.aclose() + except Exception: + pass + try: + await resp.aclose() + except Exception: + pass + try: + await client.aclose() + except Exception: + pass + _tracker.__exit__(None, None, None) + + return StreamingResponse( + _stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + except BaseException: + _tracker.__exit__(None, None, None) + raise async def _openai_passthrough_non_streaming( @@ -3914,7 +4191,9 @@ async def _openai_passthrough_non_streaming( token counts. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body(payload) + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length + ) try: async with httpx.AsyncClient() as client: diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index a93120397e..dde597ca11 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -4,6 +4,8 @@ import type { ChatModelAdapter } from "@assistant-ui/react"; import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core"; import { toast } from "sonner"; +import { getAuthToken } from "@/features/auth/session"; +import { apiUrl } from "@/lib/api-base"; import { generateAudio, listCachedGguf, @@ -707,7 +709,42 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const toolCallParts: ToolCallMessagePart[] = []; let serverMetadata: { usage?: ServerUsage; timings?: ServerTimings } | null = null; + // Per-run cancellation token so a delayed stop POST cannot match + // the next run on the same thread. + const cancelId = + typeof crypto !== "undefined" && "randomUUID" in crypto + ? crypto.randomUUID() + : `${Date.now()}-${Math.random().toString(36).slice(2)}`; + + // Colab-style proxies can swallow fetch aborts, so also POST + // /inference/cancel explicitly on abort. + const onAbortCancel = () => { + const body: Record = { cancel_id: cancelId }; + if (resolvedThreadId) body.session_id = resolvedThreadId; + // Plain fetch, not authFetch: authFetch redirects to login on + // 401, which would kick the user out mid-stop. + const token = getAuthToken(); + // Use apiUrl so the cancel POST reaches the right origin in + // Tauri production builds (where the webview origin is not the + // backend at 127.0.0.1:). Browser/dev builds get the empty + // base, so the path is unchanged there. + void fetch(apiUrl("/api/inference/cancel"), { + method: "POST", + headers: { + "Content-Type": "application/json", + ...(token ? { Authorization: `Bearer ${token}` } : {}), + }, + body: JSON.stringify(body), + keepalive: true, + }).catch(() => {}); + }; try { + if (abortSignal.aborted) { + onAbortCancel(); + } else { + abortSignal.addEventListener("abort", onAbortCancel, { once: true }); + } + const { supportsReasoning, reasoningEnabled, @@ -730,6 +767,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { presence_penalty: params.presencePenalty, image_base64: imageBase64, audio_base64: audioBase64, + cancel_id: cancelId, + ...(resolvedThreadId ? { session_id: resolvedThreadId } : {}), ...(useAdapter === undefined ? {} : { use_adapter: useAdapter }), ...(supportsReasoning ? reasoningStyle === "reasoning_effort" @@ -750,7 +789,6 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const mins = useChatRuntimeStore.getState().toolCallTimeout; return mins >= 9999 ? 9999 : mins * 60; })(), - session_id: resolvedThreadId, } : {}), }, @@ -948,6 +986,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { } throw err; } finally { + abortSignal.removeEventListener("abort", onAbortCancel); runtime.setGeneratingStatus(null); runtime.setToolStatus(null); clearTimeout(warmupTimer); diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index ccfc8b2bf1..25957f4a7b 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -179,6 +179,7 @@ export interface OpenAIChatCompletionsRequest { max_tool_calls_per_message?: number; tool_call_timeout?: number; session_id?: string; + cancel_id?: string; } export interface OpenAIChatDelta { diff --git a/tests/studio/test_cancel_atomicity.py b/tests/studio/test_cancel_atomicity.py new file mode 100644 index 0000000000..a8d4839454 --- /dev/null +++ b/tests/studio/test_cancel_atomicity.py @@ -0,0 +1,289 @@ +""" +TOCTOU atomicity guards for the cancel path. + +Structural: cancel_inference, _cancel_by_cancel_id_or_stash, and +_TrackedCancel.__enter__ must each use a single _CANCEL_LOCK critical +section over lookup + stash / register + consume-pending. + +Behavioral: parallel cancel-POST vs __enter__ must never drop a cancel. +""" + +from __future__ import annotations + +import ast +import random +import threading +from pathlib import Path + + +SOURCE_PATH = ( + Path(__file__).resolve().parents[2] + / "studio" + / "backend" + / "routes" + / "inference.py" +) +_SRC = SOURCE_PATH.read_text() +_TREE = ast.parse(_SRC) + + +def _find_function(name: str) -> ast.FunctionDef | ast.AsyncFunctionDef: + for node in ast.walk(_TREE): + if ( + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == name + ): + return node + raise AssertionError(f"function {name!r} not found") + + +def _find_class(name: str) -> ast.ClassDef: + for node in ast.walk(_TREE): + if isinstance(node, ast.ClassDef) and node.name == name: + return node + raise AssertionError(f"class {name!r} not found") + + +def _count_with_cancel_lock_blocks(node: ast.AST) -> int: + n = 0 + for sub in ast.walk(node): + if not isinstance(sub, ast.With): + continue + for item in sub.items: + ctx = item.context_expr + if isinstance(ctx, ast.Name) and ctx.id == "_CANCEL_LOCK": + n += 1 + break + return n + + +def test_cancel_by_cancel_id_or_stash_is_single_lock_critical_section(): + fn = _find_function("_cancel_by_cancel_id_or_stash") + assert _count_with_cancel_lock_blocks(fn) == 1, ( + "_cancel_by_cancel_id_or_stash must use exactly one `with " + "_CANCEL_LOCK:` block; splitting into two acquisitions reopens " + "the TOCTOU race with _TrackedCancel.__enter__" + ) + src = ast.unparse(fn) + assert "_CANCEL_REGISTRY.get(cancel_id)" in src + assert "_PENDING_CANCELS[cancel_id]" in src + + +def test_tracked_cancel_enter_registers_and_consumes_pending_under_one_lock(): + cls = _find_class("_TrackedCancel") + enter = None + for n in cls.body: + if isinstance(n, ast.FunctionDef) and n.name == "__enter__": + enter = n + break + assert enter is not None + assert _count_with_cancel_lock_blocks(enter) == 1, ( + "_TrackedCancel.__enter__ must acquire _CANCEL_LOCK exactly once. " + "A second acquisition for consume-pending lets a concurrent " + "cancel POST stash after consume sees an empty map, silently " + "dropping the cancel" + ) + with_block = None + for sub in ast.walk(enter): + if isinstance(sub, ast.With) and any( + isinstance(i.context_expr, ast.Name) and i.context_expr.id == "_CANCEL_LOCK" + for i in sub.items + ): + with_block = sub + break + assert with_block is not None + block_src = "\n".join(ast.unparse(s) for s in with_block.body) + assert "_CANCEL_REGISTRY.setdefault" in block_src + assert "_PENDING_CANCELS.pop" in block_src, ( + "__enter__ critical section must consume from _PENDING_CANCELS " + "inside the same lock, not a later re-acquisition" + ) + + +def test_cancel_inference_uses_atomic_helper_for_cancel_id_path(): + fn = _find_function("cancel_inference") + src = ast.unparse(fn) + assert "_cancel_by_cancel_id_or_stash" in src + # The pre-fix two-step idiom must be gone. + assert "_remember_pending_cancel(cancel_id)" not in src, ( + "two-step _cancel_by_keys + _remember_pending_cancel produced " + "the TOCTOU race and must not return" + ) + + +_WANTED = { + "_CANCEL_REGISTRY", + "_CANCEL_LOCK", + "_PENDING_CANCELS", + "_PENDING_CANCEL_TTL_S", + "_prune_pending", + "_remember_pending_cancel", + "_TrackedCancel", + "_cancel_by_keys", + "_cancel_by_cancel_id_or_stash", +} + + +def _load_registry_module(): + chunks = [] + for n in _TREE.body: + seg = ast.get_source_segment(_SRC, n) + if seg is None: + continue + if isinstance(n, (ast.FunctionDef, ast.ClassDef)) and n.name in _WANTED: + chunks.append(seg) + elif isinstance(n, ast.Assign): + names = [t.id for t in n.targets if isinstance(t, ast.Name)] + if any(name in _WANTED for name in names): + chunks.append(seg) + elif ( + isinstance(n, ast.AnnAssign) + and isinstance(n.target, ast.Name) + and n.target.id in _WANTED + ): + chunks.append(seg) + mod = {} + exec( + "import threading, time\nfrom typing import Optional\n" + "\n\n".join(chunks), + mod, + ) + return mod + + +def test_parallel_cancel_vs_register_never_drops(): + m = _load_registry_module() + trials = 500 + dropped = 0 + for i in range(trials): + m["_CANCEL_REGISTRY"].clear() + m["_PENDING_CANCELS"].clear() + cid = f"cid-{i}" + ev = threading.Event() + tracker = m["_TrackedCancel"](ev, cid, "thread") + start = threading.Event() + + def do_cancel(): + start.wait() + m["_cancel_by_cancel_id_or_stash"](cid) + + def do_enter(): + start.wait() + tracker.__enter__() + + threads = [ + threading.Thread(target = do_cancel), + threading.Thread(target = do_enter), + ] + random.shuffle(threads) + for t in threads: + t.start() + start.set() + for t in threads: + t.join(timeout = 5.0) + assert not t.is_alive() + + if not ev.is_set(): + dropped += 1 + tracker.__exit__(None, None, None) + + assert dropped == 0, ( + f"TOCTOU regression: {dropped}/{trials} parallel trials silently " + f"dropped the cancel" + ) + + +def test_cancel_before_register_replays_atomically(): + m = _load_registry_module() + cid = "early-cid" + ev = threading.Event() + tracker = m["_TrackedCancel"](ev, cid, "thread-x") + + assert m["_cancel_by_cancel_id_or_stash"](cid) == 0 + assert cid in m["_PENDING_CANCELS"] + + tracker.__enter__() + assert ev.is_set() + assert cid not in m["_PENDING_CANCELS"] + tracker.__exit__(None, None, None) + + +def test_cancel_after_register_signals_without_stash(): + m = _load_registry_module() + cid = "post-cid" + ev = threading.Event() + tracker = m["_TrackedCancel"](ev, cid, "thread-y") + tracker.__enter__() + + assert m["_cancel_by_cancel_id_or_stash"](cid) == 1 + assert ev.is_set() + assert cid not in m["_PENDING_CANCELS"] + tracker.__exit__(None, None, None) + + +def test_cancel_by_keys_tolerates_empty_and_falsy_keys(): + m = _load_registry_module() + m["_CANCEL_REGISTRY"].clear() + m["_PENDING_CANCELS"].clear() + assert m["_cancel_by_keys"]([]) == 0 + assert m["_cancel_by_keys"](["", None, "unknown"]) == 0 + # Non-stashing fallback must never leak into _PENDING_CANCELS. + assert m["_PENDING_CANCELS"] == {} + + +def test_cancel_by_keys_fans_out_to_all_streams_on_same_session(): + # Compare mode and other flows launch concurrent streams under a + # shared session_id; a single session cancel POST must hit all of them. + m = _load_registry_module() + m["_CANCEL_REGISTRY"].clear() + m["_PENDING_CANCELS"].clear() + session = "shared-thread" + ev_a = threading.Event() + ev_b = threading.Event() + tracker_a = m["_TrackedCancel"](ev_a, "cancel-a", session, "chatcmpl-a") + tracker_b = m["_TrackedCancel"](ev_b, "cancel-b", session, "chatcmpl-b") + tracker_a.__enter__() + tracker_b.__enter__() + try: + assert m["_cancel_by_keys"]([session]) == 2 + assert ev_a.is_set() and ev_b.is_set() + finally: + tracker_a.__exit__(None, None, None) + tracker_b.__exit__(None, None, None) + assert session not in m["_CANCEL_REGISTRY"] + + +def test_cancel_by_cancel_id_is_exclusive_to_single_run(): + # cancel_id is per-run unique; cancelling run A must not touch run B + # even when both share a session_id. + m = _load_registry_module() + m["_CANCEL_REGISTRY"].clear() + m["_PENDING_CANCELS"].clear() + session = "shared-thread-2" + ev_a = threading.Event() + ev_b = threading.Event() + tracker_a = m["_TrackedCancel"](ev_a, "cancel-only-a", session, "chatcmpl-a") + tracker_b = m["_TrackedCancel"](ev_b, "cancel-only-b", session, "chatcmpl-b") + tracker_a.__enter__() + tracker_b.__enter__() + try: + assert m["_cancel_by_cancel_id_or_stash"]("cancel-only-a") == 1 + assert ev_a.is_set() + assert not ev_b.is_set() + finally: + tracker_a.__exit__(None, None, None) + tracker_b.__exit__(None, None, None) + + +def test_tracked_cancel_exit_is_idempotent(): + # Outer except BaseException + the generator's finally may both call + # __exit__ under certain race combos; must not raise. + m = _load_registry_module() + m["_CANCEL_REGISTRY"].clear() + m["_PENDING_CANCELS"].clear() + ev = threading.Event() + tracker = m["_TrackedCancel"](ev, "cid", "sess", "chatcmpl-x") + tracker.__enter__() + tracker.__exit__(None, None, None) + tracker.__exit__(None, None, None) + tracker.__exit__(None, None, None) + assert not m["_CANCEL_REGISTRY"] diff --git a/tests/studio/test_cancel_id_wiring.py b/tests/studio/test_cancel_id_wiring.py new file mode 100644 index 0000000000..5fd76ded9b --- /dev/null +++ b/tests/studio/test_cancel_id_wiring.py @@ -0,0 +1,169 @@ +""" +Wiring tests for the per-run cancel_id field. + +A chat-thread-scoped session_id is not safe as a cancel key because a +late stop POST can match a subsequent run on the same thread. The fix +adds cancel_id (a fresh UUID per generation) that is sent both in the +completion payload and in the /api/inference/cancel body. + +Verifies: + - ChatCompletionRequest exposes an Optional[str] `cancel_id` field. + - /api/inference/cancel accepts `cancel_id` as the first-preferred key. + - OpenAIChatCompletionsRequest (frontend type) includes cancel_id. + - chat-adapter.ts generates a per-run cancelId (crypto.randomUUID + with a Math.random fallback), sends it in the completion payload, + and includes it in the /inference/cancel body on abort. +""" + +from __future__ import annotations + +import ast +import re +from pathlib import Path + + +WORKSPACE = Path(__file__).resolve().parents[2] +MODELS_SRC = (WORKSPACE / "studio/backend/models/inference.py").read_text() +ROUTES_SRC = (WORKSPACE / "studio/backend/routes/inference.py").read_text() +ADAPTER_SRC = ( + WORKSPACE / "studio/frontend/src/features/chat/api/chat-adapter.ts" +).read_text() +API_TYPES_SRC = ( + WORKSPACE / "studio/frontend/src/features/chat/types/api.ts" +).read_text() + + +def _find_class(tree: ast.AST, name: str) -> ast.ClassDef | None: + for node in ast.walk(tree): + if isinstance(node, ast.ClassDef) and node.name == name: + return node + return None + + +def test_chat_completion_request_has_cancel_id_field(): + tree = ast.parse(MODELS_SRC) + cls = _find_class(tree, "ChatCompletionRequest") + assert cls is not None + fields = { + n.target.id + for n in cls.body + if isinstance(n, ast.AnnAssign) and isinstance(n.target, ast.Name) + } + assert "cancel_id" in fields, ( + "ChatCompletionRequest must expose a cancel_id field for per-run " + "cancellation routing" + ) + + +def test_cancel_route_matches_cancel_id_exclusively_when_present(): + # A stale cancel POST carrying cancel_id AND session_id must not + # cancel a later run on the same thread via the shared session_id. + # Enforce this by requiring the handler to early-return through an + # exclusive-cancel_id path -- either an atomic helper or a keys + # list containing ONLY cancel_id (never session_id). + for node in ast.walk(ast.parse(ROUTES_SRC)): + if isinstance(node, ast.AsyncFunctionDef) and node.name == "cancel_inference": + break + else: + raise AssertionError("cancel_inference handler missing") + + cancel_id_exclusive_branch = False + for sub in ast.walk(node): + if not isinstance(sub, ast.If): + continue + test_src = ast.unparse(sub.test) + if "cancel_id" not in test_src or "isinstance" not in test_src: + continue + branch_src = "\n".join(ast.unparse(s) for s in sub.body) + before_return = branch_src.split("return", 1)[0] + matches_cancel_id_only = ( + "_cancel_by_cancel_id_or_stash(cancel_id)" in branch_src + or "_cancel_by_keys([cancel_id])" in branch_src + ) + if matches_cancel_id_only and "session_id" not in before_return: + cancel_id_exclusive_branch = True + break + assert cancel_id_exclusive_branch, ( + "cancel_inference must early-return with an exclusive cancel_id " + "match when a cancel_id is supplied, so a stale stop POST " + "cannot cancel a later run on the same thread via session_id" + ) + + +def test_cancel_route_falls_back_to_session_or_completion_when_no_cancel_id(): + for node in ast.walk(ast.parse(ROUTES_SRC)): + if isinstance(node, ast.AsyncFunctionDef) and node.name == "cancel_inference": + break + else: + raise AssertionError("cancel_inference handler missing") + + src = ast.unparse(node) + assert "session_id" in src and "completion_id" in src, ( + "cancel_inference must still accept session_id / completion_id as " + "fallback keys when cancel_id is absent" + ) + + +def test_frontend_request_type_has_cancel_id(): + assert re.search( + r"cancel_id\?\s*:\s*string\s*;", API_TYPES_SRC + ), "OpenAIChatCompletionsRequest must expose an optional cancel_id" + + +def test_chat_adapter_generates_cancel_id_per_run(): + m = re.search( + r"const\s+cancelId\s*=\s*([^;]+);", + ADAPTER_SRC, + ) + assert m, "chat-adapter.ts must declare a per-run `cancelId` constant" + rhs = m.group(1) + assert ( + "randomUUID" in rhs + ), "cancelId should prefer crypto.randomUUID() for uniqueness" + + +def test_chat_adapter_sends_cancel_id_in_completion_payload(): + assert "cancel_id: cancelId" in ADAPTER_SRC, ( + "chat-adapter.ts must include cancel_id in the streamChatCompletions " + "request payload so the backend registers under that key" + ) + + +def test_chat_adapter_sends_cancel_id_in_abort_cancel_post(): + m = re.search( + r"const\s+onAbortCancel\s*=\s*\(\)\s*=>\s*\{(.*?)\};", + ADAPTER_SRC, + flags = re.DOTALL, + ) + assert m, "onAbortCancel arrow function missing" + body = m.group(1) + assert re.search(r"cancel_id\s*:\s*cancelId", body), ( + "onAbortCancel must include cancel_id in the /inference/cancel body " + "so a stop POST matches the specific run, not the whole thread" + ) + + +def test_abort_cancel_post_uses_plain_fetch_with_manual_auth_header(): + # authFetch redirects to login on 401, which would kick the user to + # the login page mid-stop if the access token expired during a long + # stream. Use plain fetch + manual Authorization header for a + # best-effort cancel that never triggers the refresh/redirect flow. + start = ADAPTER_SRC.find("const onAbortCancel") + assert start >= 0, "onAbortCancel handler missing" + rest = ADAPTER_SRC[start:] + end = rest.find("\n try {") + body = rest if end < 0 else rest[:end] + assert "/api/inference/cancel" in body + assert "authFetch(" not in body, ( + "onAbortCancel must NOT call authFetch; a 401 from it would " + "redirect the user to the login page during a stop click" + ) + assert "fetch(" in body, "onAbortCancel must use plain fetch(...)" + assert "getAuthToken" in body, ( + "onAbortCancel must read the bearer token via getAuthToken() " + "rather than relying on authFetch's 401 flow" + ) + assert "Authorization" in body + assert ( + "keepalive: true" in body + ), "keepalive is required so the fetch survives page unload during stop" diff --git a/tests/studio/test_llama_cpp_wall_clock_cap.py b/tests/studio/test_llama_cpp_wall_clock_cap.py new file mode 100644 index 0000000000..671abea823 --- /dev/null +++ b/tests/studio/test_llama_cpp_wall_clock_cap.py @@ -0,0 +1,123 @@ +""" +Tests for the llama-server wall-clock cap (t_max_predict_ms). + +The UI always sends max_tokens = context_length, so gating +t_max_predict_ms on `max_tokens is None` makes the safety net dead +code. The fix applies the wall-clock cap unconditionally on all three +streaming payload sites and raises the default to 10 minutes so slow +CPU / macOS / Windows installs are not cut off mid-generation. + +Verifies: + - t_max_predict_ms is assigned unconditionally at the three + payload-builder sites (not inside an `if max_tokens is None` else + branch). + - _DEFAULT_T_MAX_PREDICT_MS is at least 10 minutes (previously + 120_000). + - The default max_tokens path still applies _DEFAULT_MAX_TOKENS. + - The three payload variable names (payload x2, stream_payload x1) + each get both `max_tokens` and `t_max_predict_ms`. +""" + +from __future__ import annotations + +import ast +from pathlib import Path + + +SOURCE_PATH = ( + Path(__file__).resolve().parents[2] + / "studio" + / "backend" + / "core" + / "inference" + / "llama_cpp.py" +) +SRC = SOURCE_PATH.read_text() +TREE = ast.parse(SRC) + + +def _is_subscript_assign(stmt: ast.stmt, target_name: str, key: str) -> bool: + if not isinstance(stmt, ast.Assign) or len(stmt.targets) != 1: + return False + t = stmt.targets[0] + if not isinstance(t, ast.Subscript): + return False + if not (isinstance(t.value, ast.Name) and t.value.id == target_name): + return False + slc = t.slice + return isinstance(slc, ast.Constant) and slc.value == key + + +def _collect_assignments(tree, target_name, key): + """Return list of (node, stack_of_enclosing_ifs) for each match.""" + hits = [] + + def visit(node, stack): + if _is_subscript_assign(node, target_name, key): + hits.append((node, stack)) + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.If): + for sub in child.body: + visit(sub, stack + [(child, "body")]) + for sub in child.orelse: + visit(sub, stack + [(child, "orelse")]) + else: + visit(child, stack) + + visit(tree, []) + return hits + + +def test_default_t_max_predict_ms_is_at_least_ten_minutes(): + for node in TREE.body: + if isinstance(node, ast.Assign) and len(node.targets) == 1: + t = node.targets[0] + if isinstance(t, ast.Name) and t.id == "_DEFAULT_T_MAX_PREDICT_MS": + value = node.value + assert isinstance(value, ast.Constant) + assert value.value >= 600_000, ( + f"_DEFAULT_T_MAX_PREDICT_MS must be >= 10 minutes " + f"(600_000 ms) to avoid cutting off slow-CPU generations; " + f"got {value.value}" + ) + return + raise AssertionError("_DEFAULT_T_MAX_PREDICT_MS constant missing") + + +def test_t_max_predict_ms_set_unconditionally_at_three_sites(): + hits_payload = _collect_assignments(TREE, "payload", "t_max_predict_ms") + hits_stream = _collect_assignments(TREE, "stream_payload", "t_max_predict_ms") + total = len(hits_payload) + len(hits_stream) + assert total == 3, ( + f"expected 3 total t_max_predict_ms assignments " + f"(payload x2 + stream_payload x1), got {total}" + ) + for node, stack in hits_payload + hits_stream: + for parent_if, branch in stack: + # The assignment must not be gated by a test that checks + # `max_tokens is None` (which would make it dead code for + # the UI path where max_tokens is always set). + test_src = ast.unparse(parent_if.test) + assert "max_tokens" not in test_src, ( + f"t_max_predict_ms at line {node.lineno} is nested under " + f"`if {test_src}:` -- it must be applied unconditionally so " + f"the wall-clock cap is not dead code for callers that set " + f"max_tokens" + ) + + +def test_max_tokens_default_cap_still_applied(): + # _DEFAULT_MAX_TOKENS must still kick in when caller passes None. + # We check the conditional expression `max_tokens if max_tokens is not + # None else _DEFAULT_MAX_TOKENS` appears at each site. + matches = 0 + for node in ast.walk(TREE): + if not isinstance(node, ast.IfExp): + continue + src = ast.unparse(node) + if "max_tokens" in src and "_DEFAULT_MAX_TOKENS" in src: + matches += 1 + assert matches >= 3, ( + f"expected >=3 `max_tokens if max_tokens is not None else " + f"_DEFAULT_MAX_TOKENS` expressions; got {matches}" + ) diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py new file mode 100644 index 0000000000..40ec3d6e1f --- /dev/null +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -0,0 +1,718 @@ +""" +Tests that the cancel tracker is registered BEFORE StreamingResponse is +returned, and that cleanup runs via a `finally` block inside each +async generator. + +The zombie-generation scenario is: user clicks Stop during prefill / +warmup / proxy buffering, before the first SSE chunk. If _tracker +__enter__ lives inside the async generator body, the registry is empty +at the moment /api/inference/cancel lands -- so cancel returns 0 and +the decode runs to completion. + +The fix moves _tracker = _TrackedCancel(...) and _tracker.__enter__() +to the synchronous body of openai_chat_completions (before the +StreamingResponse is returned) and places _tracker.__exit__ inside +each generator's `finally` block. Using a generator `finally` (rather +than a Starlette BackgroundTask) guarantees cleanup on every +termination path -- normal exhaustion, CancelledError from +ClientDisconnect, and OSError / BrokenPipeError during send() -- +because Starlette skips `background` callbacks when stream_response +raises. + +Structural verifies: + - No `async def ...:` body contains `_tracker.__enter__()` in + routes/inference.py (registration moved to sync body). + - Each of the four async generators (gguf_tool_stream, + gguf_stream_chunks, stream_chunks, audio_input_stream) contains + `_tracker.__exit__(None, None, None)` inside a try/finally block. + - No StreamingResponse in openai_chat_completions passes + `background=` (cleanup now lives in the generator finally). + +Behavioral verifies (extracting `_TrackedCancel` from source and +exercising the actual runtime semantics): + - `finally: _tracker.__exit__(...)` runs on normal completion, + mid-stream exception (OSError / BrokenPipeError from send()), + and aclose() from Starlette ClientDisconnect. + - A pre-set cancel_event (from `_TrackedCancel.__enter__` replaying + a pending cancel POST) lets the GGUF while-loop break cleanly + and emit final_chunk + [DONE] instead of propagating + `GeneratorExit` out of `_stream_with_retry` into the async + generator's `except Exception` (which would not catch it). +""" + +from __future__ import annotations + +import ast +import asyncio +import threading +import time +from pathlib import Path + + +SOURCE_PATH = ( + Path(__file__).resolve().parents[2] + / "studio" + / "backend" + / "routes" + / "inference.py" +) +SRC = SOURCE_PATH.read_text() +_TREE = ast.parse(SRC) + + +# ── Structural (AST) helpers ───────────────────────────────── + + +def _collect_async_functions(tree: ast.AST): + return [n for n in ast.walk(tree) if isinstance(n, ast.AsyncFunctionDef)] + + +def _has_tracker_enter_call(node: ast.AST) -> bool: + for sub in ast.walk(node): + if not isinstance(sub, ast.Call): + continue + fn = sub.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "__enter__" + and isinstance(fn.value, ast.Name) + and fn.value.id.startswith("_tracker") + ): + return True + return False + + +def _finalbody_has_tracker_exit(finalbody) -> bool: + for stmt in finalbody: + if not isinstance(stmt, ast.Expr): + continue + call = stmt.value + if not (isinstance(call, ast.Call) and isinstance(call.func, ast.Attribute)): + continue + fn = call.func + if ( + fn.attr == "__exit__" + and isinstance(fn.value, ast.Name) + and fn.value.id.startswith("_tracker") + ): + return True + return False + + +# ── Structural tests ───────────────────────────────────────── + + +def test_no_tracker_enter_inside_async_generators(): + offenders = [] + for fn in _collect_async_functions(_TREE): + if fn.name in { + "gguf_tool_stream", + "gguf_stream_chunks", + "stream_chunks", + "audio_input_stream", + }: + if _has_tracker_enter_call(fn): + offenders.append(fn.name) + assert not offenders, ( + f"Cancel tracker registration must live OUTSIDE the async generator " + f"body so a stop POST can find the registry entry before the first " + f"SSE chunk. Offending generators: {offenders}" + ) + + +def test_tracker_enter_exists_in_sync_body_of_chat_completions(): + top = None + for n in ast.walk(_TREE): + if isinstance(n, ast.AsyncFunctionDef) and n.name == "openai_chat_completions": + top = n + break + assert top is not None, "openai_chat_completions handler missing" + count = 0 + for sub in ast.walk(top): + if not isinstance(sub, ast.Call): + continue + fn = sub.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "__enter__" + and isinstance(fn.value, ast.Name) + and fn.value.id.startswith("_tracker") + ): + count += 1 + assert count >= 3, ( + f"expected >=3 _tracker.__enter__() calls in openai_chat_completions " + f"(one per streaming path), got {count}" + ) + + +def test_async_generators_cleanup_tracker_in_finally(): + required = { + "gguf_tool_stream", + "gguf_stream_chunks", + "stream_chunks", + "audio_input_stream", + } + found: set[str] = set() + for fn in [n for n in ast.walk(_TREE) if isinstance(n, ast.AsyncFunctionDef)]: + if fn.name not in required: + continue + for sub in ast.walk(fn): + if isinstance(sub, ast.Try) and sub.finalbody: + if _finalbody_has_tracker_exit(sub.finalbody): + found.add(fn.name) + break + missing = required - found + assert not missing, ( + f"Cleanup must run via `finally: _tracker.__exit__(None, None, None)` " + f"inside each streaming generator so ClientDisconnect / OSError paths " + f"also release registry entries (Starlette skips `background` callbacks " + f"when stream_response raises). Missing in: {sorted(missing)}" + ) + + +def test_streaming_responses_have_no_background_task(): + top = None + for n in ast.walk(_TREE): + if isinstance(n, ast.AsyncFunctionDef) and n.name == "openai_chat_completions": + top = n + break + assert top is not None + for sub in ast.walk(top): + if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): + continue + if sub.func.id != "StreamingResponse": + continue + kwargs = {kw.arg for kw in sub.keywords if kw.arg} + assert "background" not in kwargs, ( + "StreamingResponse in openai_chat_completions must not pass " + "`background=` -- cleanup now lives in the generator's finally " + "block; a BackgroundTask would be skipped on abrupt disconnect" + ) + + +# ── Behavioral helpers ─────────────────────────────────────── + +_WANTED = { + "_CANCEL_REGISTRY", + "_CANCEL_LOCK", + "_PENDING_CANCELS", + "_PENDING_CANCEL_TTL_S", + "_prune_pending", + "_TrackedCancel", + "_cancel_by_keys", + "_cancel_by_cancel_id_or_stash", +} + + +def _load_registry_module(): + chunks = [] + for n in _TREE.body: + seg = ast.get_source_segment(SRC, n) + if seg is None: + continue + if isinstance(n, (ast.FunctionDef, ast.ClassDef)) and n.name in _WANTED: + chunks.append(seg) + elif isinstance(n, ast.Assign): + names = [t.id for t in n.targets if isinstance(t, ast.Name)] + if any(name in _WANTED for name in names): + chunks.append(seg) + elif ( + isinstance(n, ast.AnnAssign) + and isinstance(n.target, ast.Name) + and n.target.id in _WANTED + ): + chunks.append(seg) + mod = {} + exec("import threading, time\n" + "\n\n".join(chunks), mod) + return mod + + +def _make_stream(tracker, raise_exc): + async def gen(): + try: + try: + yield "data: first\n\n" + if raise_exc is not None: + raise raise_exc + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + raise + except Exception: + yield "data: error\n\n" + finally: + tracker.__exit__(None, None, None) + except BaseException: + raise + + return gen() + + +async def _consume(agen): + out = [] + try: + async for ch in agen: + out.append(ch) + except BaseException as e: + out.append(type(e).__name__) + return out + + +def _llama_stub_raises_on_preset_cancel(cancel_event): + # Reproduces llama_cpp.py _stream_with_retry:2240 `raise GeneratorExit` + # when cancel_event is already set at entry. + if cancel_event.is_set(): + raise GeneratorExit + yield "cumulative-1" + yield "cumulative-2" + + +async def _post_fix_gguf_loop(cancel_event): + yield "first_chunk" + gen = _llama_stub_raises_on_preset_cancel(cancel_event) + sentinel = object() + while True: + if cancel_event.is_set(): + break + cumulative = await asyncio.to_thread(next, gen, sentinel) + if cumulative is sentinel: + break + yield cumulative + yield "final_chunk" + yield "[DONE]" + + +# ── Behavioral tests ───────────────────────────────────────── + + +def test_finally_cleanup_on_normal_completion(): + m = _load_registry_module() + m["_CANCEL_REGISTRY"].clear() + ev = threading.Event() + tr = m["_TrackedCancel"](ev, "cid-ok", "sid-ok") + tr.__enter__() + assert "cid-ok" in m["_CANCEL_REGISTRY"] + chunks = asyncio.run(_consume(_make_stream(tr, None))) + assert chunks == ["data: first\n\n", "data: [DONE]\n\n"] + assert "cid-ok" not in m["_CANCEL_REGISTRY"] + assert "sid-ok" not in m["_CANCEL_REGISTRY"] + + +def test_finally_cleanup_on_mid_stream_exception(): + # Simulates OSError / BrokenPipeError from Starlette send() mid-stream -- + # the exact case where pre-fix `background = BackgroundTask(...)` was + # skipped and leaked the registry entry. + m = _load_registry_module() + m["_CANCEL_REGISTRY"].clear() + ev = threading.Event() + tr = m["_TrackedCancel"](ev, "cid-err", "sid-err") + tr.__enter__() + assert "cid-err" in m["_CANCEL_REGISTRY"] + asyncio.run(_consume(_make_stream(tr, OSError("disconnect")))) + assert "cid-err" not in m["_CANCEL_REGISTRY"] + assert "sid-err" not in m["_CANCEL_REGISTRY"] + + +def test_finally_cleanup_on_aclose(): + # Starlette calls aclose() on the async generator when the client + # disconnects mid-stream. The generator's finally block must run. + m = _load_registry_module() + m["_CANCEL_REGISTRY"].clear() + ev = threading.Event() + tr = m["_TrackedCancel"](ev, "cid-abort", "sid-abort") + tr.__enter__() + assert "cid-abort" in m["_CANCEL_REGISTRY"] + + async def run(): + gen = _make_stream(tr, None) + it = gen.__aiter__() + await it.__anext__() + await gen.aclose() + + asyncio.run(run()) + assert "cid-abort" not in m["_CANCEL_REGISTRY"] + assert "sid-abort" not in m["_CANCEL_REGISTRY"] + + +def test_preset_cancel_event_exits_cleanly_with_done(): + # Pending-replay: POST /cancel arrived before the stream registered, + # was stashed, then consumed by _TrackedCancel.__enter__ which set + # cancel_event. The generator must break out of the loop cleanly + # and emit final_chunk + [DONE] rather than calling next(gen) and + # propagating `GeneratorExit` out of the GGUF stream wrapper. + ev = threading.Event() + ev.set() + chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev))) + assert "first_chunk" in chunks + assert "final_chunk" in chunks + assert "[DONE]" in chunks + assert "GeneratorExit" not in chunks + assert "cumulative-1" not in chunks + assert "cumulative-2" not in chunks + + +def test_normal_path_streams_all_tokens(): + # Regression: the top-of-loop cancel_event check must not short-circuit + # when cancel_event is unset. + ev = threading.Event() + chunks = asyncio.run(_consume(_post_fix_gguf_loop(ev))) + assert chunks == [ + "first_chunk", + "cumulative-1", + "cumulative-2", + "final_chunk", + "[DONE]", + ] + + +def test_cancel_during_streaming_stops_iteration_promptly(): + # Setting cancel_event between yields breaks out on the next iteration + # rather than draining the stub generator. + ev = threading.Event() + + async def _run(): + gen = _post_fix_gguf_loop(ev) + seen = [] + async for ch in gen: + seen.append(ch) + if ch == "cumulative-1": + ev.set() + return seen + + seen = asyncio.run(_run()) + assert "first_chunk" in seen + assert "cumulative-1" in seen + assert "cumulative-2" not in seen + assert "final_chunk" in seen + assert "[DONE]" in seen + + +# ── Cancel-event responsiveness in the streaming loops ─────── + + +def _loop_has_cancel_event_check(fn) -> bool: + # An `if cancel_event.is_set():` statement anywhere inside a + # `while`/`for` loop body is sufficient -- without it, a cancel POST + # cannot interrupt the loop because Colab-style proxies do not + # propagate request.is_disconnected(). + for sub in ast.walk(fn): + if not isinstance(sub, (ast.While, ast.For, ast.AsyncFor)): + continue + for stmt in ast.walk(sub): + if not isinstance(stmt, ast.If): + continue + t = stmt.test + if ( + isinstance(t, ast.Call) + and isinstance(t.func, ast.Attribute) + and t.func.attr == "is_set" + and isinstance(t.func.value, ast.Name) + and t.func.value.id == "cancel_event" + ): + return True + return False + + +def test_streaming_generators_check_cancel_event_in_loop(): + required = { + "gguf_tool_stream", + "gguf_stream_chunks", + "stream_chunks", + "audio_input_stream", + } + missing = [] + for fn in [n for n in ast.walk(_TREE) if isinstance(n, ast.AsyncFunctionDef)]: + if fn.name not in required: + continue + if not _loop_has_cancel_event_check(fn): + missing.append(fn.name) + assert not missing, ( + f"Each streaming generator must check `cancel_event.is_set()` inside " + f"its main loop so `POST /api/inference/cancel` can interrupt the " + f"stream through proxies that do not forward fetch aborts. " + f"Missing in: {sorted(missing)}" + ) + + +def test_audio_input_stream_offloads_blocking_next_to_thread(): + # Guards against regression back to `for chunk_text in + # audio_input_generate():` -- which blocks the event loop on each + # whisper chunk and prevents POST /api/inference/cancel from being + # serviced until the chunk yields. + audio = None + for fn in ast.walk(_TREE): + if isinstance(fn, ast.AsyncFunctionDef) and fn.name == "audio_input_stream": + audio = fn + break + assert audio is not None, "audio_input_stream generator missing" + + for sub in ast.walk(audio): + if isinstance(sub, (ast.For, ast.AsyncFor)): + it_src = ast.unparse(sub.iter) + assert "audio_input_generate" not in it_src, ( + "audio_input_stream must not iterate audio_input_generate() " + "directly -- that blocks the event loop. Use " + "`await asyncio.to_thread(next, gen, _DONE)` inside a " + "`while True` loop instead" + ) + + found_to_thread_next = False + for sub in ast.walk(audio): + if not isinstance(sub, ast.Call): + continue + fn_expr = sub.func + if not ( + isinstance(fn_expr, ast.Attribute) + and fn_expr.attr == "to_thread" + and isinstance(fn_expr.value, ast.Name) + and fn_expr.value.id == "asyncio" + ): + continue + if sub.args and isinstance(sub.args[0], ast.Name) and sub.args[0].id == "next": + found_to_thread_next = True + break + assert found_to_thread_next, ( + "audio_input_stream must call `asyncio.to_thread(next, gen, ...)` " + "to keep the event loop free while whisper yields the next chunk" + ) + + +def test_stream_chunks_cancel_branch_resets_backend_state(): + # The Unsloth path's cancel branch must flush GPU / KV-cache state + # via `backend.reset_generation_state()` -- the orchestrator's + # internal cancel path does not do this, so a cancel-via-POST that + # only broke the loop would leave the subprocess in a dirty state + # for the next request. + fn = None + top = None + for n in ast.walk(_TREE): + if isinstance(n, ast.AsyncFunctionDef) and n.name == "openai_chat_completions": + top = n + break + assert top is not None + for n in ast.walk(top): + if isinstance(n, ast.AsyncFunctionDef) and n.name == "stream_chunks": + fn = n + break + assert fn is not None, "stream_chunks generator missing" + + for sub in ast.walk(fn): + if not isinstance(sub, ast.If): + continue + t = sub.test + if not ( + isinstance(t, ast.Call) + and isinstance(t.func, ast.Attribute) + and t.func.attr == "is_set" + and isinstance(t.func.value, ast.Name) + and t.func.value.id == "cancel_event" + ): + continue + body_src = "\n".join(ast.unparse(s) for s in sub.body) + if "backend.reset_generation_state()" in body_src: + return + raise AssertionError( + "stream_chunks `if cancel_event.is_set():` branch must call " + "backend.reset_generation_state() -- matches the existing " + "request.is_disconnected() / CancelledError cleanup paths and " + "prevents KV-cache drift after cancel-via-POST" + ) + + +# ── Behavioral simulations for the iter-1 fixes ────────────── + + +def test_unsloth_stream_loop_breaks_on_external_cancel_event(): + cancel_event = threading.Event() + reset_calls = [0] + + class _Backend: + def reset_generation_state(self): + reset_calls[0] += 1 + + backend = _Backend() + + def _generate(): + for i in range(200): + time.sleep(0.005) + yield f"cum-{i}" + + async def _loop(): + _DONE = object() + loop = asyncio.get_event_loop() + gen = _generate() + seen = [] + while True: + if cancel_event.is_set(): + backend.reset_generation_state() + break + cumulative = await loop.run_in_executor(None, next, gen, _DONE) + if cumulative is _DONE: + break + seen.append(cumulative) + return seen + + async def _fire(): + await asyncio.sleep(0.05) + cancel_event.set() + + async def _main(): + return await asyncio.gather(_loop(), _fire()) + + seen, _ = asyncio.run(_main()) + assert ( + len(seen) < 200 + ), f"loop must not drain the generator after cancel; got {len(seen)} tokens" + assert reset_calls[0] == 1, ( + f"backend.reset_generation_state() must be called exactly once on " + f"cancel-via-POST, got {reset_calls[0]}" + ) + + +def test_audio_stream_stays_responsive_under_blocking_next(): + # Regression guard: replace the post-fix loop with the pre-fix + # `for chunk in audio_input_generate()` pattern and assert it blocks + # the event loop; then confirm the post-fix pattern exits promptly. + cancel_event = threading.Event() + + def _audio_gen(): + for i in range(8): + time.sleep(0.15) + yield f"chunk-{i}" + + async def _prefix_loop(): + seen = [] + for chunk_text in _audio_gen(): + if cancel_event.is_set(): + break + seen.append(chunk_text) + return seen + + async def _postfix_loop(): + _DONE = object() + gen = _audio_gen() + seen = [] + while True: + if cancel_event.is_set(): + break + chunk_text = await asyncio.to_thread(next, gen, _DONE) + if chunk_text is _DONE: + break + seen.append(chunk_text) + return seen + + async def _fire_early(): + await asyncio.sleep(0.05) + cancel_event.set() + + async def _run(loop_coro): + return await asyncio.gather(loop_coro, _fire_early()) + + cancel_event.clear() + t0 = time.monotonic() + prefix_seen, _ = asyncio.run(_run(_prefix_loop())) + prefix_elapsed = time.monotonic() - t0 + assert prefix_elapsed >= 0.13, ( + f"pre-fix pattern should block event loop for >=1 chunk time " + f"(~150ms); got {prefix_elapsed:.3f}s, {len(prefix_seen)} chunks" + ) + + cancel_event.clear() + t0 = time.monotonic() + postfix_seen, _ = asyncio.run(_run(_postfix_loop())) + postfix_elapsed = time.monotonic() - t0 + assert postfix_elapsed < prefix_elapsed, ( + f"post-fix pattern must exit faster than pre-fix (blocking) " + f"pattern; post={postfix_elapsed:.3f}s vs pre={prefix_elapsed:.3f}s" + ) + assert ( + len(postfix_seen) < 8 + ), f"post-fix loop must not drain all chunks; got {len(postfix_seen)}" + + +def test_unsloth_stream_loop_emits_zero_tokens_on_preset_cancel(): + # Pending-cancel replay: _TrackedCancel.__enter__ already set + # cancel_event before the generator body starts iterating. The + # top-of-loop check must short-circuit the very first iteration so + # no token is emitted. Catches a regression that moves the check + # below `next()` -- the mid-loop test would still pass but this + # test would observe one extra token leak. + cancel_event = threading.Event() + cancel_event.set() + reset_calls = [0] + + class _Backend: + def reset_generation_state(self): + reset_calls[0] += 1 + + backend = _Backend() + + next_calls = [0] + + def _generate(): + while True: + next_calls[0] += 1 + yield f"cum-{next_calls[0]}" + + async def _loop(): + _DONE = object() + loop = asyncio.get_event_loop() + gen = _generate() + seen = [] + while True: + if cancel_event.is_set(): + backend.reset_generation_state() + break + cumulative = await loop.run_in_executor(None, next, gen, _DONE) + if cumulative is _DONE: + break + seen.append(cumulative) + return seen + + seen = asyncio.run(_loop()) + assert seen == [], ( + f"loop must emit zero tokens when cancel_event is pre-set " + f"(pending-replay path); got {seen}" + ) + assert next_calls[0] == 0, ( + f"loop must not call next() at all on pre-set cancel; got " + f"{next_calls[0]} calls" + ) + assert reset_calls[0] == 1, ( + f"backend.reset_generation_state() must still fire exactly once " + f"on pre-set cancel; got {reset_calls[0]}" + ) + + +def test_audio_stream_emits_zero_chunks_on_preset_cancel(): + # Symmetric to the Unsloth pre-set test: the audio loop's top-of-loop + # cancel check must skip the asyncio.to_thread(next, ...) call when + # cancel_event was already set via pending-replay. + cancel_event = threading.Event() + cancel_event.set() + + next_calls = [0] + + def _audio_gen(): + while True: + next_calls[0] += 1 + yield f"chunk-{next_calls[0]}" + + async def _loop(): + _DONE = object() + gen = _audio_gen() + seen = [] + while True: + if cancel_event.is_set(): + break + chunk_text = await asyncio.to_thread(next, gen, _DONE) + if chunk_text is _DONE: + break + seen.append(chunk_text) + return seen + + seen = asyncio.run(_loop()) + assert seen == [], f"audio loop must emit zero chunks on pre-set cancel; got {seen}" + assert next_calls[0] == 0, ( + f"audio loop must not call next() on pre-set cancel; got " + f"{next_calls[0]} calls" + ) From b09aa82a3ac9ac9794c497e4b8f747f77e52b162 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Apr 2026 12:02:03 -0700 Subject: [PATCH 09/54] Studio: add github_repo seed reader and GitHub Support Bot recipe (#5169) * Studio: add github_repo seed reader and GitHub Support Bot recipe Adds a first-party Data Designer seed reader that scrapes GitHub issues, pull requests, and commits from one or more repositories via the GraphQL API, and a learning recipe (GitHub Support Bot) that turns those rows into synthetic support Q&A pairs for fine-tuning. Backend (new plugin studio/backend/plugins/data-designer-github-repo-seed): * GitHubRepoSeedSource config: repos, token (falls back to GH_TOKEN / GITHUB_TOKEN env var), item_types (issues / pulls / commits), per-resource limit (0 means all), max_comments_per_item. * Rate-limit-aware GraphQL client (GitHubClient + RepoScraper) shared across repos; flattens each item into a uniform row with columns item_type, repo, number, title, body, state, author, created_at, closed_at, url, labels, comments. * Registered via the data_designer.plugins entry point. Frontend: * New seed_github block variant so the seed node card shows "GitHub repositories" instead of the generic "Document file" placeholder, with its own icon and inline summary (repo count + item-type list). * Rewritten seed dialog github_repo form: repos textarea pre-filled with unslothai/unsloth + unslothai/unsloth-zoo, password input for the GH token, items-per-repo number with an "All" toggle, and the noisier options (item types, max comments, include comments) tucked under an Advanced collapsible. * Local model auto-load on Run: if a recipe uses an is_local provider and the inference server is not already serving that model, the executions hook calls /api/inference/load first. Removes the "open /chat to load a model" prerequisite that users kept tripping on. * Honor the recipe's run.rows value in the Run dialog (previously the store reset to 5 regardless of what the template shipped). Recipe (studio/frontend/src/features/data-recipes/learning-recipes/ github-support-bot.json): * Defaults to the Local Model provider + unsloth/gemma-4-E2B-it-GGUF. * Scrapes unslothai/unsloth and unslothai/unsloth-zoo, issues and pulls, up to 100 items per resource. * Two LLM blocks: normalized_question (llm-text) rewrites each thread into a clean support question, support_answer (llm-structured) produces JSON with answer / diagnosis_questions / cites / confidence. * Run defaults to 10 rows for a quick smoke test. Verified end-to-end on a running Studio: card renders, source-data dialog is pre-populated, All toggle disables the limit input, the recipe executes and produces rows against a loaded local GGUF. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: improve GitHub recipe support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: speed up GitHub scraper and harden the support-bot recipe Addresses a perf issue found while demoing the github_repo seed reader: Scraper is too slow at scale. The PRs GraphQL query pulls deeply nested fields (reviewThreads, reviews, commits, timelineItems, etc.) so the page size was pinned at 3 to stay under GitHub's node-count ceiling. 100 PRs meant 34 serial round trips. Added lighter query variants (PRS_PAGE_QUERY_LIGHT, ISSUES_PAGE_QUERY_LIGHT) that drop the fields the Studio flatten layer does not use (it only reads title, body, state, author, labels, comments). With the light query PR pages can safely go to 25 per page and issues to 50. The plugin scraper now passes light=True to RepoScraper so Studio always uses the fast path; the heavy query remains available for other callers. Recipe defaults are now demo-ready with production knobs called out: - max_parallel_requests: 1 and max_tokens: 800 so small local models stay stable when running the support_answer structured column. - support_answer prompt trimmed to 80-200 words so gemma-4-E2B GGUF can actually comply with the schema. The canonical 150-300 word codex prompt is still documented in the node3 markdown note for production upgrades. * Studio: rename GitHub recipe to 'GitHub Scraper' and add Easy mode Changes the recipe framing from a single-purpose 'Support Bot' pipeline to a general-purpose scraper that produces {user_request, grounded_response} training pairs. Aligns with the canonical github_data_gatherer dataset (11 enrichment tasks mirrored in pr_requests_20 / issue_requests_20 on the input side and explain_pr / issue_fix_plan / issue_solution on the output side). Recipe JSON changes: - columns[0] renamed normalized_question -> user_request, prompt now inverts a GitHub thread into a realistic user ask instead of normalising it. - columns[1] renamed support_answer -> coauthor_response, emits {response, followups, cites, task, confidence} and branches on issue vs PR thread type. - Notes rewritten to document the 11-task catalog and the canonical production prompt to paste in for a full dataset backfill. Frontend: Easy mode for github_repo recipes. The drag-and-drop canvas is hidden behind an 'Advanced' tab; Easy mode is the default for any recipe whose seed_source_type is github_repo. The Easy form reuses the existing GithubRepoSeedForm (promoted to exported), adds a rows input bound to previewRows, a model field bound to the model_config, and a single Run button that calls runPreview() directly (no modal). Non-github recipes see the same Editor / Runs tabs as before. View mode persists per-recipe-id in localStorage under recipe-studio:view-mode:. * Studio: auto-detect server GH_TOKEN and widen Easy-mode detection The GitHub seed form now fetches /api/data-recipe/seed/github/env-token on mount and, when the server exposes a GH_TOKEN / GITHUB_TOKEN env var and the token field is blank, shows a small 'Using server env var' badge and swaps the placeholder text. The token value itself is never returned to the UI. Widens Easy-mode detection in recipe-studio-page.tsx so that recipes saved before ui.seed_source_type was persisted also get the Easy tab: falls back to recipe.seed_config.source.seed_type, which is always present for github_repo seeds. * fix: polish GitHub recipe UI * Studio: default llama-server --threads to -1 (auto) Previously we passed --threads only when the caller set an explicit value, which meant llama-server fell back to its internal default. That default has varied across llama.cpp builds (some versions use hardware concurrency including hyperthreads, which hurts throughput on CPU-heavy inference). Always passing --threads -1 pins the behaviour to llama.cpp's auto-detect (physical cores). Caller-supplied n_threads still wins when non-None. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: auto-switch Easy mode to Runs pane on run start Easy mode had no progress island or canvas overlay, so after clicking Run the only visible state was the button label flipping to "Running..." while the screen otherwise stayed identical. This reads as stuck even though the job is progressing. Wire an onExecutionStart callback from recipe-studio-page.tsx through to useRecipeExecutions so that when a run is kicked off from easy mode, the page flips to the executions view where the Runs sidebar, progress bar, rate/ETA panel, and live log are rendered. Advanced/editor mode keeps its existing behavior and stays on the canvas (it already has the floating ExecutionProgressIsland). * fix: clean up GitHub scraper layout * Studio: forward llm-structured output_format as llama-server response_format Local GGUF runs of llm-structured columns used to generate the full max_tokens budget before the prompt-level "return JSON in a ```json fence" instruction got parsed. Small models (e.g. gemma-4-E2B-it) routinely broke format, so each row took ~65s and frequently failed with "No parsable JSON structure within ```json markdown fence". For any local-provider model_config referenced by an llm-structured column, clone the model_config and inject response_format into the clone's inference_parameters. Uses llama.cpp server's flat shape (tools/server/README.md): {"type": "json_schema", "schema": } Not the OpenAI-nested form; data_designer's OpenAI adapter forwards response_format verbatim via facade._COMPLETION_REQUEST_FIELDS, and llama-server's documented schema path expects the flat variant. The clone is per (model_alias, column) so: - llm-text / llm-judge columns that share the same alias keep free-form sampling. - Each structured column gets its own schema, so columns with different output_formats don't collide. Effect on gemma-4-E2B-it demos: every row parses cleanly, and the model terminates immediately after the closing brace instead of running to max_tokens. Net wall-clock is usually faster even though grammar-constrained sampling is slightly slower per token. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: flip Easy to Runs pane before validation scrape, not after Previously onExecutionStart fired inside runExecution, which runs AFTER validateRecipe() -- and validation re-invokes the seed reader. For the github_repo reader that is a full GraphQL scrape, so the user sat on a "Running..." button with an otherwise unchanged Easy form for 10-15s before anything moved. Call onExecutionStart at the top of runWithValidation, right after we have a payload to send. The view flips immediately; ensureLocalModelLoaded + validateRecipe now run against the Runs pane instead of a frozen Easy form. runExecution still calls onExecutionStart downstream, but the callback is idempotent (the page's easy -> executions guard skips the second call), so no behaviour change for runs that pass validation. If validation fails the toast + runErrors path still fires; the Easy form's error banner still reads runErrors when the user switches back. * Studio: unify data-recipe workflow auth on sk-unsloth-* keys The previous commit (a61b4cc9) assumed storage.create_api_key(..., internal=True) and storage.revoke_internal_api_key(key_id) existed, but those helpers were only in the working tree, never committed. Recipe runs in local-model mode were therefore crashing with 500 when _inject_local_providers tried to mint a workflow key. This commit ships the missing pieces. auth/storage.py: - api_keys schema gains is_internal INTEGER DEFAULT 0 (with a guarded ALTER TABLE migration so existing auth.db files upgrade in place). - create_api_key takes an internal=False kwarg; internal keys are flagged so they can be hidden from user-facing listings. - list_api_keys takes include_internal=False so UIs never see workflow keys. - New revoke_internal_api_key(key_id): id-only revoke for keys minted by non-user subjects (the JobManager does not know a username). core/data_recipe/jobs/manager.py: - JobManager.start accepts internal_api_key_id and stores it on Job so lifecycle handlers can revoke eagerly. - _handle_event revokes on EVENT_JOB_COMPLETED / _ERROR / _CANCELLED. - _pump_loop subprocess-died fallback also retires the key so a crashed worker cannot leak a live sk-unsloth-* beyond its TTL. - Revocation is best-effort (swallow exceptions) -- the 24h TTL is the safety net if storage hiccups. core/data_recipe/jobs/types.py: - Job dataclass gains internal_api_key_id: int | None = None. Replaces the bespoke 24h JWT path that jobs.py used to mint for local providers. One mint/revoke/verify surface for every API key the server issues, and revocation is now eager (seconds, not 24h) instead of TTL-only. * Studio: plug workflow-key leak on unexpected create_job errors Review follow-up on the sk-unsloth-* workflow-key lifecycle in create_job. Previously the revoke handlers wrapped mgr.start(...) but only caught RuntimeError and ValueError, and get_job_manager() sat outside the try block entirely. Any other exception type (TypeError from a mismatched kwarg, OSError from the queue write, etc.) would bubble up to FastAPI and leave the minted key live until its 24h TTL. Fix: one try block covers both get_job_manager() and mgr.start(), with a trailing except Exception that revokes and re-raises. The RuntimeError -> 409 and ValueError -> 400 paths are unchanged so specific client-facing status codes still surface. Revocation is still best-effort (_revoke_internal_api_key_safe swallows errors) because we never want revoke failures to mask the original crash. Severity is low -- the key can't bootstrap longer access and the 24h TTL bounds the window -- but the reviewer's point stands: eager revoke on every failure path is the right invariant. * Studio: nest response_format under extra_body so pydantic accepts it The previous commit dropped response_format at the top level of a cloned model_config's inference_parameters, which BuilderConfig rejected with: ValidationError: Extra inputs are not permitted [type=extra_forbidden] data_designer.model_configs.1.inference_parameters.response_format data_designer's BaseInferenceParams is a pydantic model with extra=forbid and only a fixed set of fields (temperature, top_p, max_tokens, max_parallel_requests, timeout, extra_body). The pass-through path for anything the schema doesn't know about is `extra_body`, which the OpenAI SDK spreads into the chat-completions request body at the top level -- which is exactly where llama-server reads response_format from. Inject under extra_body (merging with any existing extra_body contents) so the clone validates. llama-server still receives {"type": "json_schema", "schema": } at the top level of the request body, which is the flat shape llama.cpp's server expects. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: forward response_format to llama-server and fence-wrap the reply Two-part fix for the llm-structured data-recipe path: (1) The /v1/chat/completions proxy was dropping response_format. The route's passthrough branch only triggered on tools / tool messages, so requests carrying a JSON schema fell into the non-passthrough GGUF path which calls generate_chat_completion (no response_format kwarg). The schema never reached llama-server, so guided decoding was a no-op and the model emitted free-form text that happened to parse a fraction of the time. Widen the passthrough trigger and teach _build_passthrough_payload to forward response_format so llama-server's GBNF grammar actually runs. Guided decoding does not require supports_tools, so split the condition: a request is now passthrough-routed if it carries tools/tool messages (existing behavior) OR carries response_format (new). The vision guard, streaming fork, and tools-choice defaulting are unchanged. (2) data_designer's llm-structured parser looks for a ```json ... ``` markdown fence and discards anything else. Guided decoding emits only the JSON object (the GBNF grammar has no fence tokens), so a 100%-valid schema-constrained run still ended up 0 ok / N failed with "No parsable JSON structure within ```json markdown fence". In _openai_passthrough_non_streaming, wrap each choice's content in the expected fence when the caller asked for guided decoding. Already-fenced content is left alone so other clients that prefer raw JSON are not affected; the wrap is scoped to requests that carried response_format. Net effect on the GitHub Support Bot recipe on a local GGUF: schema actually binds during sampling, content arrives wrapped in the fence data_designer expects, and generation terminates immediately after the closing brace instead of running out to max_tokens. * Studio: Easy mode runs a full run, capped at the user's row count Easy mode used to call runPreview, which produces a test run: no artifact persisted, reduced progress tracking, and framed in the Runs pane as "Test run". The whole point of the form is to let a user kick off a real dataset build with one click, so wire it to runFull instead and bind the Rows input to fullRows (not previewRows). runFull requires a non-empty fullRunName. The Easy form has no run-name input, so seed a default on mount whenever Easy is active and fullRunName is still empty. Uses ` ` so each Easy run gets a stable-ish default that still sorts chronologically in the Runs pane. User can override it from the Advanced run dialog before clicking Run. Rename GithubScraperEasyView's rows props from previewRows/setPreviewRows to rows/setRows so the view stays agnostic to which hook state the page chooses to bind. Loading indicator now follows fullLoading. * Studio: clamp GitHub scrape page size and memoize the materialization Two wins for the "before Generating fires" gap on small previews: (1) scrape_{issues,prs,commits} hardcoded per_page (50 / 25 / 100) and only checked the trial limit AFTER the page was written, so a 1-row Easy run still asked GitHub for a full 50-issue + 25-PR page, wrote them all to JSONL, and then stopped because total_new already exceeded the trial cap. Cap per_page at min(page_cap, trial_limit) so github_limit=1 actually asks for first:1. (2) GitHubRepoSeedReader.get_dataset_uri used to scrape fresh on every invocation. data_designer calls the seed reader multiple times per recipe job (validation, preview, per-column sampling), so a 2-repo Easy preview ran the full GraphQL scrape three times back-to-back, burning ~15s of dead air before any LLM generation began. Added a module-level in-process cache keyed on (repos, item_types, limit, include_comments, max_comments_per_item, sha256(token)[:16]) that stores the JSONL path of the first materialization. Subsequent calls with the same signature return the cached path, guarded by a staleness check that drops the entry if the file was tmp-cleaned. Raw token values never land in the key. Net effect on a 1-row Easy run, 2 repos, limit=1: 2 GraphQL round trips instead of ~12, and the first-to-Generating gap collapses from ~15s to roughly 2-3s. * Studio: make Easy mode Rows input editable instead of snapping to 1 The Rows to generate input used type="number" with value bound directly to the rows state and an onChange that coerced any non-positive parse result back to 1. The moment the user pressed backspace to clear the field, the parent re-rendered with value=1 and the caret jumped, making it impossible to change the value without arrowing the browser's +/- spinner. Switch to a text input with inputMode="numeric" and pattern="[0-9]*" (so mobile still shows a numeric keyboard, and the browser drops the spinner buttons the user did not want). Add a local rowsText buffer so the field can hold transient empty / partial digit strings while editing without fighting the parent state; the canonical rows value only advances when the buffer parses to a valid integer in [1, 10000], and onBlur clamps back to 1 or 10000 if the user left it out of range. No behavior change for valid numeric edits - the downstream runFull() still sees a clean positive integer. * Studio: expand dataset cells horizontally by column on click Click a long cell to expand that whole column. Click again to collapse. Replaces the prior row-level vertical expansion which made it hard to compare cells across columns. State is scoped per execution and per column; the row itself is no longer a click target. * Studio: force expanded dataset column to grow wide enough to read * Studio: disable thinking for local recipe inference and plumb the kwarg Reasoning-capable models (gemma-3n, qwen3.5, etc.) emit a ... preamble ahead of the answer by default, which roughly doubles the generated token count per row on a local GGUF and pushes the actual answer past data_designer's json-fence regex on llm-structured columns. Recipes want the terse answer, not the scratchpad. Two halves of the fix: (1) routes/data_recipe/jobs.py: when _inject_local_providers walks the recipe's model_configs to point them at the local endpoint, also stash chat_template_kwargs={"enable_thinking": false} under each config's inference_parameters.extra_body. OpenAI SDK spreads extra_body into the top-level request body, so llama-server and the Studio /v1/chat/completions route both see it. (2) routes/inference.py: the chat-completions route previously dropped chat_template_kwargs on the floor because the whitelist body builder only forwarded known fields. - At the top of openai_chat_completions, lift chat_template_kwargs.enable_thinking from payload.model_extra onto the typed payload.enable_thinking field when the caller did not set the latter, so the non-passthrough GGUF path's generate_chat_completion(...) call honors the override. - Teach _build_passthrough_payload to forward a chat_template_kwargs dict, and have _build_openai_passthrough_body derive that dict from payload.enable_thinking so response_format requests (structured columns) also land at llama-server with the reasoning preamble suppressed. Net effect on a 10-row support-bot run with gemma-4-E2B-it-GGUF: responses arrive without tags, wall-clock per call drops roughly in half, and structured columns stop leaking reasoning tokens through the GBNF-constrained output. * Studio: update GitHub Support Bot learning recipe with maintainer layout Replace the template with the hand-laid-out export from the maintainer so note nodes ship with real x/y positions (scattered around the graph instead of all stacked at x=480) and the edges / canvas pan look correct on first load. Also picks up the maintainer's prompt tweaks and output schema names (coauthor_response / user_request / followups / task / cites / confidence). Diff is mostly ui.nodes positions and prompt bodies; runtime shape is unchanged (seed_config / columns still target model_1 against the Local Model provider). * Studio: auto-size dataset sample columns; wide text gets a wide column Drop the per-column click-to-expand toggle and the 180-char truncation. Every column now renders its full value. Columns with long text get a min-w of 48rem so the text is readable without wrapping into a tall block; narrow-content columns get a 12rem min-w. The table wrapper already has overflow-x-auto, so wide-column totals cause a horizontal scrollbar instead of cramming everything into the viewport. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix GitHub scrape progress * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add resetApiBase export for test setup * Studio: rename github-support-bot output columns to User / Assistant Previously emitted user_request and coauthor_response, which did not match the canonical User / Assistant chat-pair shape that downstream SFT consumers expect. Renamed the columns in the recipe JSON (columns, UI node ids, edges, notes, prompt Jinja refs) and the matching copy in the learning-recipes index, data-recipes-page, and easy view. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: wasimysaid --- pyproject.toml | 1 + studio/backend/auth/storage.py | 84 +- .../core/data_recipe/jobs/constants.py | 1 + .../backend/core/data_recipe/jobs/manager.py | 105 ++- studio/backend/core/data_recipe/jobs/parse.py | 224 ++++- studio/backend/core/data_recipe/jobs/types.py | 25 + .../backend/core/data_recipe/jobs/worker.py | 29 +- studio/backend/core/inference/llama_cpp.py | 6 +- .../data-designer-github-repo-seed/README.md | 73 ++ .../pyproject.toml | 25 + .../__init__.py | 7 + .../data_designer_github_repo_seed/config.py | 64 ++ .../data_designer_github_repo_seed/impl.py | 83 ++ .../data_designer_github_repo_seed/plugin.py | 10 + .../data_designer_github_repo_seed/scraper.py | 236 +++++ .../scraper_impl/__init__.py | 2 + .../scraper_impl/gh_client.py | 248 ++++++ .../scraper_impl/queries.py | 685 +++++++++++++++ .../scraper_impl/scraper.py | 756 ++++++++++++++++ .../scraper_impl/state_store.py | 105 +++ .../single-env/data-designer-deps.txt | 3 +- studio/backend/routes/data_recipe/jobs.py | 199 ++++- studio/backend/routes/data_recipe/seed.py | 12 + studio/backend/routes/data_recipe/validate.py | 67 ++ studio/backend/routes/inference.py | 99 ++- .../tests/test_data_recipe_github_progress.py | 91 ++ .../learning-recipes/github-support-bot.json | 238 +++++ .../data-recipes/learning-recipes/index.ts | 11 + .../data-recipes/pages/data-recipes-page.tsx | 17 + .../src/features/recipe-studio/api/index.ts | 70 +- .../recipe-studio/blocks/definitions.ts | 18 +- .../executions/execution-data-tab.tsx | 178 +++- .../executions/execution-overview-tab.tsx | 66 ++ .../components/executions/executions-view.tsx | 87 +- .../components/inline/inline-seed.tsx | 83 +- .../components/recipe-studio-header.tsx | 11 +- .../runtime/execution-progress-island.tsx | 98 ++- .../recipe-studio/dialogs/preview-dialog.tsx | 9 +- .../dialogs/seed/seed-dialog.tsx | 810 ++++++++++++++---- .../easy/github-crawler-easy-view.tsx | 191 +++++ .../features/recipe-studio/execution-types.ts | 26 +- .../recipe-studio/executions/runtime.ts | 11 + .../hooks/use-recipe-executions.ts | 117 ++- .../hooks/use-recipe-studio-actions.ts | 4 + .../recipe-studio/recipe-studio-page.tsx | 73 +- .../recipe-studio/stores/recipe-studio.ts | 10 +- .../src/features/recipe-studio/types/index.ts | 19 +- .../import/parsers/seed-config-parser.ts | 32 + .../features/recipe-studio/utils/node-data.ts | 4 +- .../utils/payload/builders-seed.ts | 108 ++- .../recipe-studio/utils/payload/types.ts | 2 +- .../recipe-studio/utils/validation.ts | 51 +- studio/frontend/src/lib/api-base.ts | 6 + studio/install_python_stack.py | 39 +- 54 files changed, 5210 insertions(+), 419 deletions(-) create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/README.md create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py create mode 100644 studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py create mode 100644 studio/backend/tests/test_data_recipe_github_progress.py create mode 100644 studio/frontend/src/features/data-recipes/learning-recipes/github-support-bot.json create mode 100644 studio/frontend/src/features/recipe-studio/easy/github-crawler-easy-view.tsx diff --git a/pyproject.toml b/pyproject.toml index 815c6ee119..5687ea12f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ studio = [ "frontend/*.yaml", "frontend/.git*", "backend/requirements/**/*", + "backend/plugins/**/*", "backend/core/data_recipe/oxc-validator/*.json", "backend/core/data_recipe/oxc-validator/*.mjs", ] diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 2b0e359d39..9a03f5f542 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -146,10 +146,18 @@ def get_connection() -> sqlite3.Connection: created_at TEXT NOT NULL, last_used_at TEXT, expires_at TEXT, - is_active INTEGER NOT NULL DEFAULT 1 + is_active INTEGER NOT NULL DEFAULT 1, + is_internal INTEGER NOT NULL DEFAULT 0 ); """ ) + api_key_columns = { + row["name"] for row in conn.execute("PRAGMA table_info(api_keys)") + } + if "is_internal" not in api_key_columns: + conn.execute( + "ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS app_secrets ( @@ -592,11 +600,15 @@ def create_api_key( username: str, name: str, expires_at: Optional[str] = None, + internal: bool = False, ) -> Tuple[str, dict]: """Create a new API key for *username*. Returns ``(raw_key, row_dict)`` where *raw_key* is shown to the user - exactly once. The database only stores the SHA-256 hash. + exactly once. The database only stores the PBKDF2 hash. + + Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe + runs) that should not appear in user-facing key listings. """ raw_key = API_KEY_PREFIX + secrets.token_hex(16) key_hash = _pbkdf2_api_key(raw_key) @@ -607,10 +619,18 @@ def create_api_key( try: conn.execute( """ - INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal) + VALUES (?, ?, ?, ?, ?, ?, ?) """, - (username, key_prefix, key_hash, name, now, expires_at), + ( + username, + key_prefix, + key_hash, + name, + now, + expires_at, + 1 if internal else 0, + ), ) conn.commit() cur = conn.execute("SELECT * FROM api_keys WHERE key_hash = ?", (key_hash,)) @@ -620,19 +640,33 @@ def create_api_key( conn.close() -def list_api_keys(username: str) -> list: - """Return all API keys for *username* (never exposes ``key_hash``).""" +def list_api_keys(username: str, include_internal: bool = False) -> list: + """Return API keys for *username*. Internal workflow keys are hidden + by default so they do not clutter user-facing UIs.""" conn = get_connection() try: - cur = conn.execute( - """ - SELECT id, username, key_prefix, name, created_at, last_used_at, expires_at, is_active - FROM api_keys - WHERE username = ? - ORDER BY created_at DESC - """, - (username,), - ) + if include_internal: + cur = conn.execute( + """ + SELECT id, username, key_prefix, name, created_at, last_used_at, + expires_at, is_active, is_internal + FROM api_keys + WHERE username = ? + ORDER BY created_at DESC + """, + (username,), + ) + else: + cur = conn.execute( + """ + SELECT id, username, key_prefix, name, created_at, last_used_at, + expires_at, is_active, is_internal + FROM api_keys + WHERE username = ? AND is_internal = 0 + ORDER BY created_at DESC + """, + (username,), + ) return [dict(row) for row in cur.fetchall()] finally: conn.close() @@ -652,6 +686,24 @@ def revoke_api_key(username: str, key_id: int) -> bool: conn.close() +def revoke_internal_api_key(key_id: int) -> bool: + """Revoke an internal workflow-minted key without requiring a username. + + Used by the recipe runner to retire its sk-unsloth-* key once the job + terminates, shrinking the window a leaked key could be abused. + """ + conn = get_connection() + try: + cursor = conn.execute( + "UPDATE api_keys SET is_active = 0 WHERE id = ? AND is_internal = 1", + (key_id,), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + def validate_api_key(raw_key: str) -> Optional[str]: """Validate *raw_key* and return the owning username, or ``None``. diff --git a/studio/backend/core/data_recipe/jobs/constants.py b/studio/backend/core/data_recipe/jobs/constants.py index 08237326f8..0045276e20 100644 --- a/studio/backend/core/data_recipe/jobs/constants.py +++ b/studio/backend/core/data_recipe/jobs/constants.py @@ -9,6 +9,7 @@ STAGE_PREVIEW = "preview" STAGE_DAG = "dag" STAGE_HEALTHCHECK = "healthcheck" STAGE_SAMPLING = "sampling" +STAGE_SOURCE = "source" STAGE_COLUMN_CONFIG = "column_config" STAGE_GENERATING = "generating" STAGE_BATCH = "batch" diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 3d7cf2dbe6..db59d573a9 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -33,6 +33,60 @@ from .worker import run_job_process _CTX = mp.get_context("spawn") +def _github_source_estimated_total(recipe: dict) -> int | None: + seed_config = recipe.get("seed_config") + if not isinstance(seed_config, dict): + return None + source = seed_config.get("source") + if not isinstance(source, dict) or source.get("seed_type") != "github_repo": + return None + + repos_raw = source.get("repos") + repos = ( + [repo for repo in repos_raw if isinstance(repo, str) and repo.strip()] + if isinstance(repos_raw, list) + else [] + ) + item_types_raw = source.get("item_types") + item_types = ( + [ + item + for item in item_types_raw + if isinstance(item, str) and item in {"issues", "pulls", "commits"} + ] + if isinstance(item_types_raw, list) + else [] + ) + try: + limit = int(source.get("limit") or 0) + except (TypeError, ValueError): + return None + if not repos or not item_types or limit <= 0: + return None + return len(repos) * len(item_types) * limit + + +def _source_progress_status(job: Job) -> dict[str, Any] | None: + progress = job.source_progress + if progress is None: + return None + return { + "source": progress.source, + "status": progress.status, + "repo": progress.repo, + "resource": progress.resource, + "page": progress.page, + "page_items": progress.page_items, + "fetched_items": progress.fetched_items, + "estimated_total": progress.estimated_total, + "percent": progress.percent, + "rate_remaining": progress.rate_remaining, + "retry_after_sec": progress.retry_after_sec, + "message": progress.message, + "updated_at": progress.updated_at, + } + + @dataclass class Subscription: replay: list[dict] @@ -71,8 +125,20 @@ class JobManager: self._pump_thread: threading.Thread | None = None self._seq: int = 0 - def start(self, *, recipe: dict, run: dict) -> str: - """Spawn the job subprocess (one at a time, no cap).""" + def start( + self, + *, + recipe: dict, + run: dict, + internal_api_key_id: int | None = None, + ) -> str: + """Spawn the job subprocess (one at a time, no cap). + + ``internal_api_key_id`` is the row id of a workflow-scoped + sk-unsloth-* key minted by the route layer for local providers. + JobManager revokes it when the job reaches a terminal state so the + key's live window is no longer than the run. + """ llm_columns = recipe.get("columns") or [] llm_column_count = 0 if isinstance(llm_columns, list): @@ -92,6 +158,10 @@ class JobManager: job_id = uuid.uuid4().hex self._job = Job(job_id = job_id, status = "pending", started_at = time.time()) self._job.progress_columns_total = llm_column_count + self._job.source_progress_estimated_total = _github_source_estimated_total( + recipe + ) + self._job.internal_api_key_id = internal_api_key_id self._events.clear() self._seq = 0 @@ -163,6 +233,7 @@ class JobManager: "ok": job.column_progress.ok, "failed": job.column_progress.failed, }, + "source_progress": _source_progress_status(job), "model_usage": { name: { "model": usage.model, @@ -405,6 +476,7 @@ class JobManager: for e in self._drain_queue(mp_q): self._handle_event(job, e) + retired_job: Job | None = None with self._lock: if self._job and self._job.status in { "pending", @@ -429,6 +501,9 @@ class JobManager: "job_id": self._job.job_id, } ) + retired_job = self._job + if retired_job is not None: + self._retire_workflow_key(retired_job) return def _handle_event(self, job: Job, event: dict) -> None: @@ -436,6 +511,7 @@ class JobManager: et = event.get("type") msg = event.get("message") if et == "log" else None + terminal = False with self._lock: if self._job is None or self._job.job_id != job.job_id: return @@ -452,18 +528,43 @@ class JobManager: if self._job.progress.total and self._job.progress.total > 0: self._job.progress.done = self._job.progress.total self._job.progress.percent = 100.0 + terminal = True if et == EVENT_JOB_ERROR: self._job.status = "error" self._job.finished_at = time.time() self._job.error = event.get("error") or "error" + terminal = True + if et == EVENT_JOB_CANCELLED: + terminal = True if msg: upd = parse_log_message(msg) if upd: apply_update(self._job, upd) + if terminal: + self._retire_workflow_key(job) + self._emit(event) + def _retire_workflow_key(self, job: Job) -> None: + """Revoke the workflow-scoped sk-unsloth-* key, if one was minted. + + Best-effort: revocation failures are swallowed. The key would + expire on its own after 24h, so a missed revoke is a latency + concern, not a correctness one. + """ + key_id = getattr(job, "internal_api_key_id", None) + if not key_id: + return + try: + from auth import storage # deferred: avoids circular import + + storage.revoke_internal_api_key(int(key_id)) + except Exception: + pass + job.internal_api_key_id = None + _JOB_MANAGER: JobManager | None = None diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 324b62a92e..cea6d8ea64 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -4,6 +4,7 @@ from __future__ import annotations import re +import time from dataclasses import dataclass from typing import Any @@ -17,9 +18,10 @@ from .constants import ( STAGE_PREVIEW, STAGE_PROFILING, STAGE_SAMPLING, + STAGE_SOURCE, USAGE_RESET_STAGES, ) -from .types import Job, ModelUsage, Progress +from .types import Job, ModelUsage, Progress, SourceProgress @dataclass(frozen = True) @@ -41,6 +43,7 @@ class ParsedUpdate: usage_requests_total: int | None = None usage_rpm: float | None = None usage_section_start: bool | None = None + source_progress: SourceProgress | None = None # kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI. @@ -61,9 +64,165 @@ _RE_USAGE_TOKENS = re.compile( _RE_USAGE_REQUESTS = re.compile( r"requests:\s*success=(?P\d+),\s*failed=(?P\d+),\s*total=(?P\d+),\s*rpm=(?P[0-9.]+)" ) +_RE_GITHUB_PAGE = re.compile( + r"^\[(?P[^\]\s]+/[^\]\s]+)\]\s+" + r"(?Pissues|PRs|commits)\s+page\s+(?P\d+)\s+" + r"\(\+(?P\d+)\).*?\bremaining=(?P\d+)", + re.IGNORECASE, +) +_RE_GITHUB_RATE_LIMIT = re.compile( + r"Rate limit hit\. Sleeping (?P\d+)s until reset\.", + re.IGNORECASE, +) +_RE_GITHUB_SECONDARY_RATE_LIMIT = re.compile( + r"Secondary rate limit(?: on REST)?\. Sleep (?P\d+)s\.", + re.IGNORECASE, +) +_RE_GITHUB_REST_RATE_LIMIT = re.compile( + r"REST 403/429, sleep (?P\d+)", + re.IGNORECASE, +) +_RE_GITHUB_TRANSIENT = re.compile( + r"^(?PGraphQL|REST) (?P\d{3}) transient, retrying", + re.IGNORECASE, +) +_RE_GITHUB_NETWORK_RETRY = re.compile( + r"^(?PGraphQL|REST) network error: .* Retry\.", + re.IGNORECASE, +) +_RE_GITHUB_TRIAL_LIMIT = re.compile( + r"Trial limit reached for (?Pissues|PRs|commits) \((?P\d+)\)", + re.IGNORECASE, +) +_RE_GITHUB_COMPLETE = re.compile( + r"Scraper complete\. GraphQL calls=\d+ REST calls=\d+", + re.IGNORECASE, +) def parse_log_message(msg: str) -> ParsedUpdate | None: + m = _RE_GITHUB_PAGE.search(msg) + if m: + resource_raw = m.group("resource") + resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower() + repo = m.group("repo") + page = int(m.group("page")) + page_items = int(m.group("items")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "fetching", + repo = repo, + resource = resource, + page = page, + page_items = page_items, + rate_remaining = int(m.group("remaining")), + message = ( + f"Scraping GitHub source: {repo} " + f"{resource} page {page} (+{page_items})" + ), + ), + ) + + m = _RE_GITHUB_RATE_LIMIT.search(msg) + if m: + seconds = int(m.group("seconds")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "rate_limited", + retry_after_sec = seconds, + message = ( + "Waiting for GitHub rate limit. " + "Studio will resume automatically." + ), + ), + ) + + m = _RE_GITHUB_SECONDARY_RATE_LIMIT.search(msg) + if m: + seconds = int(m.group("seconds")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "rate_limited", + retry_after_sec = seconds, + message = ( + "Waiting for GitHub secondary rate limit. " + "Studio will resume automatically." + ), + ), + ) + + m = _RE_GITHUB_REST_RATE_LIMIT.search(msg) + if m: + seconds = int(m.group("seconds")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "rate_limited", + retry_after_sec = seconds, + message = ( + "Waiting for GitHub rate limit. " + "Studio will resume automatically." + ), + ), + ) + + m = _RE_GITHUB_TRIAL_LIMIT.search(msg) + if m: + resource_raw = m.group("resource") + resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower() + items = int(m.group("items")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "fetching", + resource = resource, + message = f"GitHub {resource} trial limit reached ({items}).", + ), + ) + + m = _RE_GITHUB_TRANSIENT.search(msg) + if m: + api = m.group("api") + code = m.group("code") + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "retrying", + message = f"GitHub {api} returned {code}; retrying automatically.", + ), + ) + + m = _RE_GITHUB_NETWORK_RETRY.search(msg) + if m: + api = m.group("api") + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "retrying", + message = f"GitHub {api} request failed; retrying automatically.", + ), + ) + + if _RE_GITHUB_COMPLETE.search(msg): + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "completed", + message = "GitHub source scrape complete.", + ), + ) + m = _RE_SAMPLERS.search(msg) if m: return ParsedUpdate( @@ -172,6 +331,8 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: job.batch.idx = update.batch_idx if update.batch_total is not None: job.batch.total = update.batch_total + if update.source_progress is not None: + _apply_source_progress(job, update.source_progress) if update.stage in USAGE_RESET_STAGES: # usage summary is a short block so we reset once we move into the next stage. @@ -216,6 +377,67 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: usage.rpm = update.usage_rpm +def _apply_source_progress(job: Job, progress: SourceProgress) -> None: + previous = job.source_progress + now = time.time() + + page_items = progress.page_items + if progress.repo and progress.resource and progress.page is not None: + page_key = f"{progress.repo}:{progress.resource}:{progress.page}" + count_key = f"{progress.repo}:{progress.resource}" + if page_key not in job._source_seen_pages: + job._source_seen_pages.add(page_key) + job._source_counts[count_key] = int( + job._source_counts.get(count_key, 0) + ) + int(page_items or 0) + + fetched_items = sum(job._source_counts.values()) + if fetched_items <= 0: + fetched_items = progress.fetched_items or ( + previous.fetched_items if previous else None + ) + + estimated_total = ( + progress.estimated_total + or job.source_progress_estimated_total + or (previous.estimated_total if previous else None) + ) + percent: float | None = progress.percent + if percent is None and estimated_total and fetched_items is not None: + raw_percent = (float(fetched_items) / float(max(1, estimated_total))) * 100.0 + percent = 100.0 if progress.status == "completed" else min(99.0, raw_percent) + if percent is None and previous is not None: + percent = previous.percent + + job.source_progress = SourceProgress( + source = "github", + status = progress.status or (previous.status if previous else None), + repo = progress.repo or (previous.repo if previous else None), + resource = progress.resource or (previous.resource if previous else None), + page = ( + progress.page + if progress.page is not None + else (previous.page if previous else None) + ), + page_items = ( + page_items + if page_items is not None + else (previous.page_items if previous else None) + ), + fetched_items = fetched_items, + estimated_total = estimated_total, + percent = percent, + rate_remaining = ( + progress.rate_remaining + if progress.rate_remaining is not None + else (previous.rate_remaining if previous else None) + ), + retry_after_sec = progress.retry_after_sec, + message = progress.message or (previous.message if previous else None), + updated_at = now, + ) + + def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress: if not job.rows: return column_progress diff --git a/studio/backend/core/data_recipe/jobs/types.py b/studio/backend/core/data_recipe/jobs/types.py index 8d77903238..3d3ddb974e 100644 --- a/studio/backend/core/data_recipe/jobs/types.py +++ b/studio/backend/core/data_recipe/jobs/types.py @@ -35,6 +35,23 @@ class BatchProgress: total: int | None = None +@dataclass +class SourceProgress: + source: str = "github" + status: str | None = None + repo: str | None = None + resource: str | None = None + page: int | None = None + page_items: int | None = None + fetched_items: int | None = None + estimated_total: int | None = None + percent: float | None = None + rate_remaining: int | None = None + retry_after_sec: int | None = None + message: str | None = None + updated_at: float | None = None + + @dataclass class ModelUsage: model: str @@ -57,6 +74,7 @@ class Job: progress: Progress = field(default_factory = Progress) column_progress: Progress = field(default_factory = Progress) batch: BatchProgress = field(default_factory = BatchProgress) + source_progress: SourceProgress | None = None rows: int | None = None cols: int | None = None error: str | None = None @@ -70,8 +88,15 @@ class Job: processor_artifacts: dict[str, Any] | None = None model_usage: dict[str, ModelUsage] = field(default_factory = dict) progress_columns_total: int | None = None + source_progress_estimated_total: int | None = None completed_columns: list[str] = field(default_factory = list) + # Id of the internal sk-unsloth-* API key minted for a local-model + # workflow. Revoked when the job terminates so the key's live window + # matches the run rather than its 24h TTL. + internal_api_key_id: int | None = None _current_usage_model: str | None = None _in_usage_summary: bool = False _seen_generation_columns: list[str] = field(default_factory = list) _column_done: dict[str, int] = field(default_factory = dict) + _source_counts: dict[str, int] = field(default_factory = dict) + _source_seen_pages: set[str] = field(default_factory = set) diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index 63e38bd18d..8c5c7fe657 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -21,6 +21,15 @@ from ..service import build_config_builder, create_data_designer from utils.paths import ensure_dir, recipe_datasets_root _ARTIFACT_ROOT = recipe_datasets_root() +_RE_GITHUB_CURSOR = re.compile(r"\bcursor=[^\s,]+") +_RE_SECRET_TOKEN = re.compile( + r"\b(?:(?:ghp|gho|ghu|ghs|ghr|github_pat)_[A-Za-z0-9_]+|sk-unsloth-[A-Za-z0-9]+)" +) + + +def _sanitize_log_message(message: str) -> str: + message = _RE_GITHUB_CURSOR.sub("cursor=", message) + return _RE_SECRET_TOKEN.sub("", message) class _QueueLogHandler(logging.Handler): @@ -35,7 +44,7 @@ class _QueueLogHandler(logging.Handler): "ts": record.created, "level": record.levelname, "logger": record.name, - "message": record.getMessage(), + "message": _sanitize_log_message(record.getMessage()), } self._q.put(event) except (OSError, RuntimeError, ValueError): @@ -119,10 +128,16 @@ def run_job_process( # Attach queue logger directly to `data_designer` so parser events survive root resets. handler = _QueueLogHandler(event_queue) handler.setLevel(logging.INFO) - data_designer_logger = logging.getLogger("data_designer") - data_designer_logger.addHandler(handler) - data_designer_logger.setLevel(logging.INFO) - data_designer_logger.propagate = True + for logger_name in ( + "data_designer", + "scraper", + "gh_client", + "data_designer_github_repo_seed", + ): + logger = logging.getLogger(logger_name) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = True if run_config_raw: designer.set_run_config(RunConfig.model_validate(run_config_raw)) @@ -180,8 +195,8 @@ def run_job_process( { "type": EVENT_JOB_ERROR, "ts": time.time(), - "error": str(exc), - "stack": traceback.format_exc(limit = 20), + "error": _sanitize_log_message(str(exc)), + "stack": _sanitize_log_message(traceback.format_exc(limit = 20)), } ) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 81ff9ae4e5..33722f103c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1618,8 +1618,10 @@ class LlamaCppBackend: # Model fits on selected GPU(s) -- offload all layers cmd.extend(["-ngl", "-1"]) - if n_threads is not None: - cmd.extend(["--threads", str(n_threads)]) + # -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we + # do not inherit llama-server's internal default, which has historically + # varied (hardware concurrency incl. hyperthreads on some builds). + cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)]) # Always enable Jinja chat template rendering for proper template support cmd.extend(["--jinja"]) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/README.md b/studio/backend/plugins/data-designer-github-repo-seed/README.md new file mode 100644 index 0000000000..346d94b305 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/README.md @@ -0,0 +1,73 @@ +# data-designer-github-repo-seed + +A Data Designer seed-reader plugin for **Unsloth Studio** that scrapes real +GitHub data (issues, pull requests, commits) from one or more repositories +and hands it to the recipe pipeline as a seed dataset. + +Designed to ship with Studio as a default seed source so any user with a +GitHub token can build training datasets straight from live repos. + +## What it does + +Given a list of `owner/name` repos, a GitHub token, and a per-resource +`limit`, the plugin uses GitHub's GraphQL API to fetch issues, pull +requests, and/or commits, with labels, state, authors, and the first N +comments of each item, and materialises a single JSONL with uniform +columns so the rest of the recipe (LLM text / LLM structured / processors) +can treat it like any other seed table. + +| Column | Description | +|---------------|------------------------------------------------| +| `item_type` | `issue` / `pull` / `commit` | +| `repo` | `owner/name` | +| `number` | Issue/PR number, or commit SHA | +| `title` | Title (or commit message headline) | +| `body` | Issue/PR body (or full commit message) | +| `state` | `OPEN` / `CLOSED` / `MERGED` (empty for commit)| +| `author` | GitHub login of the author | +| `created_at` | ISO8601 | +| `closed_at` | ISO8601 (empty for commits) | +| `url` | Permalink | +| `labels` | List of label names | +| `comments` | First N comments concatenated | + +## Usage in a recipe + +```json +{ + "seed_config": { + "source": { + "seed_type": "github_repo", + "repos": ["unslothai/unsloth", "unslothai/unsloth-zoo"], + "token": "", + "item_types": ["issues", "pulls"], + "limit": 100, + "include_comments": true, + "max_comments_per_item": 30 + }, + "sampling_strategy": "shuffle", + "selection_strategy": null + } +} +``` + +Leave `token` empty to fall back to the server's `GH_TOKEN` / `GITHUB_TOKEN` +environment variable, useful when the recipe is published and shouldn't +carry a secret. + +## Auth + +A GitHub personal access token with `public_repo` scope is enough for public +repositories; `repo` scope is required for private ones. GraphQL requests +are rate-limit aware: the client inspects `x-ratelimit-*` headers and +sleeps until reset when the budget drops below a safety threshold. + +## Install + +Shipped as a default Studio plugin. For development: + +```bash +pip install -e . +``` + +Registered automatically via the `data_designer.plugins` entry point. diff --git a/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml new file mode 100644 index 0000000000..e232adc60c --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "data-designer-github-repo-seed" +version = "0.1.0" +description = "Unsloth Studio seed plugin that scrapes GitHub issues, PRs, and commits." +requires-python = ">=3.11" +dependencies = [ + "data-designer-engine>=0.5.4,<0.6", + "requests>=2.31", +] + +[project.entry-points."data_designer.plugins"] +github_repo_seed = "data_designer_github_repo_seed.plugin:github_repo_seed_plugin" + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py new file mode 100644 index 0000000000..f57af4c6c3 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +# Intentionally empty. Data-designer loads submodules lazily via qualified names +# (impl_qualified_name / config_qualified_name in plugin.py), so importing this +# package must NOT touch modules that depend on data_designer.engine.* during +# Studio's bootstrap (circular import). diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py new file mode 100644 index 0000000000..6b347c4f83 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from typing import Literal + +from pydantic import Field, field_validator, model_validator + +from data_designer.config.seed_source import SeedSource + + +class GitHubRepoSeedSource(SeedSource): + seed_type: Literal["github_repo"] = "github_repo" + + repos: list[str] = Field( + default_factory = list, + description = "List of GitHub repositories to scrape, each in `owner/name` form.", + ) + token: str = Field( + default = "", + description = "Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.", + ) + item_types: list[Literal["issues", "pulls", "commits"]] = Field( + default = ["issues", "pulls"], + description = "Which GitHub item types to fetch per repo.", + ) + limit: int = Field( + default = 100, + ge = 1, + le = 5000, + description = "Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).", + ) + include_comments: bool = Field( + default = True, + description = "Fetch the first N comments of each issue/PR and include them in the `comments` column.", + ) + max_comments_per_item: int = Field(default = 30, ge = 0, le = 200) + + @field_validator("repos") + @classmethod + def _validate_repos(cls, v: list[str]) -> list[str]: + out: list[str] = [] + for r in v or []: + r = r.strip() + if not r: + continue + if r.count("/") != 1 or not all(r.split("/")): + raise ValueError(f"Each repo must be `owner/name`; got {r!r}") + out.append(r) + return out + + @field_validator("item_types") + @classmethod + def _validate_item_types(cls, v: list[str]) -> list[str]: + if not v: + raise ValueError("item_types must not be empty") + return list(dict.fromkeys(v)) + + @model_validator(mode = "after") + def _ensure_repos(self) -> "GitHubRepoSeedSource": + if not self.repos: + raise ValueError("At least one repo is required") + return self diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py new file mode 100644 index 0000000000..5a38e26d6b --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import hashlib +import tempfile +import threading +from pathlib import Path +from typing import Optional + +import data_designer.lazy_heavy_imports as lazy +from data_designer.engine.resources.seed_reader import SeedReader + +from .config import GitHubRepoSeedSource +from .scraper import ScrapeConfig, materialize_to_jsonl + + +# In-process cache mapping a stable config signature to the JSONL materialization +# path. A single recipe job invokes the seed reader multiple times (validation, +# preview, per-column sampling), and the default flow re-scrapes the repo on +# every call: for a 2-repo preview that is ~15s of redundant GitHub GraphQL +# traffic before any generation fires. Memoize the materialization so the second +# and third passes reuse the file the first pass wrote. Cache key excludes the +# raw token and uses a short SHA-256 digest so token values never hit memory +# twice and token rotation invalidates cleanly. +_SCRAPE_CACHE: dict[tuple, str] = {} +_SCRAPE_CACHE_LOCK = threading.Lock() + + +def _scrape_cache_key(cfg: ScrapeConfig) -> tuple: + token_digest = hashlib.sha256( + (cfg.token or "").encode("utf-8"), + ).hexdigest()[:16] + return ( + tuple(cfg.repos), + tuple(cfg.item_types), + cfg.limit, + bool(cfg.include_comments), + cfg.max_comments_per_item, + token_digest, + ) + + +def _lookup_cached_scrape(key: tuple) -> Optional[str]: + with _SCRAPE_CACHE_LOCK: + path = _SCRAPE_CACHE.get(key) + if path and Path(path).exists(): + return path + # Stale entry (tmp cleanup, user restarted, ...); drop it so the caller + # materializes a fresh file rather than returning a dangling path. + if path: + with _SCRAPE_CACHE_LOCK: + _SCRAPE_CACHE.pop(key, None) + return None + + +def _store_cached_scrape(key: tuple, path: str) -> None: + with _SCRAPE_CACHE_LOCK: + _SCRAPE_CACHE[key] = path + + +class GitHubRepoSeedReader(SeedReader[GitHubRepoSeedSource]): + def create_duckdb_connection(self): + return lazy.duckdb.connect() + + def get_dataset_uri(self) -> str: + out_dir = Path(tempfile.gettempdir()) / "studio-github-repo-seed" + cfg = ScrapeConfig( + repos = list(self.source.repos), + token = self.source.token, + item_types = list(self.source.item_types), + limit = self.source.limit, + include_comments = self.source.include_comments, + max_comments_per_item = self.source.max_comments_per_item, + ) + cache_key = _scrape_cache_key(cfg) + cached_path = _lookup_cached_scrape(cache_key) + if cached_path is not None: + return cached_path + path = materialize_to_jsonl(cfg, out_dir) + _store_cached_scrape(cache_key, str(path)) + return str(path) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py new file mode 100644 index 0000000000..f87dbd0507 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from data_designer.plugins.plugin import Plugin, PluginType + +github_repo_seed_plugin = Plugin( + impl_qualified_name = "data_designer_github_repo_seed.impl.GitHubRepoSeedReader", + config_qualified_name = "data_designer_github_repo_seed.config.GitHubRepoSeedSource", + plugin_type = PluginType.SEED_READER, +) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py new file mode 100644 index 0000000000..d768fe37be --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Multi-repo GitHub scraper for the Studio seed plugin. + +Drives the GraphQL-based scraper in `scraper_impl/` per repo. Each repo is +scraped with a trial_limits cap so we stop at `limit` items per resource. +After scraping, we read the per-resource JSONL shards and flatten them into +a single unified JSONL with stable columns (`item_type`, `repo`, `number`, +`title`, `body`, ...). +""" + +from __future__ import annotations + +import json +import os +import sys +import time +import uuid +from dataclasses import dataclass +from pathlib import Path + +# Defer scraper_impl imports until `scrape()` runs with a resolved token. +_IMPL_DIR = Path(__file__).parent / "scraper_impl" + + +def _ensure_impl_on_path() -> None: + if str(_IMPL_DIR) not in sys.path: + sys.path.insert(0, str(_IMPL_DIR)) + + +def _load_impl(): + _ensure_impl_on_path() + import importlib + + gh_client = importlib.import_module("gh_client") # type: ignore + scraper_mod = importlib.import_module("scraper") # type: ignore + return gh_client.GitHubClient, scraper_mod.RepoScraper + + +@dataclass +class ScrapeConfig: + repos: list[str] + token: str + item_types: list[str] + limit: int + include_comments: bool + max_comments_per_item: int + + +def _resolve_token(token: str) -> str: + tok = token or os.environ.get("GH_TOKEN", "") or os.environ.get("GITHUB_TOKEN", "") + if not tok: + raise ValueError( + "GitHub token is required. Set it in the recipe config or the GH_TOKEN / GITHUB_TOKEN env var." + ) + return tok + + +def _read_jsonl(path: Path, max_rows: int | None = None): + if not path.exists(): + return + with path.open(encoding = "utf-8") as f: + for i, line in enumerate(f): + if not line.strip(): + continue + if max_rows is not None and i >= max_rows: + return + try: + yield json.loads(line) + except json.JSONDecodeError: + continue + + +def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] + comments_nodes = (r.get("comments") or {}).get("nodes") or [] + comments_text = "" + if include_comments and comments_nodes: + kept = comments_nodes[:max_c] + comments_text = "\n\n".join( + f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" + for c in kept + ) + return { + "item_type": "issue", + "repo": repo, + "number": r.get("number"), + "title": r.get("title") or "", + "body": r.get("body") or "", + "state": r.get("state") or "", + "author": (r.get("author") or {}).get("login", ""), + "created_at": r.get("createdAt") or "", + "closed_at": r.get("closedAt") or "", + "url": r.get("url") or r.get("permalink") or "", + "labels": labels, + "comments": comments_text, + } + + +def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] + comments_nodes = (r.get("comments") or {}).get("nodes") or [] + comments_text = "" + if include_comments and comments_nodes: + kept = comments_nodes[:max_c] + comments_text = "\n\n".join( + f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" + for c in kept + ) + return { + "item_type": "pull", + "repo": repo, + "number": r.get("number"), + "title": r.get("title") or "", + "body": r.get("body") or "", + "state": r.get("state") or "", + "author": (r.get("author") or {}).get("login", ""), + "created_at": r.get("createdAt") or "", + "closed_at": r.get("closedAt") or "", + "url": r.get("url") or r.get("permalink") or "", + "labels": labels, + "comments": comments_text, + } + + +def _flatten_commit_row(r: dict, repo: str) -> dict: + msg = r.get("messageHeadline") or r.get("message") or "" + body = r.get("messageBody") or r.get("message") or msg + author = r.get("author") or {} + return { + "item_type": "commit", + "repo": repo, + "number": r.get("oid") or r.get("sha") or "", + "title": msg, + "body": body, + "state": "", + "author": (author.get("user") or {}).get("login") or author.get("name", ""), + "created_at": (author.get("date") or r.get("committedDate") or ""), + "closed_at": "", + "url": r.get("url") or "", + "labels": [], + "comments": "", + } + + +def scrape(cfg: ScrapeConfig, base_dir: Path): + token = _resolve_token(cfg.token) + GitHubClient, RepoScraper = _load_impl() + client = GitHubClient(token = token) + base_dir.mkdir(parents = True, exist_ok = True) + + # Per-resource trial limits. limit <= 0 means "all": use a very large cap. + effective_limit = cfg.limit if cfg.limit and cfg.limit > 0 else 1_000_000 + trial_limits: dict[str, int] = {} + if "issues" in cfg.item_types: + trial_limits["issues"] = effective_limit + if "pulls" in cfg.item_types: + trial_limits["pull_requests"] = effective_limit + if "commits" in cfg.item_types: + trial_limits["commits"] = effective_limit + + all_rows: list[dict] = [] + for repo in cfg.repos: + owner, name = repo.split("/", 1) + scraper = RepoScraper( + owner = owner, + name = name, + base_dir = base_dir, + client = client, + trial_limits = trial_limits, + light = True, + ) + try: + repo_meta = scraper.scrape_repo_meta() + if "issues" in cfg.item_types: + scraper.scrape_issues() + if "pulls" in cfg.item_types: + scraper.scrape_prs() + if "commits" in cfg.item_types: + default_ref = repo_meta.get("defaultBranchRef") or {} + default_branch = ( + default_ref.get("name") if isinstance(default_ref, dict) else None + ) + branch = ( + f"refs/heads/{default_branch}" + if default_branch + else "refs/heads/main" + ) + scraper.scrape_commits(branch = branch) + finally: + scraper.close() + + read_cap = cfg.limit if cfg.limit and cfg.limit > 0 else None + repo_dir = base_dir / f"{owner}__{name}" + if "issues" in cfg.item_types: + for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap): + all_rows.append( + _flatten_issue_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) + ) + if "pulls" in cfg.item_types: + for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap): + all_rows.append( + _flatten_pr_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) + ) + if "commits" in cfg.item_types: + for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap): + all_rows.append(_flatten_commit_row(row, repo)) + + return all_rows + + +def materialize_to_jsonl(cfg: ScrapeConfig, out_dir: Path) -> Path: + out_dir.mkdir(parents = True, exist_ok = True) + tag = "-".join(r.replace("/", "__") for r in cfg.repos)[:120] + kinds = "-".join(cfg.item_types) + run_id = f"{int(time.time())}-{uuid.uuid4().hex[:12]}" + fname = f"github_{tag}__{kinds}__{cfg.limit}_{run_id}.jsonl" + out = out_dir / fname + rows = scrape(cfg, out_dir / "raw-runs" / run_id) + with out.open("w", encoding = "utf-8") as f: + for r in rows: + f.write(json.dumps(r, ensure_ascii = False) + "\n") + return out diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py new file mode 100644 index 0000000000..dd2de2f5ce --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GitHub API client with rate-limit awareness, retry, and dual REST/GraphQL support.""" + +from __future__ import annotations + +import json +import os +import time +import logging +from typing import Any, Dict, Iterable, Iterator, List, Optional + +import requests + +log = logging.getLogger("gh_client") + +GRAPHQL_URL = "https://api.github.com/graphql" +REST_BASE = "https://api.github.com" + +BASE_HEADERS = { + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + "User-Agent": "github-data-gatherer/1.0", +} + + +class RateLimitError(Exception): + pass + + +class GitHubClient: + def __init__( + self, + min_remaining_graphql: int = 100, + min_remaining_rest: int = 100, + token: str | None = None, + ): + token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN") + if not token: + raise RuntimeError("GH_TOKEN not set in environment") + self.session = requests.Session() + self.session.headers.update( + {**BASE_HEADERS, "Authorization": f"Bearer {token}"} + ) + self.min_remaining_graphql = min_remaining_graphql + self.min_remaining_rest = min_remaining_rest + self.graphql_remaining: Optional[int] = None + self.graphql_reset: Optional[int] = None + self.rest_remaining: Optional[int] = None + self.rest_reset: Optional[int] = None + self.calls_graphql = 0 + self.calls_rest = 0 + self.retry_count = 0 + + def _sleep_until(self, reset_ts: int, buffer_s: int = 10) -> None: + now = int(time.time()) + wait = max(0, reset_ts - now) + buffer_s + log.warning("Rate limit hit. Sleeping %ds until reset.", wait) + time.sleep(wait) + + def _check_rate_and_wait(self, kind: str) -> None: + if kind == "graphql": + remaining = self.graphql_remaining + reset = self.graphql_reset + min_remaining = self.min_remaining_graphql + else: + remaining = self.rest_remaining + reset = self.rest_reset + min_remaining = self.min_remaining_rest + if remaining is not None and remaining < min_remaining: + if reset: + self._sleep_until(reset) + # Reset remaining so we don't spin + if kind == "graphql": + self.graphql_remaining = None + else: + self.rest_remaining = None + + def graphql( + self, + query: str, + variables: Optional[Dict[str, Any]] = None, + max_retries: int = 20, + ) -> Dict[str, Any]: + self._check_rate_and_wait("graphql") + backoff = 2 + last_err = None + for attempt in range(max_retries): + try: + r = self.session.post( + GRAPHQL_URL, + json = {"query": query, "variables": variables or {}}, + timeout = 120, + ) + self.calls_graphql += 1 + # Update rate info from response headers + rem = r.headers.get("X-RateLimit-Remaining") + rst = r.headers.get("X-RateLimit-Reset") + if rem is not None: + try: + self.graphql_remaining = int(rem) + except ValueError: + pass + if rst is not None: + try: + self.graphql_reset = int(rst) + except ValueError: + pass + if r.status_code in (502, 503, 504): + log.warning("GraphQL %s transient, retrying", r.status_code) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + continue + if r.status_code == 403 or r.status_code == 429: + # Check for secondary/abuse + retry_after = r.headers.get("Retry-After") + if retry_after: + t = int(retry_after) + log.warning("Secondary rate limit. Sleep %ds.", t) + time.sleep(t + 2) + continue + if self.graphql_reset: + self._sleep_until(self.graphql_reset) + continue + time.sleep(60) + continue + r.raise_for_status() + data = r.json() + if "errors" in data and data["errors"]: + # Surface errors but allow partial data + errs = data["errors"] + # Retry on RATE_LIMITED + for e in errs: + if e.get("type") == "RATE_LIMITED": + self._sleep_until( + (self.graphql_reset or int(time.time()) + 60) + ) + break + else: + # No rate-limit error, log and return partial + log.warning("GraphQL errors: %s", json.dumps(errs)[:400]) + return data + continue + return data + except requests.RequestException as e: + last_err = e + log.warning("GraphQL network error: %s. Retry.", e) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + raise RuntimeError(f"GraphQL failed after {max_retries} retries: {last_err}") + + def rest( + self, + method: str, + path: str, + params: Optional[Dict[str, Any]] = None, + json_body: Optional[Dict[str, Any]] = None, + max_retries: int = 6, + ) -> requests.Response: + self._check_rate_and_wait("rest") + if path.startswith("http"): + url = path + else: + url = REST_BASE + path + backoff = 2 + last_err = None + for attempt in range(max_retries): + try: + r = self.session.request( + method, url, params = params, json = json_body, timeout = 120 + ) + self.calls_rest += 1 + rem = r.headers.get("X-RateLimit-Remaining") + rst = r.headers.get("X-RateLimit-Reset") + if rem is not None: + try: + self.rest_remaining = int(rem) + except ValueError: + pass + if rst is not None: + try: + self.rest_reset = int(rst) + except ValueError: + pass + if r.status_code in (502, 503, 504): + log.warning("REST %s transient, retrying", r.status_code) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + continue + if r.status_code in (403, 429): + retry_after = r.headers.get("Retry-After") + if retry_after: + t = int(retry_after) + log.warning("Secondary rate limit on REST. Sleep %ds.", t) + time.sleep(t + 2) + continue + # Check if primary rate + if self.rest_remaining == 0 and self.rest_reset: + self._sleep_until(self.rest_reset) + continue + log.warning("REST 403/429, sleep 60") + time.sleep(60) + continue + return r + except requests.RequestException as e: + last_err = e + log.warning("REST network error: %s. Retry.", e) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + raise RuntimeError(f"REST failed after {max_retries} retries: {last_err}") + + def rest_paginate( + self, path: str, params: Optional[Dict[str, Any]] = None, per_page: int = 100 + ) -> Iterator[dict]: + params = dict(params or {}) + params.setdefault("per_page", per_page) + url = path + while True: + r = self.rest("GET", url, params = params if url == path else None) + if r.status_code != 200: + log.error( + "REST paginate got %s at %s: %s", r.status_code, url, r.text[:200] + ) + return + items = r.json() + if isinstance(items, dict): + # Some endpoints return dict with list field + items = items.get("items", []) + for it in items: + yield it + # Follow link header + link = r.headers.get("Link", "") + nxt = None + for part in link.split(","): + if 'rel="next"' in part: + nxt = part.split(";")[0].strip().strip("<>") + break + if not nxt: + return + url = nxt + params = None + + def rate_snapshot(self) -> Dict[str, Any]: + r = self.rest("GET", "/rate_limit") + if r.status_code == 200: + return r.json() + return {} diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py new file mode 100644 index 0000000000..9dc7613db5 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py @@ -0,0 +1,685 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GraphQL queries for GitHub data scraping. + +GitHub's GraphQL rejects queries that define unused fragments, so each query +only includes the fragments it actually references. +""" + +# ---- Fragments (kept as raw strings, composed per query) ---- +F_ACTOR = """ +fragment ActorFields on Actor { + __typename + login + url + avatarUrl + ... on User { id databaseId name } + ... on Bot { id databaseId } + ... on Organization { id databaseId name } +} +""" + +F_LABEL = """ +fragment LabelFields on Label { + id + name + color + description + createdAt +} +""" + +F_TIMELINE = """ +fragment TimelineItem on IssueTimelineItems { + __typename + ... on Node { id } + ... on AddedToProjectEvent { createdAt actor { ...ActorFields } } + ... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on ClosedEvent { createdAt actor { ...ActorFields } stateReason closer { __typename ... on Commit { oid url } ... on PullRequest { number url } } } + ... on CommentDeletedEvent { createdAt actor { ...ActorFields } } + ... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url repository { nameWithOwner } } ... on PullRequest { number url repository { nameWithOwner } } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on ConvertedNoteToIssueEvent { createdAt actor { ...ActorFields } } + ... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } } + ... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } } + ... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on LockedEvent { createdAt actor { ...ActorFields } lockReason } + ... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on MentionedEvent { createdAt actor { ...ActorFields } } + ... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } } + ... on PinnedEvent { createdAt actor { ...ActorFields } } + ... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } } + ... on RemovedFromProjectEvent { createdAt actor { ...ActorFields } } + ... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle } + ... on ReopenedEvent { createdAt actor { ...ActorFields } } + ... on SubscribedEvent { createdAt actor { ...ActorFields } } + ... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } } + ... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on UnlockedEvent { createdAt actor { ...ActorFields } } + ... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } } + ... on UnpinnedEvent { createdAt actor { ...ActorFields } } + ... on UnsubscribedEvent { createdAt actor { ...ActorFields } } + ... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration } +} +""" + +F_PR_TIMELINE = """ +fragment PRTimelineItem on PullRequestTimelineItems { + __typename + ... on Node { id } + ... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on AutoMergeDisabledEvent { createdAt actor { ...ActorFields } reason } + ... on AutoMergeEnabledEvent { createdAt actor { ...ActorFields } } + ... on AutoRebaseEnabledEvent { createdAt actor { ...ActorFields } } + ... on AutoSquashEnabledEvent { createdAt actor { ...ActorFields } } + ... on AutomaticBaseChangeFailedEvent { createdAt actor { ...ActorFields } oldBase newBase } + ... on AutomaticBaseChangeSucceededEvent { createdAt actor { ...ActorFields } oldBase newBase } + ... on BaseRefChangedEvent { createdAt actor { ...ActorFields } previousRefName currentRefName } + ... on BaseRefDeletedEvent { createdAt actor { ...ActorFields } baseRefName } + ... on BaseRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } } + ... on ClosedEvent { createdAt actor { ...ActorFields } stateReason } + ... on CommentDeletedEvent { createdAt actor { ...ActorFields } } + ... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url } ... on PullRequest { number url } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on ConvertToDraftEvent { createdAt actor { ...ActorFields } } + ... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } } + ... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on DeployedEvent { createdAt actor { ...ActorFields } } + ... on DeploymentEnvironmentChangedEvent { createdAt actor { ...ActorFields } } + ... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on HeadRefDeletedEvent { createdAt actor { ...ActorFields } headRefName } + ... on HeadRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } } + ... on HeadRefRestoredEvent { createdAt actor { ...ActorFields } } + ... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } } + ... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on LockedEvent { createdAt actor { ...ActorFields } lockReason } + ... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on MentionedEvent { createdAt actor { ...ActorFields } } + ... on MergedEvent { createdAt actor { ...ActorFields } commit { oid url } mergeRefName } + ... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } } + ... on PinnedEvent { createdAt actor { ...ActorFields } } + ... on PullRequestCommit { commit { oid url message author { user { login } date } committedDate } } + ... on PullRequestCommitCommentThread { commit { oid } } + ... on PullRequestReview { id databaseId createdAt submittedAt author { ...ActorFields } body state url reactionGroups { content reactors { totalCount } } } + ... on PullRequestReviewThread { id isResolved isOutdated path line diffSide } + ... on PullRequestRevisionMarker { createdAt lastSeenCommit { oid } } + ... on ReadyForReviewEvent { createdAt actor { ...ActorFields } } + ... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } } + ... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle } + ... on ReopenedEvent { createdAt actor { ...ActorFields } } + ... on ReviewDismissedEvent { createdAt actor { ...ActorFields } dismissalMessage previousReviewState } + ... on ReviewRequestRemovedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } } + ... on ReviewRequestedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } } + ... on SubscribedEvent { createdAt actor { ...ActorFields } } + ... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } } + ... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on UnlockedEvent { createdAt actor { ...ActorFields } } + ... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } } + ... on UnpinnedEvent { createdAt actor { ...ActorFields } } + ... on UnsubscribedEvent { createdAt actor { ...ActorFields } } + ... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration } +} +""" + + +def _q(parts: list[str], body: str) -> str: + return "\n".join(parts + [body]) + + +ISSUES_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL, F_TIMELINE], + """ +query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state stateReason + createdAt updatedAt closedAt + url + author { ...ActorFields } + editor { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + assignees(first: 20) { nodes { login id } } + milestone { title number state dueOn } + reactionGroups { content reactors { totalCount } } + comments(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + timelineItems(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { ...TimelineItem } + } + trackedInIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } } + trackedIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PRS_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL, F_PR_TIMELINE], + """ +query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state isDraft + createdAt updatedAt closedAt mergedAt + url + headRefName headRefOid + baseRefName baseRefOid + additions deletions changedFiles + mergeable merged mergeStateStatus + author { ...ActorFields } + editor { ...ActorFields } + mergedBy { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + assignees(first: 20) { nodes { login id } } + milestone { title number state dueOn } + reactionGroups { content reactors { totalCount } } + closingIssuesReferences(first: 20) { totalCount nodes { number url repository { nameWithOwner } title } } + comments(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + reviewThreads(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id isResolved isOutdated path line diffSide + comments(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body path diffHunk + author { ...ActorFields } + editor { ...ActorFields } + position originalPosition line originalLine + commit { oid } + reactionGroups { content reactors { totalCount } } + } + } + } + } + reviews(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId state createdAt submittedAt body url + author { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + commits(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + commit { + oid + message + messageHeadline + committedDate + authoredDate + author { name email user { login } date } + committer { name email user { login } date } + additions deletions changedFilesIfAvailable + parents(first: 3) { nodes { oid } } + } + } + } + files(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + path additions deletions changeType + } + } + timelineItems(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { ...PRTimelineItem } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PRS_PAGE_QUERY_LIGHT = _q( + [F_ACTOR, F_LABEL], + """ +query PRsPageLight($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state isDraft + createdAt updatedAt closedAt mergedAt + url + author { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + comments(first: 30) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +ISSUES_PAGE_QUERY_LIGHT = _q( + [F_ACTOR, F_LABEL], + """ +query IssuesPageLight($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state + createdAt updatedAt closedAt + url + author { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + comments(first: 30) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +ISSUE_COMMENTS_QUERY = _q( + [F_ACTOR], + """ +query IssueComments($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issueOrPullRequest(number: $number) { + __typename + ... on Issue { + comments(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + ... on PullRequest { + comments(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +ISSUE_TIMELINE_QUERY = _q( + [F_ACTOR, F_TIMELINE], + """ +query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issue(number: $number) { + timelineItems(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { ...TimelineItem } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PR_TIMELINE_QUERY = _q( + [F_ACTOR, F_PR_TIMELINE], + """ +query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + timelineItems(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { ...PRTimelineItem } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PR_COMMITS_QUERY = """ +query PRCommits($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + commits(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + commit { + oid message messageHeadline committedDate authoredDate + author { name email user { login } date } + committer { name email user { login } date } + additions deletions changedFilesIfAvailable + parents(first: 3) { nodes { oid } } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +PR_FILES_QUERY = """ +query PRFiles($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + files(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { path additions deletions changeType } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +PR_REVIEW_THREADS_QUERY = _q( + [F_ACTOR], + """ +query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + reviewThreads(first: 50, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id isResolved isOutdated path line diffSide + comments(first: 50) { + totalCount + nodes { + id databaseId createdAt updatedAt url body path diffHunk + author { ...ActorFields } + editor { ...ActorFields } + position originalPosition line originalLine + commit { oid } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +DISCUSSIONS_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL], + """ +query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + discussions(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body + createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + locked + answerChosenAt + closed closedAt + category { id name emoji description isAnswerable } + labels(first: 30) { nodes { ...LabelFields } } + upvoteCount + answer { id databaseId body author { ...ActorFields } createdAt url } + reactionGroups { content reactors { totalCount } } + comments(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + upvoteCount + isAnswer + reactionGroups { content reactors { totalCount } } + replies(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +DISCUSSION_COMMENTS_QUERY = _q( + [F_ACTOR], + """ +query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + discussion(number: $number) { + comments(first: 50, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + upvoteCount + isAnswer + reactionGroups { content reactors { totalCount } } + replies(first: 50) { + totalCount + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +DISCUSSION_REPLIES_QUERY = _q( + [F_ACTOR], + """ +query DiscussionReplies($commentId: ID!, $after: String) { + node(id: $commentId) { + ... on DiscussionComment { + replies(first: 50, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +COMMITS_PAGE_QUERY = """ +query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String, $branch: String!) { + repository(owner: $owner, name: $name) { + ref(qualifiedName: $branch) { + target { + ... on Commit { + history(first: $first, after: $after) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + oid + message + messageHeadline + committedDate + authoredDate + url + additions deletions changedFilesIfAvailable + author { name email date user { login id } } + committer { name email date user { login id } } + parents(first: 3) { nodes { oid } } + associatedPullRequests(first: 5) { nodes { number url state } } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +RELEASES_QUERY = _q( + [F_ACTOR], + """ +query Releases($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + releases(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId name tagName description + createdAt publishedAt updatedAt + isDraft isPrerelease isLatest + url + author { ...ActorFields } + tagCommit { oid url } + reactionGroups { content reactors { totalCount } } + releaseAssets(first: 50) { + nodes { name contentType size downloadUrl createdAt updatedAt } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +LABELS_QUERY = _q( + [F_LABEL], + """ +query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + labels(first: $first, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { ...LabelFields } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +MILESTONES_QUERY = """ +query Milestones($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + milestones(first: $first, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id number title description state + createdAt updatedAt closedAt dueOn + creator { login } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +REPO_META_QUERY = """ +query RepoMeta($owner: String!, $name: String!) { + repository(owner: $owner, name: $name) { + id databaseId name nameWithOwner description url + createdAt updatedAt pushedAt + isArchived isDisabled isFork isPrivate + primaryLanguage { name } + languages(first: 20, orderBy: {field: SIZE, direction: DESC}) { + edges { size node { name } } + totalSize + } + stargazerCount forkCount watchers { totalCount } + diskUsage + licenseInfo { key name } + homepageUrl + defaultBranchRef { name } + } + rateLimit { cost remaining resetAt } +} +""" diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py new file mode 100644 index 0000000000..127129e18b --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py @@ -0,0 +1,756 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Main scraper orchestration. Collects issues, PRs, discussions, commits, releases, etc. + +Resumable via state file. Writes JSONL shards under data/{repo}/{resource}.jsonl. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple + +# Allow running as a module or script +THIS_DIR = Path(__file__).resolve().parent +if str(THIS_DIR) not in sys.path: + sys.path.insert(0, str(THIS_DIR)) + +from gh_client import GitHubClient +from state_store import JsonlWriter, StateStore +import queries as Q + +log = logging.getLogger("scraper") + + +def ts() -> str: + return time.strftime("%Y-%m-%d %H:%M:%S") + + +class RepoScraper: + def __init__( + self, + owner: str, + name: str, + base_dir: Path, + client: GitHubClient, + trial_limits: Optional[Dict[str, int]] = None, + light: bool = False, + ): + self.owner = owner + self.name = name + self.base_dir = base_dir + self.client = client + self.trial_limits = trial_limits or {} + # When light=True, use trimmed GraphQL queries (no reviewThreads, + # reviews, commits, timelineItems, files) so PR pages can be much + # larger without blowing GitHub's node-count ceiling. + self.light = light + self.repo_dir = base_dir / f"{owner}__{name}" + self.repo_dir.mkdir(parents = True, exist_ok = True) + self.state = StateStore(base_dir / "state" / f"{owner}__{name}.json") + + # Writers + self.writers: Dict[str, JsonlWriter] = {} + for key in ( + "issues", + "pull_requests", + "discussions", + "commits", + "releases", + "labels", + "milestones", + "pr_extra_comments", + "pr_extra_timeline", + "pr_extra_reviews", + "issue_extra_comments", + "issue_extra_timeline", + "discussion_extra_comments", + "discussion_extra_replies", + "repo_meta", + ): + self.writers[key] = JsonlWriter(self.repo_dir / f"{key}.jsonl") + + # ----- helpers ----- + def _trial_stop(self, key: str, counter: int) -> bool: + lim = self.trial_limits.get(key) + if lim is None: + return False + return counter >= lim + + def _log_rate(self, where: str, data: Dict[str, Any]) -> None: + rl = ( + data.get("data", {}).get("rateLimit") + if isinstance(data.get("data"), dict) + else None + ) + if rl: + log.debug( + "[%s] rate cost=%s remaining=%s resetAt=%s", + where, + rl.get("cost"), + rl.get("remaining"), + rl.get("resetAt"), + ) + + # ----- repo meta ----- + def scrape_repo_meta(self) -> Dict[str, Any]: + data = self.client.graphql( + Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name} + ) + self._log_rate("repo_meta", data) + repo = data.get("data", {}).get("repository") or {} + repo["_fetchedAt"] = ts() + self.writers["repo_meta"].write(repo) + return repo + + # ----- issues ----- + def scrape_issues(self) -> int: + key = "issues" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + log.info("%s/%s issues already complete", self.owner, self.name) + return 0 + total_new = 0 + page = 0 + # Light query skips heavy nested fields; safe at 50 per page. + # Clamp by trial_limit so e.g. limit=1 asks GitHub for first:1 + # instead of fetching a full 50-item page and discarding 49. + page_cap = 50 if self.light else 15 + trial_cap = self.trial_limits.get(key) + per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } + query = Q.ISSUES_PAGE_QUERY_LIGHT if self.light else Q.ISSUES_PAGE_QUERY + data = self.client.graphql(query, vars_) + self._log_rate("issues", data) + repo = (data.get("data") or {}).get("repository") or {} + issues = repo.get("issues") or {} + nodes = issues.get("nodes") or [] + for it in nodes: + it["_owner"] = self.owner + it["_repo"] = self.name + it["_fetchedAt"] = ts() + if not self.light: + if it.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_issue_comments( + it["number"], it["comments"]["pageInfo"]["endCursor"] + ) + if ( + it.get("timelineItems", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): + self._paginate_issue_timeline( + it["number"], + it["timelineItems"]["pageInfo"]["endCursor"], + ) + if self.writers[key].write(it): + total_new += 1 + info = issues.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] issues page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + log.info("Trial limit reached for issues (%d)", total_new) + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + def _paginate_issue_comments(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "issueOrPullRequest" + ) or {} + comments = item.get("comments") or {} + for c in comments.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_issueNumber"] = number + self.writers["issue_extra_comments"].write(c) + info = comments.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_issue_timeline(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.ISSUE_TIMELINE_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get("issue") or {} + tl = item.get("timelineItems") or {} + for ev in tl.get("nodes") or []: + ev["_owner"] = self.owner + ev["_repo"] = self.name + ev["_issueNumber"] = number + self.writers["issue_extra_timeline"].write(ev) + info = tl.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + # ----- PRs ----- + def scrape_prs(self) -> int: + key = "pull_requests" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + log.info("%s/%s PRs already complete", self.owner, self.name) + return 0 + total_new = 0 + page = 0 + # Heavy nested PR query is capped at 3 per page (GitHub node-count + # ceiling); light query skips reviewThreads/reviews/commits/etc and + # can safely go to 25 per page. Clamp by trial_limit for small + # previews so limit=1 does not fetch a whole 25-item page. + page_cap = 25 if self.light else 3 + trial_cap = self.trial_limits.get(key) + per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } + query = Q.PRS_PAGE_QUERY_LIGHT if self.light else Q.PRS_PAGE_QUERY + data = self.client.graphql(query, vars_) + self._log_rate("prs", data) + repo = (data.get("data") or {}).get("repository") or {} + prs = repo.get("pullRequests") or {} + nodes = prs.get("nodes") or [] + for pr in nodes: + pr["_owner"] = self.owner + pr["_repo"] = self.name + pr["_fetchedAt"] = ts() + num = pr["number"] + if not self.light: + if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_comments( + num, pr["comments"]["pageInfo"]["endCursor"] + ) + if ( + pr.get("timelineItems", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): + self._paginate_pr_timeline( + num, pr["timelineItems"]["pageInfo"]["endCursor"] + ) + if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_commits( + num, pr["commits"]["pageInfo"]["endCursor"] + ) + if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_files( + num, pr["files"]["pageInfo"]["endCursor"] + ) + if ( + pr.get("reviewThreads", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): + self._paginate_pr_review_threads( + num, pr["reviewThreads"]["pageInfo"]["endCursor"] + ) + if self.writers[key].write(pr): + total_new += 1 + info = prs.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + log.info("Trial limit reached for PRs (%d)", total_new) + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + def _paginate_pr_comments(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "issueOrPullRequest" + ) or {} + comments = item.get("comments") or {} + for c in comments.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_prNumber"] = number + self.writers["pr_extra_comments"].write(c) + info = comments.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_timeline(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + tl = item.get("timelineItems") or {} + for ev in tl.get("nodes") or []: + ev["_owner"] = self.owner + ev["_repo"] = self.name + ev["_prNumber"] = number + self.writers["pr_extra_timeline"].write(ev) + info = tl.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_commits(self, number: int, after: str) -> None: + cur = after + out_key = "pr_extra_commits" + if out_key not in self.writers: + self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + cc = item.get("commits") or {} + for c in cc.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_prNumber"] = number + self.writers[out_key].write(c) + info = cc.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_files(self, number: int, after: str) -> None: + cur = after + out_key = "pr_extra_files" + if out_key not in self.writers: + self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_FILES_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + ff = item.get("files") or {} + for f in ff.get("nodes") or []: + f["_owner"] = self.owner + f["_repo"] = self.name + f["_prNumber"] = number + # files don't have id, synthesize one + f["_syntheticId"] = f"{self.owner}/{self.name}#{number}:{f.get('path')}" + self.writers[out_key].write(f) + info = ff.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_review_threads(self, number: int, after: str) -> None: + cur = after + out_key = "pr_extra_review_threads" + if out_key not in self.writers: + self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + rt = item.get("reviewThreads") or {} + for th in rt.get("nodes") or []: + th["_owner"] = self.owner + th["_repo"] = self.name + th["_prNumber"] = number + self.writers[out_key].write(th) + info = rt.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + # ----- Discussions ----- + def scrape_discussions(self) -> int: + key = "discussions" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + log.info("%s/%s discussions already complete", self.owner, self.name) + return 0 + total_new = 0 + page = 0 + per_page = 15 + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } + data = self.client.graphql(Q.DISCUSSIONS_PAGE_QUERY, vars_) + self._log_rate("discussions", data) + repo = (data.get("data") or {}).get("repository") or {} + dd = repo.get("discussions") or {} + nodes = dd.get("nodes") or [] + for d in nodes: + d["_owner"] = self.owner + d["_repo"] = self.name + d["_fetchedAt"] = ts() + num = d["number"] + if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_discussion_comments( + num, d["comments"]["pageInfo"]["endCursor"] + ) + # paginate replies per comment if needed + for c in d.get("comments", {}).get("nodes", []) or []: + if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_discussion_replies( + c["id"], c["replies"]["pageInfo"]["endCursor"], num + ) + if self.writers[key].write(d): + total_new += 1 + info = dd.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + def _paginate_discussion_comments(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_) + disc = ((data.get("data") or {}).get("repository") or {}).get( + "discussion" + ) or {} + cc = disc.get("comments") or {} + for c in cc.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_discussionNumber"] = number + self.writers["discussion_extra_comments"].write(c) + info = cc.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_discussion_replies( + self, comment_id: str, after: str, disc_number: int + ) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "commentId": comment_id, + "after": cur, + } + data = self.client.graphql(Q.DISCUSSION_REPLIES_QUERY, vars_) + node = (data.get("data") or {}).get("node") or {} + replies = node.get("replies") or {} + for r in replies.get("nodes") or []: + r["_owner"] = self.owner + r["_repo"] = self.name + r["_discussionNumber"] = disc_number + r["_commentId"] = comment_id + self.writers["discussion_extra_replies"].write(r) + info = replies.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + # ----- Commits ----- + def scrape_commits(self, branch: str = "refs/heads/main") -> int: + key = "commits" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + return 0 + total_new = 0 + page = 0 + page_cap = 100 + trial_cap = self.trial_limits.get(key) + per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + "branch": branch, + } + data = self.client.graphql(Q.COMMITS_PAGE_QUERY, vars_) + self._log_rate("commits", data) + ref = ((data.get("data") or {}).get("repository") or {}).get("ref") or {} + tgt = ref.get("target") or {} + hist = tgt.get("history") or {} + nodes = hist.get("nodes") or [] + for c in nodes: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_fetchedAt"] = ts() + if self.writers[key].write(c): + total_new += 1 + info = hist.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] commits page %d (+%d) remaining=%s", + self.owner, + self.name, + page, + len(nodes), + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + # ----- Releases/Labels/Milestones ----- + def scrape_releases(self) -> int: + return self._scrape_simple("releases", Q.RELEASES_QUERY, "releases") + + def scrape_labels(self) -> int: + return self._scrape_simple("labels", Q.LABELS_QUERY, "labels") + + def scrape_milestones(self) -> int: + return self._scrape_simple("milestones", Q.MILESTONES_QUERY, "milestones") + + def _scrape_simple(self, key: str, query: str, field: str) -> int: + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + return 0 + total_new = 0 + while True: + vars_ = { + "owner": self.owner, + "name": self.name, + "first": 50, + "after": cursor, + } + data = self.client.graphql(query, vars_) + repo = (data.get("data") or {}).get("repository") or {} + col = repo.get(field) or {} + for it in col.get("nodes") or []: + it["_owner"] = self.owner + it["_repo"] = self.name + it["_fetchedAt"] = ts() + if self.writers[key].write(it): + total_new += 1 + info = col.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + if self._trial_stop(key, total_new): + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + log.info("[%s/%s] %s done +%d", self.owner, self.name, key, total_new) + return total_new + + def close(self) -> None: + for w in self.writers.values(): + try: + w.close() + except Exception: + pass + + +def setup_logging(log_file: Path) -> None: + log_file.parent.mkdir(parents = True, exist_ok = True) + fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s" + handlers = [ + logging.StreamHandler(sys.stdout), + logging.FileHandler(log_file, mode = "a", encoding = "utf-8"), + ] + logging.basicConfig(level = logging.INFO, format = fmt, handlers = handlers, force = True) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper" + ) + ap.add_argument( + "--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"] + ) + ap.add_argument("--trial", action = "store_true", help = "Small trial run") + ap.add_argument( + "--only", + nargs = "+", + default = None, + help = "Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta", + ) + ap.add_argument( + "--hf-upload-interval", + type = int, + default = 900, + help = "Seconds between HF uploads (0 to disable)", + ) + args = ap.parse_args() + + base = Path(args.base_dir) + data_dir = base / "data" + data_dir.mkdir(parents = True, exist_ok = True) + setup_logging(base / "logs" / f"scraper_{time.strftime('%Y%m%d_%H%M%S')}.log") + log.info("Scraper starting: repos=%s trial=%s", args.repos, args.trial) + + client = GitHubClient(min_remaining_graphql = 80, min_remaining_rest = 80) + rl = client.rate_snapshot() + log.info( + "Rate limit snapshot: %s", + json.dumps(rl.get("resources", {}), default = str)[:400], + ) + + # Start HF uploader in background if requested + uploader = None + if args.hf_upload_interval > 0: + from hf_uploader import HFUploader + + uploader = HFUploader(data_dir, interval_s = args.hf_upload_interval) + uploader.start() + + trial_limits = None + if args.trial: + trial_limits = { + "issues": 5, + "pull_requests": 5, + "discussions": 3, + "commits": 20, + "releases": 3, + "labels": 20, + "milestones": 20, + } + + only = set(args.only or []) + + try: + for repo_spec in args.repos: + owner, name = repo_spec.split("/") + scraper = RepoScraper(owner, name, data_dir, client, trial_limits) + try: + repo_meta: Dict[str, Any] = {} + if not only or "meta" in only or "commits" in only: + repo_meta = scraper.scrape_repo_meta() + if not only or "labels" in only: + scraper.scrape_labels() + if not only or "milestones" in only: + scraper.scrape_milestones() + if not only or "releases" in only: + scraper.scrape_releases() + if not only or "discussions" in only: + scraper.scrape_discussions() + if not only or "issues" in only: + scraper.scrape_issues() + if not only or "pulls" in only: + scraper.scrape_prs() + if not only or "commits" in only: + default_ref = repo_meta.get("defaultBranchRef") or {} + default_branch = ( + default_ref.get("name") + if isinstance(default_ref, dict) + else None + ) + branch = ( + f"refs/heads/{default_branch}" + if default_branch + else "refs/heads/main" + ) + scraper.scrape_commits(branch = branch) + finally: + scraper.close() + finally: + if uploader: + log.info("Stopping uploader and final sync...") + uploader.stop(final_upload = True) + log.info( + "Scraper complete. GraphQL calls=%d REST calls=%d", + client.calls_graphql, + client.calls_rest, + ) + + +if __name__ == "__main__": + main() diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py new file mode 100644 index 0000000000..efa663db2f --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Checkpoint state management for resumable scraping.""" + +from __future__ import annotations + +import json +import os +import threading +from pathlib import Path +from typing import Any, Dict + + +class StateStore: + def __init__(self, path: str | Path): + self.path = Path(path) + self.path.parent.mkdir(parents = True, exist_ok = True) + self._lock = threading.Lock() + self._data: Dict[str, Any] = {} + if self.path.exists(): + try: + with self.path.open() as f: + self._data = json.load(f) + except Exception: + self._data = {} + + def get(self, key: str, default: Any = None) -> Any: + with self._lock: + return self._data.get(key, default) + + def set(self, key: str, value: Any) -> None: + with self._lock: + self._data[key] = value + self._flush() + + def update(self, key: str, **kwargs) -> None: + with self._lock: + sub = dict(self._data.get(key, {})) + sub.update(kwargs) + self._data[key] = sub + self._flush() + + def all(self) -> Dict[str, Any]: + with self._lock: + return dict(self._data) + + def _flush(self) -> None: + tmp = self.path.with_suffix(self.path.suffix + ".tmp") + with tmp.open("w") as f: + json.dump(self._data, f, indent = 2, default = str) + os.replace(tmp, self.path) + + +class JsonlWriter: + """Append-only JSONL writer, thread-safe, with line buffering.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + self.path.parent.mkdir(parents = True, exist_ok = True) + self._lock = threading.Lock() + self._fh = self.path.open("a", buffering = 1) + self._count_seen_keys: set[str] = set() + # Preload seen keys if file exists (for dedup across resumes) + if self.path.exists() and self.path.stat().st_size > 0: + try: + with self.path.open() as f: + for line in f: + try: + obj = json.loads(line) + k = self._key(obj) + if k is not None: + self._count_seen_keys.add(k) + except Exception: + pass + except Exception: + pass + + def _key(self, obj: dict) -> str | None: + for k in ("id", "node_id", "number", "sha", "url"): + if k in obj: + return f"{k}:{obj[k]}" + return None + + def has(self, key: str) -> bool: + return key in self._count_seen_keys + + def write(self, obj: dict) -> bool: + """Return True if newly written, False if already present.""" + k = self._key(obj) + with self._lock: + if k is not None and k in self._count_seen_keys: + return False + if k is not None: + self._count_seen_keys.add(k) + self._fh.write(json.dumps(obj, default = str, ensure_ascii = False)) + self._fh.write("\n") + self._fh.flush() + return True + + def close(self) -> None: + try: + self._fh.close() + except Exception: + pass diff --git a/studio/backend/requirements/single-env/data-designer-deps.txt b/studio/backend/requirements/single-env/data-designer-deps.txt index fc63230922..f63c076621 100644 --- a/studio/backend/requirements/single-env/data-designer-deps.txt +++ b/studio/backend/requirements/single-env/data-designer-deps.txt @@ -19,7 +19,8 @@ ruff<1,>=0.14.10 scipy<2,>=1.11.0 sqlfluff<4,>=3.2.0 tiktoken<1,>=0.8.0 -# Unstructured-seed plugin deps (plugin installed with --no-deps) +# Local seed plugin deps (plugins installed with --no-deps) +requests>=2.31 pymupdf>=1.24.0 pymupdf4llm>=0.0.17 mammoth>=1.8.0 diff --git a/studio/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index 606ef1832c..da6416e324 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -5,8 +5,9 @@ from __future__ import annotations -from datetime import timedelta -from typing import Any +import copy +from datetime import datetime, timedelta, timezone +from typing import Any, Optional from urllib.parse import urlparse from fastapi import APIRouter, HTTPException, Query, Request @@ -94,14 +95,111 @@ def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]: return aliases -def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: +def _inject_local_structured_response_format( + recipe: dict[str, Any], local_provider_names: set[str] +) -> None: + """For each llm-structured column that targets a local-provider model_config, + clone the model_config and inject an OpenAI ``response_format`` with the + column's ``output_format`` JSON schema. The column is rewritten to point at + the clone so llm-text / llm-judge columns that share the same alias keep + free-form sampling. + + Without this, data_designer only injects a prompt-level "return JSON in a + ```json fence" instruction. Small GGUF models frequently break format, + wasting the full ``max_tokens`` budget per row and then failing to parse. + Forwarding ``response_format`` lets llama-server apply grammar-constrained + sampling from the JSON schema, which guarantees a parseable response and + terminates early. + """ + columns = recipe.get("columns") + model_configs = recipe.get("model_configs") + if not isinstance(columns, list) or not isinstance(model_configs, list): + return + + # alias -> model_config (only configs referencing a local provider qualify). + alias_to_local_mc: dict[str, dict[str, Any]] = {} + for mc in model_configs: + if not isinstance(mc, dict): + continue + if mc.get("provider") in local_provider_names and isinstance( + mc.get("alias"), str + ): + alias_to_local_mc[mc["alias"]] = mc + + if not alias_to_local_mc: + return + + # Clone per (alias, column) so each llm-structured column gets its own + # schema without leaking response_format onto other columns that share the + # same base alias. + seen_clone_aliases: set[str] = { + mc.get("alias") for mc in model_configs if isinstance(mc.get("alias"), str) + } + new_configs: list[dict[str, Any]] = [] + for column in columns: + if not isinstance(column, dict): + continue + if column.get("column_type") != "llm-structured": + continue + alias = column.get("model_alias") + if not isinstance(alias, str) or alias not in alias_to_local_mc: + continue + output_format = column.get("output_format") + if not isinstance(output_format, dict) or not output_format: + continue + base_mc = alias_to_local_mc[alias] + column_name = column.get("name") or "structured" + clone_alias_base = f"{alias}__{column_name}_structured" + clone_alias = clone_alias_base + counter = 1 + while clone_alias in seen_clone_aliases: + counter += 1 + clone_alias = f"{clone_alias_base}_{counter}" + seen_clone_aliases.add(clone_alias) + + clone = copy.deepcopy(base_mc) + clone["alias"] = clone_alias + params = clone.get("inference_parameters") + if not isinstance(params, dict): + params = {} + clone["inference_parameters"] = params + # data_designer's BaseInferenceParams is a pydantic model with + # extra="forbid", so response_format cannot sit at the top level of + # inference_parameters. It does expose an `extra_body: dict` pass- + # through that the OpenAI client spreads into the request body at the + # top level, which is where llama-server reads response_format from. + # llama.cpp server shape (tools/server/README.md): the schema sits + # directly under response_format, not nested in a json_schema object + # the way OpenAI's Chat Completions API expects. llama-server converts + # the schema to a GBNF grammar and applies it during sampling. + extra_body = params.get("extra_body") + if not isinstance(extra_body, dict): + extra_body = {} + extra_body["response_format"] = { + "type": "json_schema", + "schema": output_format, + } + params["extra_body"] = extra_body + new_configs.append(clone) + column["model_alias"] = clone_alias + + if new_configs: + model_configs.extend(new_configs) + + +def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optional[int]: """ Mutate recipe dict in-place: for any provider with is_local=True, - generate a JWT and fill in the endpoint pointing at this server. + fill in the endpoint pointing at this server and inject a short-lived + internal sk-unsloth-* API key for workflow auth. + + Returns the row id of the minted internal key (so the caller can + revoke it on job completion) or ``None`` when no local provider is + actually reachable from an LLM column. """ providers = recipe.get("model_providers") if not providers: - return + return None # Collect local providers and pop is_local from ALL dicts unconditionally. # Strict `is True` guard so malformed payloads (is_local: 1, @@ -115,7 +213,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: local_indices.append(i) if not local_indices: - return + return None endpoint = _resolve_local_v1_endpoint(request) @@ -138,6 +236,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: } token = "" + internal_key_id: Optional[int] = None if local_names & referenced_providers: # Verify a model is loaded. # NOTE: This is a point-in-time check (TOCTOU). The model could be unloaded @@ -158,18 +257,21 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: "No model loaded in Chat. Load a model first, then run the recipe." ) - from auth.authentication import ( - create_access_token, - ) # deferred: avoids circular import + from auth import storage # deferred: avoids circular import - # Uses the "unsloth" admin subject. If the user changes their password, - # the JWT secret rotates and this token becomes invalid mid-run. - # Acceptable for v1 - recipes typically finish well within one session. - token = create_access_token( - subject = "unsloth", - expires_delta = timedelta(hours = 24), - desktop = _request_has_desktop_access_token(request), + # Mint an internal sk-unsloth-* key scoped to this workflow run. + # Uses the unified API-key issuance path (one mint/revoke/verify + # surface instead of a second JWT code path). The key is marked + # internal so it is hidden from the user's API-key list, and the + # caller revokes it when the job terminates. + expires_at = (datetime.now(timezone.utc) + timedelta(hours = 24)).isoformat() + token, row = storage.create_api_key( + username = "unsloth", + name = "data-recipe workflow", + expires_at = expires_at, + internal = True, ) + internal_key_id = int(row["id"]) # Defensively strip any stale "external"-only fields the frontend may # have left on the dict (extra_headers/extra_body/api_key_env). The UI @@ -196,6 +298,37 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: continue if mc.get("provider") in local_names: mc["skip_health_check"] = True + # Disable thinking for data-recipe inference on local providers. + # Reasoning models emit a ... preamble before the + # answer, which roughly doubles generated token count per row and + # pushes the visible answer past data_designer's json-fence + # regex. Forward chat_template_kwargs={enable_thinking: False} + # through the OpenAI SDK's extra_body passthrough so llama-server + # renders the template without the reasoning preamble. Free-form + # llm-text columns benefit from the latency cut, and structured + # columns also stop leaking think tags into the grammar- + # constrained JSON (llama-server's GBNF path still enforces the + # schema either way). + params = mc.get("inference_parameters") + if not isinstance(params, dict): + params = {} + mc["inference_parameters"] = params + extra_body = params.get("extra_body") + if not isinstance(extra_body, dict): + extra_body = {} + tpl_kwargs = extra_body.get("chat_template_kwargs") + if not isinstance(tpl_kwargs, dict): + tpl_kwargs = {} + tpl_kwargs.setdefault("enable_thinking", False) + extra_body["chat_template_kwargs"] = tpl_kwargs + params["extra_body"] = extra_body + + # Forward each llm-structured column's output_format as an OpenAI + # response_format so llama-server uses grammar-constrained sampling and + # small GGUFs stop wasting the full max_tokens budget on broken JSON. + _inject_local_structured_response_format(recipe, local_names) + + return internal_key_id def _normalize_run_name(value: Any) -> str | None: @@ -240,21 +373,49 @@ def create_job(payload: RecipePayload, request: Request): ) from exc try: - _inject_local_providers(recipe, request) + internal_api_key_id = _inject_local_providers(recipe, request) except ValueError as exc: raise HTTPException(status_code = 400, detail = str(exc)) from exc - mgr = get_job_manager() + # Single try block covers get_job_manager() AND mgr.start() so a workflow + # key minted above never outlives the request even when an unexpected + # exception type (TypeError from a stale kwarg, OSError from a queue + # write, etc.) bubbles up. Without the bare except, such exceptions let + # the sk-unsloth-* key live until its 24h TTL. try: - job_id = mgr.start(recipe = recipe, run = run) + mgr = get_job_manager() + job_id = mgr.start( + recipe = recipe, + run = run, + internal_api_key_id = internal_api_key_id, + ) except RuntimeError as exc: + if internal_api_key_id is not None: + _revoke_internal_api_key_safe(internal_api_key_id) raise HTTPException(status_code = 409, detail = str(exc)) from exc except ValueError as exc: + if internal_api_key_id is not None: + _revoke_internal_api_key_safe(internal_api_key_id) raise HTTPException(status_code = 400, detail = str(exc)) from exc + except Exception: + if internal_api_key_id is not None: + _revoke_internal_api_key_safe(internal_api_key_id) + raise return {"job_id": job_id} +def _revoke_internal_api_key_safe(key_id: int) -> None: + """Best-effort revoke of a workflow-minted key; swallow any error so + that revocation failures never mask the caller's own error path.""" + try: + from auth import storage # deferred: avoids circular import + + storage.revoke_internal_api_key(key_id) + except Exception: + pass + + @router.get("/jobs/{job_id}/status") def job_status(job_id: str): mgr = get_job_manager() diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index e9cf828610..91cf718e6e 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -8,6 +8,7 @@ from __future__ import annotations import base64 import binascii import json +import os import re from itertools import islice from pathlib import Path @@ -627,3 +628,14 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons split = None, subset = None, ) + + +@router.get("/seed/github/env-token") +def get_github_env_token_status() -> dict: + """Report whether the server has a GH_TOKEN / GITHUB_TOKEN env var. + + The value is never returned; the UI uses this to tell the user they + can leave the token field blank. + """ + has_token = bool(os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")) + return {"has_token": has_token} diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py index 555e3eaa06..87eef939b4 100644 --- a/studio/backend/routes/data_recipe/validate.py +++ b/studio/backend/routes/data_recipe/validate.py @@ -18,6 +18,57 @@ from models.data_recipe import RecipePayload, ValidateError, ValidateResponse router = APIRouter() +_GITHUB_VALIDATE_NOTE = "Recipe shape is valid. GitHub access and rate limits are checked when the run starts." +_GITHUB_ITEM_TYPES = {"issues", "pulls", "commits"} + + +def _github_seed_source(recipe: dict[str, Any]) -> dict[str, Any] | None: + seed_config = recipe.get("seed_config") + if not isinstance(seed_config, dict): + return None + source = seed_config.get("source") + if not isinstance(source, dict) or source.get("seed_type") != "github_repo": + return None + return source + + +def _validate_github_seed_static(source: dict[str, Any]) -> list[ValidateError]: + errors: list[ValidateError] = [] + + repos = source.get("repos") + if not isinstance(repos, list) or not repos: + errors.append(ValidateError(message = "GitHub seed requires at least one repo.")) + else: + for repo in repos: + if not isinstance(repo, str) or not repo.strip() or "/" not in repo: + errors.append( + ValidateError(message = "GitHub repos must be owner/name strings.") + ) + break + + item_types = source.get("item_types") + if not isinstance(item_types, list) or not item_types: + errors.append( + ValidateError(message = "GitHub seed requires at least one item type.") + ) + else: + invalid_items = [item for item in item_types if item not in _GITHUB_ITEM_TYPES] + if invalid_items: + errors.append( + ValidateError( + message = "GitHub item types must be issues, pulls, or commits." + ) + ) + + try: + limit = int(source.get("limit")) + except (TypeError, ValueError): + limit = 0 + if limit < 1 or limit > 5000: + errors.append(ValidateError(message = "GitHub limit must be from 1 to 5000.")) + + return errors + def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]: try: @@ -93,6 +144,22 @@ def validate(payload: RecipePayload) -> ValidateResponse: _patch_local_providers(recipe) + github_source = _github_seed_source(recipe) + if github_source is not None: + static_errors = _validate_github_seed_static(github_source) + if static_errors: + return ValidateResponse(valid = False, errors = static_errors) + try: + build_config_builder(recipe) + except Exception as exc: + detail = str(exc).strip() or "Validation failed." + return ValidateResponse( + valid = False, + errors = [ValidateError(message = detail)], + raw_detail = detail, + ) + return ValidateResponse(valid = True, raw_detail = _GITHUB_VALIDATE_NOTE) + try: validate_recipe(recipe) except RuntimeError as exc: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index cf3b37a2fd..b13eb08967 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1286,6 +1286,20 @@ async def openai_chat_completions( llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded + # OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, + # which the SDK spreads into the request body at the top level. Studio's + # ChatCompletionRequest has ``extra="allow"`` so pydantic stashes them in + # ``model_extra``, but the typed ``payload.enable_thinking`` path is what + # downstream generators actually consume. Lift ``enable_thinking`` from + # the extra-body chat_template_kwargs onto the typed field so clients + # that only know the OpenAI shape (data_designer recipe runs, etc.) + # can still control the reasoning preamble. + _extra = getattr(payload, "model_extra", None) + if payload.enable_thinking is None and isinstance(_extra, dict): + _tpl_kw = _extra.get("chat_template_kwargs") + if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: + payload.enable_thinking = bool(_tpl_kw["enable_thinking"]) + # ── Determine which backend is active ───────────────────── if using_gguf: model_name = llama_backend.model_identifier or payload.model @@ -1440,11 +1454,22 @@ async def openai_chat_completions( # carry `tool_calls` (content=None) — both of which are valid in # multi-turn client-side tool loops. _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) + # Route guided-decoding requests through the verbatim passthrough so + # ``response_format`` (JSON schema) actually reaches llama-server and + # the model's GBNF-constrained output comes back unmodified. The + # non-passthrough GGUF path below calls ``generate_chat_completion`` + # which has no response_format kwarg, so the schema gets silently + # dropped and data_designer falls back to free-form sampling. Guided + # decoding does not require ``supports_tools`` - the grammar machinery + # is independent of tool-call parsing. + _has_response_format = _extract_response_format(payload) is not None + _tools_passthrough = llama_backend.supports_tools and ( + (payload.tools and len(payload.tools) > 0) or _has_tool_messages + ) if ( using_gguf - and llama_backend.supports_tools and not payload.enable_tools - and ((payload.tools and len(payload.tools) > 0) or _has_tool_messages) + and (_tools_passthrough or _has_response_format) ): # Preserve the vision guard that would otherwise run in the # non-passthrough path below: text-only tool-capable GGUFs @@ -3652,6 +3677,8 @@ def _build_passthrough_payload( repetition_penalty = None, presence_penalty = None, tool_choice = "auto", + response_format = None, + chat_template_kwargs = None, backend_ctx = None, ): body = { @@ -3680,6 +3707,17 @@ def _build_passthrough_payload( body["repeat_penalty"] = repetition_penalty if presence_penalty is not None: body["presence_penalty"] = presence_penalty + if response_format is not None: + # llama-server applies a GBNF grammar derived from the JSON schema + # when response_format is present. Field is documented flat at the + # request root (tools/server/README.md), which is also what the + # OpenAI SDK produces by spreading extra_body into the body top. + body["response_format"] = response_format + if chat_template_kwargs is not None: + # Propagate reasoning / template overrides (e.g. enable_thinking) + # so llama-server renders the Jinja template in the mode the caller + # asked for instead of whatever default the model was loaded with. + body["chat_template_kwargs"] = chat_template_kwargs return body @@ -3990,6 +4028,20 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages +def _extract_response_format(payload): + """Return the ``response_format`` field on an incoming ChatCompletionRequest + (or None). The model is declared with ``extra="allow"`` so pydantic stashes + unknown top-level fields in ``model_extra``; OpenAI-SDK clients spread + ``extra_body`` into the request body top level, which is where guided- + decoding recipes park their JSON-schema response_format. + """ + extra = getattr(payload, "model_extra", None) + if not isinstance(extra, dict): + return None + rf = extra.get("response_format") + return rf if isinstance(rf, dict) else None + + def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: """Assemble the llama-server request body from a ChatCompletionRequest. @@ -3999,6 +4051,12 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: """ messages = _openai_messages_for_passthrough(payload) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" + # When the caller asked for a specific reasoning mode, forward it to + # llama-server via chat_template_kwargs so the Jinja template renders + # with (or without) the reasoning preamble. + tpl_kwargs = None + if payload.enable_thinking is not None: + tpl_kwargs = {"enable_thinking": bool(payload.enable_thinking)} return _build_passthrough_payload( messages, payload.tools, @@ -4012,6 +4070,8 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: repetition_penalty = payload.repetition_penalty, presence_penalty = payload.presence_penalty, tool_choice = tool_choice, + response_format = _extract_response_format(payload), + chat_template_kwargs = tpl_kwargs, backend_ctx = backend_ctx, ) @@ -4214,6 +4274,41 @@ async def _openai_passthrough_non_streaming( detail = f"llama-server error: {resp.text[:500]}", ) + # Guided-decoding fence wrap. llama-server returns raw JSON that matches + # the schema (no surrounding markdown) because the GBNF grammar only + # emits the JSON object itself. data_designer's llm-structured parser + # looks for a ```json ... ``` markdown fence and discards unfenced + # output, which collapses a 100%-valid guided-decoding run to 0/N. + # Wrap each choice's content in the expected fence when the caller + # asked for guided decoding, leaving already-fenced content alone. + if _extract_response_format(payload) is not None: + try: + data = resp.json() + changed = False + for choice in data.get("choices", []): + if not isinstance(choice, dict): + continue + msg = choice.get("message") + if not isinstance(msg, dict): + continue + content = msg.get("content") + if not isinstance(content, str): + continue + stripped = content.strip() + if not stripped or stripped.startswith("```"): + continue + msg["content"] = f"```json\n{stripped}\n```" + changed = True + if changed: + return JSONResponse(content = data) + except Exception as exc: + # Wrap is best-effort; fall through to the verbatim body if + # the response is not JSON-shaped or the structure is unusual. + logger.warning( + "response_format fence wrap skipped: %s", + exc, + ) + # Pass the upstream body through as raw bytes — skips a redundant # parse+re-serialize round-trip and keeps the response truly # verbatim (matches the docstring). Status is guaranteed 200 by diff --git a/studio/backend/tests/test_data_recipe_github_progress.py b/studio/backend/tests/test_data_recipe_github_progress.py new file mode 100644 index 0000000000..8e8c3995f4 --- /dev/null +++ b/studio/backend/tests/test_data_recipe_github_progress.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from core.data_recipe.jobs.parse import apply_update, parse_log_message +from core.data_recipe.jobs.types import Job +from routes.data_recipe.validate import _GITHUB_VALIDATE_NOTE, validate +from models.data_recipe import RecipePayload + + +def test_github_page_log_updates_source_progress_without_cursor(): + job = Job(job_id = "job-1") + job.source_progress_estimated_total = 200 + + update = parse_log_message( + "[unslothai/unsloth] issues page 2 (+15) cursor=abc123 remaining=2960" + ) + + assert update is not None + apply_update(job, update) + + progress = job.source_progress + assert progress is not None + assert progress.source == "github" + assert progress.status == "fetching" + assert progress.repo == "unslothai/unsloth" + assert progress.resource == "issues" + assert progress.page == 2 + assert progress.page_items == 15 + assert progress.fetched_items == 15 + assert progress.estimated_total == 200 + assert progress.rate_remaining == 2960 + assert progress.message is not None + assert "cursor" not in progress.message + assert "abc123" not in progress.message + + +def test_github_rate_limit_log_updates_source_progress(): + job = Job(job_id = "job-1") + + update = parse_log_message("Rate limit hit. Sleeping 123s until reset.") + + assert update is not None + apply_update(job, update) + + progress = job.source_progress + assert progress is not None + assert progress.status == "rate_limited" + assert progress.retry_after_sec == 123 + assert "resume automatically" in (progress.message or "") + + +def test_github_real_sample_prs_and_trial_limit_are_parsed(): + job = Job(job_id = "job-1") + + for message in ( + "[unslothai/unsloth] PRs page 4 (+25) cursor=abc123 remaining=4983", + "Trial limit reached for PRs (100)", + ): + update = parse_log_message(message) + assert update is not None + apply_update(job, update) + + progress = job.source_progress + assert progress is not None + assert progress.repo == "unslothai/unsloth" + assert progress.resource == "pulls" + assert progress.page == 4 + assert progress.fetched_items == 25 + assert progress.rate_remaining == 4983 + assert progress.message == "GitHub pulls trial limit reached (100)." + + +def test_github_validate_skips_live_access_with_honest_note(): + response = validate( + RecipePayload( + recipe = { + "seed_config": { + "source": { + "seed_type": "github_repo", + "repos": ["unslothai/unsloth"], + "item_types": ["issues"], + "limit": 1, + } + }, + "columns": [{"column_type": "expression", "name": "x", "expr": "1"}], + } + ) + ) + + assert response.valid is True + assert response.raw_detail == _GITHUB_VALIDATE_NOTE diff --git a/studio/frontend/src/features/data-recipes/learning-recipes/github-support-bot.json b/studio/frontend/src/features/data-recipes/learning-recipes/github-support-bot.json new file mode 100644 index 0000000000..b984f84d2f --- /dev/null +++ b/studio/frontend/src/features/data-recipes/learning-recipes/github-support-bot.json @@ -0,0 +1,238 @@ +{ + "recipe": { + "model_providers": [ + { + "name": "Local Model", + "endpoint": "", + "provider_type": "openai", + "extra_headers": {}, + "extra_body": {}, + "is_local": true + } + ], + "mcp_providers": [], + "model_configs": [ + { + "alias": "model_1", + "model": "unsloth/gemma-4-E2B-it-GGUF", + "provider": "Local Model", + "inference_parameters": { + "temperature": 0.4, + "max_tokens": 800 + } + } + ], + "seed_config": { + "source": { + "seed_type": "github_repo", + "repos": [ + "unslothai/unsloth", + "unslothai/unsloth-zoo" + ], + "item_types": [ + "issues", + "pulls" + ], + "limit": 100, + "include_comments": true, + "max_comments_per_item": 20 + }, + "sampling_strategy": "shuffle", + "selection_strategy": null + }, + "tool_configs": [], + "columns": [ + { + "column_type": "llm-text", + "name": "User", + "drop": false, + "model_alias": "model_1", + "prompt": "Read the GitHub {{ item_type }} below and write ONE realistic user request that could have produced it. Imagine a developer asking a GitHub co-author model to either file this {{ item_type }} or draft a PR that resolves it. Use first-person imperative phrasing (\"Open an issue...\", \"Draft a PR that...\", \"Investigate why...\"). Preserve concrete technical details (model names, flags, file paths, tracebacks) that appear in the thread. Keep it 1-3 sentences. Output ONLY the user request, no preamble.\n\n--- INPUT ---\nRepo: {{ repo }}\nType: {{ item_type }}\nTitle: {{ title }}\nBody:\n{{ body }}\n\nFirst comments:\n{{ comments }}", + "system_prompt": "You invert real GitHub threads into the user request that would have produced them. Faithful to the thread, no invented facts, no em-dashes, no emojis.", + "with_trace": "none", + "extract_reasoning_content": false + }, + { + "column_type": "llm-structured", + "name": "Assistant", + "drop": false, + "model_alias": "model_1", + "prompt": "You are generating one training row for an Unsloth GitHub co-author model. Given the real GitHub thread and a synthesized user request, produce a grounded structured response.\n\nSource thread:\n- Repo: {{ repo }}\n- Type: {{ item_type }}\n- Title: {{ title }}\n- URL: {{ url }}\n- State: {{ state }}\n- Labels: {{ labels }}\n- Body: {{ body }}\n- First comments: {{ comments }}\n\nUser request:\n{{ User }}\n\nRules:\n- `response`: 100-250 words of Markdown grounded in the thread. If the thread is a closed / resolved issue, follow the `issue_fix_plan` shape: brief diagnosis, numbered fix steps, and a short repro. If the thread is a PR, follow the `explain_pr` shape: what changed, why, and which files or symbols were touched. If the thread is open / unresolved, answer honestly and ask for the missing info.\n- Cite the source URL at least once inline as `[source: {{ url }}]`.\n- Name at least one concrete symbol (function, class, flag, env var, or file path) from the thread when available.\n- Include a short ```bash or ```python code block ONLY if the thread itself contains that code or command.\n- Never recommend `rm -rf`, force push, or other destructive commands without an explicit warning.\n- No em-dashes, no emojis, no AI-disclaimer phrases. Only cite URLs / paths that appear in the thread.\n- `followups`: 0-4 follow-up questions when the thread is missing info (versions, GPU, traceback). Empty list if the response is complete.\n- `cites`: URLs / file paths actually used. Always include `{{ url }}`.\n- `task`: one of `explain_pr`, `issue_fix_plan`, `issue_solution`, `discussion_qa`. Pick the closest match.\n- `confidence`: `high` / `medium` / `low`. Use `low` when ambiguous or out of scope.", + "system_prompt": "You write grounded GitHub co-author responses for Unsloth. Faithful to the thread, no invented facts, no em-dashes, no emojis, no AI-disclaimer phrases.", + "with_trace": "none", + "extract_reasoning_content": false, + "output_format": { + "type": "object", + "properties": { + "response": { + "type": "string", + "minLength": 1 + }, + "followups": { + "type": "array", + "items": { + "type": "string" + }, + "maxItems": 4 + }, + "cites": { + "type": "array", + "items": { + "type": "string" + }, + "maxItems": 6 + }, + "task": { + "type": "string", + "enum": [ + "explain_pr", + "issue_fix_plan", + "issue_solution", + "discussion_qa" + ] + }, + "confidence": { + "type": "string", + "enum": [ + "high", + "medium", + "low" + ] + } + }, + "required": [ + "response", + "followups", + "cites", + "task", + "confidence" + ], + "additionalProperties": false + } + } + ], + "processors": [] + }, + "run": { + "rows": 5, + "preview": true, + "output_formats": [ + "jsonl" + ] + }, + "ui": { + "nodes": [ + { + "id": "note_1", + "x": 460.4272962489311, + "y": 123.10114554082236, + "width": 400, + "node_type": "markdown_note", + "name": "note_1", + "markdown": "### GitHub Crawler\nReal GitHub issues and PRs turned into `{User, Assistant}` training pairs. Mirrors two of the eleven canonical enrichment tasks in the `github_data_gatherer` dataset: `pr_requests_20` / `issue_requests_20` for the input side, and `explain_pr` / `issue_fix_plan` / `issue_solution` for the output side.\n\n**Click `Run` below for 10 sample rows.** Defaults point at `unslothai/unsloth` + `unslothai/unsloth-zoo`, use the server's `GH_TOKEN` / `GITHUB_TOKEN` env var, and run the bundled local model. Paste a PAT only when you need private repos.\n\n**Configure source data**\n- Paste `owner/name` values or GitHub URLs; the editor normalizes and dedupes them.\n- Keep the seed `limit` around 100 for previews. Increase toward 5000 per item type for larger backfills.\n- Comments and large repos can make Check or Run take minutes; watch logs for GitHub page and rate-limit messages.\n\n**Upgrade to production**\n- Swap `unsloth/gemma-4-E2B-it-GGUF` for a stronger model (`gpt-5.4-mini` with `reasoning_effort=medium` is what the reference dataset uses).\n- Replace the demo prompts with the task-specific prompts from the reference dataset (see Note 3).\n- Raise `max_parallel_requests` to 4 once the inference server can handle it.", + "note_color": "#E0F2FE", + "note_opacity": "35" + }, + { + "id": "note_2", + "x": -514.4586648521598, + "y": 869.2785862193363, + "width": 400, + "node_type": "markdown_note", + "name": "note_2", + "markdown": "The **User** column inverts each GitHub thread into a realistic request a developer would give a co-author model (`\"Draft a PR that...\"`, `\"Investigate why...\"`). Same shape as the `pr_requests_20` and `issue_requests_20` enrichments.\n\nTweak the prompt to:\n- always keep the traceback verbatim\n- vary persona (newcomer, maintainer, ops)\n- split one thread into multiple alternative phrasings for data augmentation.", + "note_color": "#E0F2FE", + "note_opacity": "35" + }, + { + "id": "note_3", + "x": -536.3163554442993, + "y": -84.18780704666585, + "width": 400, + "node_type": "markdown_note", + "name": "note_3", + "markdown": "The **Assistant** block emits `{response, followups, cites, task, confidence}` and branches on thread type: closed issues become `issue_fix_plan` rows, PRs become `explain_pr` rows, everything else becomes `issue_solution` or `discussion_qa`.\n\n**Demo default**: 100-250 word response, one inline `[source: ]` cite, `max_parallel_requests=1` so a small local model stays stable.\n\n**Production prompt (paste in):**\n- Match the reference dataset's per-task prompts (`explain_pr`, `issue_fix_plan`, `pr_review_critique`, `pr_test_plan`, etc.).\n- Require 300 words for explanations, 6-12 bullets for test plans.\n- Enforce named symbols (function / flag / env var / file path).\n- Code fences only for content already in the thread.\n- Reject rows with em-dashes, emojis, or AI-disclaimer phrases.\n\nSee the `github_data_gatherer` dataset card for the full task catalog and the codex prompts used to train the reference GitHub model.", + "note_color": "#E0F2FE", + "note_opacity": "35" + }, + { + "id": "seed", + "x": 0, + "y": 140, + "width": 400 + }, + { + "id": "Local Model", + "x": -1056, + "y": 520, + "width": 400 + }, + { + "id": "model_1", + "x": -544, + "y": 488, + "width": 400 + }, + { + "id": "User", + "x": 0, + "y": 440, + "width": 400 + }, + { + "id": "Assistant", + "x": 0, + "y": 740, + "width": 400 + } + ], + "edges": [ + { + "from": "seed", + "to": "User", + "type": "canvas", + "source_handle": "data-out-bottom", + "target_handle": "data-in-top" + }, + { + "from": "User", + "to": "Assistant", + "type": "canvas", + "source_handle": "data-out-bottom", + "target_handle": "data-in-top" + }, + { + "from": "Local Model", + "to": "model_1", + "type": "semantic", + "source_handle": "semantic-out", + "target_handle": "semantic-in" + }, + { + "from": "model_1", + "to": "User", + "type": "semantic", + "source_handle": "semantic-out", + "target_handle": "data-in" + }, + { + "from": "model_1", + "to": "Assistant", + "type": "semantic", + "source_handle": "semantic-out-bottom", + "target_handle": "data-in" + } + ], + "layout_direction": "LR", + "seed_source_type": "github_repo", + "seed_columns": [], + "seed_drop_columns": [], + "seed_preview_rows": [], + "local_file_name": "", + "unstructured_file_ids": [], + "unstructured_file_names": [], + "unstructured_file_sizes": [], + "unstructured_chunk_size": "1200", + "unstructured_chunk_overlap": "200" + } +} \ No newline at end of file diff --git a/studio/frontend/src/features/data-recipes/learning-recipes/index.ts b/studio/frontend/src/features/data-recipes/learning-recipes/index.ts index d7a7e66e0a..8607d3faeb 100644 --- a/studio/frontend/src/features/data-recipes/learning-recipes/index.ts +++ b/studio/frontend/src/features/data-recipes/learning-recipes/index.ts @@ -19,6 +19,10 @@ const ocrDocumentExtractionUrl = new URL( "./ocr-document-extraction.json", import.meta.url, ).href; +const githubSupportBotUrl = new URL( + "./github-support-bot.json", + import.meta.url, +).href; function isRecord(value: unknown): value is Record { return !!value && typeof value === "object" && !Array.isArray(value); @@ -137,4 +141,11 @@ export const LEARNING_RECIPES: LearningRecipeDef[] = [ "Use image context to generate OCR-style document extraction output.", loadPayload: () => loadPayloadFromUrl(ocrDocumentExtractionUrl), }, + { + id: "github-support-bot", + title: "GitHub Crawler", + description: + "Crawl real GitHub issues and PRs and turn each thread into a {User, Assistant} training pair.", + loadPayload: () => loadPayloadFromUrl(githubSupportBotUrl), + }, ]; diff --git a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx index 9148b9e0da..0a51f9df70 100644 --- a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx +++ b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx @@ -35,6 +35,7 @@ import { Delete02Icon, DocumentAttachmentIcon, FunctionIcon, + GithubIcon, Plant01Icon, PlusSignIcon, } from "@hugeicons/core-free-icons"; @@ -162,6 +163,22 @@ const TEMPLATE_CARDS: TemplateCard[] = [ ], learningRecipeId: "structured-outputs-jinja", }, + { + title: "GitHub Crawler", + description: + "Crawl real GitHub issues and PRs and invert each thread into a {User, Assistant} training pair.", + icon: GithubIcon, + difficulty: "Intermediate", + learningBadges: ["GitHub", "LLM Text", "Structured LLM"], + surfaceClassName: + "from-slate-500/15 via-zinc-500/5 to-transparent dark:from-slate-400/30 dark:via-zinc-400/14 dark:to-slate-950/16", + shineColor: [ + "rgb(71 85 105 / 0.45)", + "rgb(100 116 139 / 0.4)", + "rgb(148 163 184 / 0.45)", + ], + learningRecipeId: "github-support-bot", + }, ]; const LEARNING_RECIPE_BY_ID = new Map( diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 9212e4db9b..273d4aea8d 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -27,6 +27,28 @@ export type PublishRecipeJobResponse = { message: string; }; +export type SourceProgressResponse = { + source?: string | null; + status?: string | null; + repo?: string | null; + resource?: string | null; + page?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + page_items?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + fetched_items?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + estimated_total?: number | null; + percent?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + rate_remaining?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + retry_after_sec?: number | null; + message?: string | null; + // biome-ignore lint/style/useNamingConvention: api schema + updated_at?: number | null; +}; + export type JobStatusResponse = { // biome-ignore lint/style/useNamingConvention: api schema job_id: string; @@ -62,6 +84,8 @@ export type JobStatusResponse = { failed?: number | null; }; // biome-ignore lint/style/useNamingConvention: api schema + source_progress?: SourceProgressResponse | null; + // biome-ignore lint/style/useNamingConvention: api schema model_usage?: Record; rows?: number | null; cols?: number | null; @@ -183,12 +207,7 @@ async function parseErrorResponse(response: Response): Promise { // biome-ignore lint/style/useNamingConvention: api schema raw_detail?: string; }; - return ( - parsed.detail ?? - parsed.message ?? - parsed.raw_detail ?? - text - ); + return parsed.detail ?? parsed.message ?? parsed.raw_detail ?? text; } catch { return text; } @@ -264,11 +283,15 @@ export async function validateRecipe( return postJson("/validate", payload); } -export async function createRecipeJob(payload: unknown): Promise { +export async function createRecipeJob( + payload: unknown, +): Promise { return postJson("/jobs", payload); } -export async function getRecipeJobStatus(jobId: string): Promise { +export async function getRecipeJobStatus( + jobId: string, +): Promise { return getJson(`/jobs/${jobId}/status`); } @@ -292,7 +315,9 @@ export async function getRecipeJobDataset( ); } -export async function cancelRecipeJob(jobId: string): Promise { +export async function cancelRecipeJob( + jobId: string, +): Promise { return postJson(`/jobs/${jobId}/cancel`, {}); } @@ -315,6 +340,13 @@ export async function inspectSeedUpload( return postJson("/seed/inspect-upload", payload); } +// biome-ignore lint/style/useNamingConvention: api schema +export type GithubEnvTokenStatus = { has_token: boolean }; + +export async function getGithubEnvTokenStatus(): Promise { + return getJson("/seed/github/env-token"); +} + export async function listMcpTools( payload: McpToolsListRequest, ): Promise { @@ -407,11 +439,14 @@ export async function uploadUnstructuredFile( formData.append("existing_file_ids", existingFileIds.join(",")); } - const res = await authFetch(`${DATA_DESIGNER_API_BASE}/seed/upload-unstructured-file`, { - method: "POST", - body: formData, - signal, - }); + const res = await authFetch( + `${DATA_DESIGNER_API_BASE}/seed/upload-unstructured-file`, + { + method: "POST", + body: formData, + signal, + }, + ); if (res.status === 413) { const detail = await res.json().catch(() => ({ detail: "File too large" })); @@ -420,13 +455,16 @@ export async function uploadUnstructuredFile( filename: file.name, size_bytes: file.size, status: "error", - error: typeof detail.detail === "string" ? detail.detail : "File too large", + error: + typeof detail.detail === "string" ? detail.detail : "File too large", }; } if (!res.ok) { const detail = await res.json().catch(() => ({ detail: "Upload failed" })); - throw new Error(typeof detail.detail === "string" ? detail.detail : "Upload failed"); + throw new Error( + typeof detail.detail === "string" ? detail.detail : "Upload failed", + ); } return res.json(); diff --git a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts index 363cd85be5..ed58d1418d 100644 --- a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts +++ b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts @@ -12,6 +12,7 @@ import { EqualSignIcon, FingerPrintIcon, FunctionIcon, + GithubIcon, Plug01Icon, Parabola02Icon, PencilEdit02Icon, @@ -58,11 +59,16 @@ export type BlockType = | "seed_hf" | "seed_local" | "seed_unstructured" + | "seed_github" | "model_provider" | "model_config" | "tool_config"; -export type SeedBlockType = "seed_hf" | "seed_local" | "seed_unstructured"; +export type SeedBlockType = + | "seed_hf" + | "seed_local" + | "seed_unstructured" + | "seed_github"; type IconType = typeof CodeIcon; @@ -169,6 +175,15 @@ const BLOCK_DEFINITIONS: BlockDefinition[] = [ dialogKey: "seed", createConfig: (id, existing) => makeSeedConfig(id, existing, "unstructured"), }, + { + kind: "seed", + type: "seed_github", + title: "GitHub repositories", + description: "Crawl issues, pull requests, and commits from one or more GitHub repos.", + icon: GithubIcon, + dialogKey: "seed", + createConfig: (id, existing) => makeSeedConfig(id, existing, "github_repo"), + }, { kind: "sampler", type: "category", @@ -388,6 +403,7 @@ export function getBlockDefinitionForConfig( hf: "seed_hf", local: "seed_local", unstructured: "seed_unstructured", + github_repo: "seed_github", }; return getBlockDefinition("seed", seedType[config.seed_source_type ?? "hf"]); } diff --git a/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx b/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx index a8e12ac1ca..a3e5dbe962 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx @@ -2,6 +2,8 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { ReactElement } from "react"; +import { GithubIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import type { ColumnDef } from "@tanstack/react-table"; import { Button } from "@/components/ui/button"; import { DataTable } from "@/components/ui/data-table"; @@ -12,11 +14,136 @@ import { DropdownMenuLabel, DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; +import { Progress } from "@/components/ui/progress"; import { Spinner } from "@/components/ui/spinner"; -import { cn } from "@/lib/utils"; import { isExecutionInProgress } from "../../executions/execution-helpers"; import type { RecipeExecutionRecord } from "../../execution-types"; -import { hasExpandableTextCell } from "./executions-view-helpers"; +import { formatMetricValue } from "./executions-view-helpers"; + +function formatSourceResource(value: string | null | undefined): string { + if (value === "pulls") { + return "PRs"; + } + return value ?? "--"; +} + +function formatFetchedValue(execution: RecipeExecutionRecord): string { + const source = execution.source_progress; + if (!source) { + return "--"; + } + const fetched = formatMetricValue(source.fetched_items); + if (typeof source.estimated_total !== "number" || source.estimated_total <= 0) { + return fetched; + } + return `${fetched} / ${formatMetricValue(source.estimated_total)}`; +} + +function formatGitHubSourceMessage(execution: RecipeExecutionRecord): string { + const source = execution.source_progress; + if (!source) { + return "Collecting repository threads before rows are available."; + } + if (source.status === "rate_limited") { + return source.message ?? "Waiting for GitHub rate limit. Studio will resume automatically."; + } + return source.message ?? "Collecting repository threads before rows are available."; +} + +function RunningDatasetEmptyState({ + execution, + onOpenOverview, +}: { + execution: RecipeExecutionRecord; + onOpenOverview: () => void; +}): ReactElement { + const source = execution.source_progress; + if (source?.source === "github") { + const title = + source.status === "rate_limited" + ? "Waiting for GitHub rate limit" + : "Crawling GitHub source"; + const showProgress = typeof source.percent === "number"; + + return ( +
+
+
+
+ {showProgress ? ( + + ) : ( + + )} +
+
+

{title}

+

+ {formatGitHubSourceMessage(execution)} +

+
+
+ +
+ {showProgress && } +
+

+ Repo + {source.repo ?? "--"} +

+

+ Resource + + {formatSourceResource(source.resource)} + {typeof source.page === "number" ? ` page ${source.page}` : ""} + +

+

+ Fetched + {formatFetchedValue(execution)} +

+

+ Rate remaining + {formatMetricValue(source.rate_remaining)} +

+

+ Retry wait + + {typeof source.retry_after_sec === "number" + ? `${formatMetricValue(source.retry_after_sec)}s` + : "--"} + +

+
+
+ ); + } + + return ( +
+
+ +
+

+ Generating data… +

+

+ {execution.current_column + ? `Current column: ${execution.current_column}` + : "Rows will appear here once the run produces a dataset sample."} +

+
+ +
+
+ ); +} type ExecutionDataTabProps = { execution: RecipeExecutionRecord; @@ -27,13 +154,10 @@ type ExecutionDataTabProps = { totalPages: number; tableColumns: ColumnDef>[]; datasetRowsForTable: Record[]; - visibleDatasetColumnNames: string[]; - expandedDatasetRows: Record; - selectedExecutionIdSafe: string | null; + onOpenOverview: () => void; onSetHiddenColumns: (updater: (current: string[]) => string[]) => void; onPrevPage: () => void; onNextPage: () => void; - onToggleRowExpanded: (rowId: string) => void; }; export function ExecutionDataTab({ @@ -45,13 +169,10 @@ export function ExecutionDataTab({ totalPages, tableColumns, datasetRowsForTable, - visibleDatasetColumnNames, - expandedDatasetRows, - selectedExecutionIdSafe, + onOpenOverview, onSetHiddenColumns, onPrevPage, onNextPage, - onToggleRowExpanded, }: ExecutionDataTabProps): ReactElement { return (
@@ -123,17 +244,10 @@ export function ExecutionDataTab({
{execution.dataset.length === 0 ? ( isExecutionInProgress(execution.status) ? ( -
- -
-

- Generating data… -

-

- Check the Overview tab for live terminal logs. -

-
-
+ ) : (

No rows returned.

) @@ -143,27 +257,7 @@ export function ExecutionDataTab({

) : (
- { - const canExpand = hasExpandableTextCell(row, visibleDatasetColumnNames); - if (!canExpand) { - return undefined; - } - return cn( - "cursor-pointer", - expandedDatasetRows[rowId] ? "bg-primary/[0.05]" : "hover:bg-primary/[0.06]", - ); - }} - onRowClick={(row, _rowIndex, rowId) => { - const canExpand = hasExpandableTextCell(row, visibleDatasetColumnNames); - if (!canExpand || !selectedExecutionIdSafe) { - return; - } - onToggleRowExpanded(rowId); - }} - /> +
)} diff --git a/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx b/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx index 8e2e710299..72624738fc 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx @@ -6,6 +6,7 @@ import { Database01Icon, Database02Icon, Flag02Icon, + GithubIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Badge } from "@/components/ui/badge"; @@ -23,6 +24,28 @@ import type { RecipeExecutionRecord } from "../../execution-types"; import type { ModelUsageRow } from "./executions-view-helpers"; import { formatMetricValue } from "./executions-view-helpers"; +function formatSourceResource(value: string | null | undefined): string { + if (value === "pulls") { + return "PRs"; + } + return value ?? "--"; +} + +function formatSourceMessage(execution: RecipeExecutionRecord): string { + const source = execution.source_progress; + if (!source) { + return "No source progress captured."; + } + if (source.status === "rate_limited") { + const wait = + typeof source.retry_after_sec === "number" && source.retry_after_sec > 0 + ? ` Waiting ~${formatMetricValue(source.retry_after_sec)}s.` + : ""; + return `Waiting for GitHub rate limit. Studio will resume automatically.${wait}`; + } + return source.message ?? "Crawling GitHub source."; +} + type ExecutionOverviewTabProps = { execution: RecipeExecutionRecord; showSummaryCards: boolean; @@ -60,6 +83,8 @@ export function ExecutionOverviewTab({ canPublish, onOpenPublish, }: ExecutionOverviewTabProps): ReactElement { + const sourceProgress = execution.source_progress; + return (
{showSummaryCards && ( @@ -204,6 +229,47 @@ export function ExecutionOverviewTab({ )}
)} + {sourceProgress?.source === "github" && ( +
+
+

Source data

+ +
+

+ {sourceProgress.status === "completed" + ? "GitHub source complete" + : "Crawling GitHub source"} +

+

+ {formatSourceMessage(execution)} +

+
+

+ Repo + {sourceProgress.repo ?? "--"} +

+

+ Resource + + {formatSourceResource(sourceProgress.resource)} + {typeof sourceProgress.page === "number" ? ` page ${sourceProgress.page}` : ""} + +

+

+ Fetched + + {formatMetricValue(sourceProgress.fetched_items)} + +

+

+ Rate remaining + + {formatMetricValue(sourceProgress.rate_remaining)} + +

+
+
+ )}

Terminal output

diff --git a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx index f442d8f98e..cfc0e61338 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx @@ -32,10 +32,9 @@ import { formatCellValue, formatDuration, formatPercent, - hasExpandableTextCell, + isExpandableCellValue, parseAnalysisColumns, parseModelUsageRows, - truncateCellValue, } from "./executions-view-helpers"; type ExecutionsViewProps = { @@ -63,9 +62,6 @@ export function ExecutionsView({ const [hiddenDatasetColumnsByExecution, setHiddenDatasetColumnsByExecution] = useState< Record >({}); - const [expandedDatasetRowsByExecution, setExpandedDatasetRowsByExecution] = useState< - Record> - >({}); const [previewDatasetPageByExecution, setPreviewDatasetPageByExecution] = useState< Record >({}); @@ -91,13 +87,6 @@ export function ExecutionsView({ } return hiddenDatasetColumnsByExecution[selectedExecutionIdSafe] ?? []; }, [hiddenDatasetColumnsByExecution, selectedExecutionIdSafe]); - const expandedDatasetRows = useMemo(() => { - if (!selectedExecutionIdSafe) { - return {}; - } - return expandedDatasetRowsByExecution[selectedExecutionIdSafe] ?? {}; - }, [expandedDatasetRowsByExecution, selectedExecutionIdSafe]); - const datasetColumnNames = useMemo(() => { if (!selectedExecution) { return []; @@ -119,6 +108,36 @@ export function ExecutionsView({ [datasetColumnNames, hiddenDatasetColumns], ); + // Columns where at least one row has text long enough that it would wrap at + // the default narrow width. We give those columns a wider min-width so the + // text is readable without clicking anything. The table's wrapper already + // scrolls horizontally, so a few wide columns just add a horizontal + // scrollbar instead of squeezing everything into the viewport. + const wideColumns = useMemo(() => { + const result = new Set(); + if (!selectedExecution) { + return result; + } + for (const row of selectedExecution.dataset) { + for (const name of visibleDatasetColumnNames) { + if (result.has(name)) { + continue; + } + const raw = row[name]; + if (resolveImagePreview(raw)) { + continue; + } + if (isExpandableCellValue(formatCellValue(raw))) { + result.add(name); + } + } + if (result.size === visibleDatasetColumnNames.length) { + break; + } + } + return result; + }, [selectedExecution, visibleDatasetColumnNames]); + const tableColumns = useMemo>[]>(() => { if (!selectedExecution) { return []; @@ -126,12 +145,12 @@ export function ExecutionsView({ return visibleDatasetColumnNames.map((name) => ({ accessorKey: name, header: name, - cell: ({ getValue, row }) => { + cell: ({ getValue }) => { const rawValue = getValue(); const imagePreview = resolveImagePreview(rawValue); if (imagePreview?.kind === "ready") { return ( -
+
{`${name} -

- Image too large to preview -

-
+

+ Image too large to preview +

); } const value = formatCellValue(rawValue); - const rowExpanded = Boolean(expandedDatasetRows[row.id]); - const rowHasExpandableCell = hasExpandableTextCell( - row.original, - visibleDatasetColumnNames, - ); - const showTruncated = rowHasExpandableCell && !rowExpanded; - + const isWide = wideColumns.has(name); return ( -
-

- {showTruncated ? truncateCellValue(value) : value} -

+
+

{value}

); }, })); - }, [expandedDatasetRows, selectedExecution, visibleDatasetColumnNames]); + }, [selectedExecution, visibleDatasetColumnNames, wideColumns]); const analysisColumns = useMemo( () => parseAnalysisColumns(selectedExecution?.analysis ?? null), @@ -512,9 +521,7 @@ export function ExecutionsView({ totalPages={totalPages} tableColumns={tableColumns} datasetRowsForTable={datasetRowsForTable} - visibleDatasetColumnNames={visibleDatasetColumnNames} - expandedDatasetRows={expandedDatasetRows} - selectedExecutionIdSafe={selectedExecutionIdSafe} + onOpenOverview={() => setDetailTab("overview")} onSetHiddenColumns={(updater) => { const selectedId = selectedExecution.id; setHiddenDatasetColumnsByExecution((current) => { @@ -547,18 +554,6 @@ export function ExecutionsView({ } onLoadDatasetPage(selectedExecution.id, currentDatasetPage + 1); }} - onToggleRowExpanded={(rowId) => { - setExpandedDatasetRowsByExecution((current) => { - const rows = current[selectedExecution.id] ?? {}; - return { - ...current, - [selectedExecution.id]: { - ...rows, - [rowId]: !rows[rowId], - }, - }; - }); - }} /> diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx index bf022fa028..82ec2f1de8 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx @@ -1,7 +1,12 @@ // 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 { DocumentAttachmentIcon, DocumentCodeIcon, Plant01Icon } from "@hugeicons/core-free-icons"; +import { + DocumentAttachmentIcon, + DocumentCodeIcon, + GithubIcon, + Plant01Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import type { ReactElement } from "react"; import type { SeedConfig } from "../../types"; @@ -13,9 +18,75 @@ type InlineSeedProps = { onUpdate: (patch: Partial) => void; }; -export function InlineSeed({ config, onUpdate }: InlineSeedProps): ReactElement { +const GITHUB_REPO_LIST_SPLIT_RE = /[\n,]+/; + +function parseGithubRepos(value: string | undefined): string[] { + return (value ?? "") + .split(GITHUB_REPO_LIST_SPLIT_RE) + .map((repo) => repo.trim()) + .filter(Boolean); +} + +function hasInvalidGithubRepo(repo: string): boolean { + const parts = repo.split("/"); + return parts.length !== 2 || parts.some((part) => !part.trim()); +} + +export function InlineSeed({ + config, + onUpdate, +}: InlineSeedProps): ReactElement { const mode = config.seed_source_type ?? "hf"; + if (mode === "github_repo") { + const repos = parseGithubRepos(config.github_repo_slug); + const invalidCount = repos.filter(hasInvalidGithubRepo).length; + const summary = + repos.length === 0 + ? "Add GitHub repos" + : repos.length === 1 + ? repos[0] + : `${repos.length} repositories`; + const configuredItems = config.github_item_types; + const items = + configuredItems && configuredItems.length > 0 + ? configuredItems + : ["issues", "pulls"]; + const itemsLabel = items.join(" · "); + const limit = (config.github_limit ?? "100").trim() || "100"; + const commentsLabel = + config.github_include_comments === false + ? "comments off" + : `comments ≤ ${config.github_max_comments_per_item ?? "30"}`; + const tokenLabel = config.github_token?.trim() ? "PAT set" : "server token"; + const warning = + repos.length === 0 + ? "No repos configured" + : invalidCount > 0 + ? `${invalidCount} invalid repo${invalidCount === 1 ? "" : "s"}` + : null; + return ( +
+
+ +
+
+

{summary}

+

+ {warning ?? + `${itemsLabel} · limit ${limit} · ${commentsLabel} · ${tokenLabel}`} +

+
+ +
+ ); + } + if (mode === "hf") { return (
@@ -62,10 +133,14 @@ export function InlineSeed({ config, onUpdate }: InlineSeedProps): ReactElement {fileName || "No file selected"}

- {isLocal ? "Structured file" : "Unstructured document"} · configure in dialog + {isLocal ? "Structured file" : "Unstructured document"} · configure in + dialog

- +
); } diff --git a/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx b/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx index acbda2d7d8..385bef2b77 100644 --- a/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx +++ b/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx @@ -34,6 +34,7 @@ type RecipeStudioHeaderProps = { savedAtLabel: string; workflowName: string; warnings?: GraphWarning[]; + supportsEasyMode?: boolean; onWorkflowNameChange: (value: string) => void; onViewChange: (view: RecipeStudioView) => void; onSaveRecipe: () => void; @@ -51,6 +52,7 @@ export function RecipeStudioHeader({ savedAtLabel, workflowName, warnings = [], + supportsEasyMode = false, onWorkflowNameChange, onViewChange, onSaveRecipe, @@ -58,7 +60,7 @@ export function RecipeStudioHeader({ const [editingWorkflowName, setEditingWorkflowName] = useState(false); function handleViewValueChange(value: string): void { - if (value === "editor" || value === "executions") { + if (value === "easy" || value === "editor" || value === "executions") { onViewChange(value); } } @@ -130,7 +132,12 @@ export function RecipeStudioHeader({
- Editor + {supportsEasyMode && ( + Easy + )} + + {supportsEasyMode ? "Advanced" : "Editor"} + Runs diff --git a/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx b/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx index 43f7437848..563280c55f 100644 --- a/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx +++ b/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx @@ -49,6 +49,52 @@ function statusLabel(input: { return "Run status"; } +function formatSourceResource(value: string | null | undefined): string { + if (value === "pulls") { + return "PRs"; + } + return value ?? "source"; +} + +function formatGitHubSourceSummary( + execution: RecipeExecutionRecord, +): string | null { + const source = execution.source_progress; + if (!source || source.source !== "github") { + return null; + } + if (source.status === "rate_limited") { + const wait = + typeof source.retry_after_sec === "number" && source.retry_after_sec > 0 + ? ` ~${formatMetricValue(source.retry_after_sec)}s` + : ""; + return `Waiting for GitHub rate limit${wait}. Studio will resume automatically.`; + } + if (source.status === "retrying") { + return source.message ?? "GitHub request failed; retrying automatically."; + } + + const parts = [ + source.repo, + source.resource + ? `${formatSourceResource(source.resource)}${ + typeof source.page === "number" ? ` page ${source.page}` : "" + }` + : null, + typeof source.fetched_items === "number" + ? `${formatMetricValue(source.fetched_items)} fetched` + : null, + typeof source.rate_remaining === "number" + ? `remaining ${formatMetricValue(source.rate_remaining)}` + : null, + ].filter(Boolean); + + if (parts.length === 0) { + return source.message ?? "Crawling GitHub source"; + } + return parts.join(" · "); +} + export function ExecutionProgressIsland({ execution, currentColumnIcon, @@ -58,14 +104,24 @@ export function ExecutionProgressIsland({ }: ExecutionProgressIslandProps): ReactElement { const complete = execution.status === "completed"; const inProgress = isExecutionInProgress(execution.status); - const progressPercent = execution.progress?.percent ?? (complete ? 100 : 0); + const sourceSummary = formatGitHubSourceSummary(execution); + const showSourceProgress = Boolean( + sourceSummary && + inProgress && + (execution.stage === "source" || + execution.source_progress?.status === "rate_limited" || + execution.source_progress?.status === "retrying"), + ); + const sourcePercent = showSourceProgress ? execution.source_progress?.percent : null; + const progressPercent = sourcePercent ?? execution.progress?.percent ?? (complete ? 100 : 0); const hasProgressSignal = Boolean( - execution.progress && + (execution.progress && (typeof execution.progress.done === "number" || typeof execution.progress.total === "number" || typeof execution.progress.percent === "number" || typeof execution.progress.rate === "number" || - typeof execution.progress.eta_sec === "number"), + typeof execution.progress.eta_sec === "number")) || + typeof sourcePercent === "number", ); const showLoadingSpinner = inProgress && !hasProgressSignal; const batchTotal = execution.batch?.total ?? null; @@ -136,18 +192,30 @@ export function ExecutionProgressIsland({ ETA: {formatEta(execution.progress?.eta_sec)}

-
- -

- Column: {execution.current_column ?? "--"} -

-
+ {showSourceProgress ? ( +
+ +

+ GitHub source: {sourceSummary} +

+
+ ) : ( +
+ +

+ Column: {execution.current_column ?? "--"} +

+
+ )} {showBatch && (
)} - {!validateResult.valid && validateResult.rawDetail && ( -

+ {validateResult.rawDetail && ( +

{validateResult.rawDetail}

)} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 868974acb5..54eae08f7c 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -31,18 +31,27 @@ import { TableHeader, TableRow, } from "@/components/ui/table"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { - Tabs, - TabsContent, - TabsList, - TabsTrigger, -} from "@/components/ui/tabs"; -import { type ReactElement, useCallback, useEffect, useMemo, useRef, useState } from "react"; + type KeyboardEvent, + type ReactElement, + useCallback, + useEffect, + useId, + useMemo, + useRef, + useState, +} from "react"; import { cn } from "@/lib/utils"; import { UnstructuredDropZone, type FileEntry } from "./unstructured-drop-zone"; -import { inspectSeedDataset, inspectSeedUpload } from "../../api"; +import { + getGithubEnvTokenStatus, + inspectSeedDataset, + inspectSeedUpload, +} from "../../api"; import { resolveImagePreview } from "../../utils/image-preview"; import type { + GithubItemType, SeedConfig, SeedSamplingStrategy, SeedSelectionType, @@ -51,10 +60,11 @@ import { CollapsibleSectionTriggerButton } from "../shared/collapsible-section-t import { HfDatasetCombobox } from "../../components/shared/hf-dataset-combobox"; import { FieldLabel } from "../shared/field-label"; -const SAMPLING_OPTIONS: Array<{ value: SeedSamplingStrategy; label: string }> = [ - { value: "ordered", label: "Ordered" }, - { value: "shuffle", label: "Shuffle" }, -]; +const SAMPLING_OPTIONS: Array<{ value: SeedSamplingStrategy; label: string }> = + [ + { value: "ordered", label: "Ordered" }, + { value: "shuffle", label: "Shuffle" }, + ]; const SELECTION_OPTIONS: Array<{ value: SeedSelectionType; label: string }> = [ { value: "none", label: "None" }, @@ -68,6 +78,14 @@ const DEFAULT_CHUNK_SIZE = 1200; const DEFAULT_CHUNK_OVERLAP = 200; const MAX_CHUNK_SIZE = 20000; const PREVIEW_TRUNCATE_AT = 320; +const GITHUB_TRAILING_PUNCTUATION_RE = /[),.;]+$/g; +const GITHUB_SSH_PREFIX_RE = /^git@github\.com:/i; +const URL_PROTOCOL_RE = /^https?:\/\//i; +const WWW_PREFIX_RE = /^www\./i; +const URL_QUERY_OR_HASH_RE = /[?#]/; +const GIT_SUFFIX_RE = /\.git$/i; +const GITHUB_REPO_INPUT_SPLIT_RE = /[\s,]+/; +const GITHUB_REPO_PART_RE = /^[A-Za-z0-9_.-]+$/; type SeedDialogProps = { config: SeedConfig; @@ -75,6 +93,367 @@ type SeedDialogProps = { open: boolean; }; +function normalizeGithubRepoInput(value: string): string { + let raw = value.trim().replace(GITHUB_TRAILING_PUNCTUATION_RE, ""); + if (!raw) { + return ""; + } + raw = raw.replace(GITHUB_SSH_PREFIX_RE, "https://github.com/"); + raw = raw.replace(URL_PROTOCOL_RE, ""); + raw = raw.replace(WWW_PREFIX_RE, ""); + if (raw.toLowerCase().startsWith("github.com/")) { + raw = raw.slice("github.com/".length); + } + raw = raw.split(URL_QUERY_OR_HASH_RE)[0] ?? ""; + raw = raw.replace(GIT_SUFFIX_RE, ""); + const parts = raw.split("/").filter(Boolean); + if (parts.length >= 2) { + return `${parts[0]}/${parts[1]}`; + } + return raw; +} + +function splitGithubRepoInput(value: string): string[] { + return value + .split(GITHUB_REPO_INPUT_SPLIT_RE) + .map(normalizeGithubRepoInput) + .filter(Boolean); +} + +function getGithubRepoError(repo: string): string | null { + const parts = repo.split("/"); + if (parts.length !== 2 || parts.some((part) => !part.trim())) { + return "Use owner/name."; + } + if (parts.some((part) => !GITHUB_REPO_PART_RE.test(part))) { + return "Use only GitHub repo characters."; + } + return null; +} + +function githubLimitString(value: string | undefined): string { + return (value ?? "100").trim(); +} + +export function GithubRepoSeedForm({ + config, + onUpdate, +}: { + config: SeedConfig; + onUpdate: (patch: Partial) => void; +}): ReactElement { + const [repoDraft, setRepoDraft] = useState(""); + const [serverHasEnvToken, setServerHasEnvToken] = useState( + null, + ); + const repoInputId = useId(); + const repoHelpId = useId(); + const repoErrorId = useId(); + const tokenId = useId(); + const tokenHelpId = useId(); + const limitId = useId(); + const limitHelpId = useId(); + const commentsId = useId(); + const includeCommentsId = useId(); + const commentsHelpId = useId(); + const repos = useMemo( + () => splitGithubRepoInput(config.github_repo_slug ?? ""), + [config.github_repo_slug], + ); + const repoErrors = repos + .map((repo, index) => ({ repo, index, error: getGithubRepoError(repo) })) + .filter((item) => item.error); + const hasRepoErrors = repoErrors.length > 0; + const configuredItemTypes = config.github_item_types; + const itemTypes: GithubItemType[] = + configuredItemTypes && configuredItemTypes.length > 0 + ? configuredItemTypes + : ["issues", "pulls"]; + const limitNum = Number.parseInt(githubLimitString(config.github_limit), 10); + const boundedLimit = Number.isFinite(limitNum) + ? Math.min(5000, Math.max(1, limitNum)) + : 100; + const estimatedItems = repos.length * itemTypes.length * boundedLimit; + const includeComments = config.github_include_comments ?? true; + const hasToken = Boolean(config.github_token?.trim()); + const usingEnvToken = !hasToken && serverHasEnvToken === true; + + useEffect(() => { + let cancelled = false; + void getGithubEnvTokenStatus() + .then((status) => { + if (!cancelled) setServerHasEnvToken(status.has_token); + }) + .catch(() => { + if (!cancelled) setServerHasEnvToken(false); + }); + return () => { + cancelled = true; + }; + }, []); + + function updateRepos(nextRepos: string[]): void { + const seen = new Set(); + const deduped = nextRepos.filter((repo) => { + const key = repo.toLowerCase(); + if (seen.has(key)) { + return false; + } + seen.add(key); + return true; + }); + onUpdate({ github_repo_slug: deduped.join("\n") }); + } + + function addRepos(raw: string): void { + const next = splitGithubRepoInput(raw); + if (next.length === 0) { + return; + } + updateRepos([...repos, ...next]); + setRepoDraft(""); + } + + function removeRepo(index: number): void { + updateRepos(repos.filter((_, i) => i !== index)); + } + + function handleRepoKeyDown(event: KeyboardEvent): void { + if (["Enter", ",", " "].includes(event.key)) { + event.preventDefault(); + addRepos(repoDraft); + } else if (event.key === "Backspace" && !repoDraft && repos.length > 0) { + event.preventDefault(); + removeRepo(repos.length - 1); + } + } + + return ( +
+
+ +
+ {repos.map((repo, index) => { + const error = getGithubRepoError(repo); + return ( + + {repo} + + + ); + })} + setRepoDraft(event.target.value)} + onKeyDown={handleRepoKeyDown} + onBlur={() => addRepos(repoDraft)} + onPaste={(event) => { + const pasted = event.clipboardData.getData("text"); + if (pasted) { + event.preventDefault(); + addRepos(pasted); + } + }} + placeholder={ + repos.length === 0 + ? "unslothai/unsloth or GitHub URL" + : "Add repo…" + } + aria-invalid={hasRepoErrors} + aria-describedby={`${repoHelpId}${hasRepoErrors ? ` ${repoErrorId}` : ""}`} + /> +
+

+ Stored as one owner/name repo per line in the recipe. +

+ {hasRepoErrors && ( +
    + {repoErrors.map((item) => ( +
  • + Row {item.index + 1}: {item.error} +
    • + ))} +
+ )} +
+ +
+
+ + {usingEnvToken && ( + + Using server env var + + )} +
+ onUpdate({ github_token: e.target.value })} + placeholder={ + usingEnvToken + ? "Using server GH_TOKEN / GITHUB_TOKEN" + : "Leave blank to use server GH_TOKEN" + } + aria-describedby={tokenHelpId} + /> +

+ {usingEnvToken + ? "Studio detected a server env token, so saved/shared recipes can leave this blank." + : "Blank is safest for saved/shared recipes because Studio will read the server environment at run time."} +

+ {hasToken && ( +

+ Personal access tokens are sensitive. Prefer server env vars when + possible, and avoid sharing recipes that contain a PAT. +

+ )} +
+ +
+ + Fetch scope + +
+ + onUpdate({ github_limit: e.target.value })} + placeholder="100" + aria-describedby={limitHelpId} + /> +

+ Estimate before comments: up to {estimatedItems.toLocaleString()}{" "} + rows ({repos.length || 0} repos × {itemTypes.length} item types ×{" "} + {boundedLimit} limit). Commit crawling uses each repo's default + branch. +

+
+ +
+ + Item types + +
+ {(["issues", "pulls", "commits"] as const).map((kind) => { + const checked = itemTypes.includes(kind); + const itemTypeId = `${repoInputId}-${kind}`; + return ( + + ); + })} +
+
+ +
+ +
+ + + onUpdate({ github_max_comments_per_item: e.target.value }) + } + aria-describedby={commentsHelpId} + /> +

+ Comments increase GraphQL cost and can make Check/Run look quiet + while GitHub pages and rate-limit waits stream in logs. +

+
+
+
+ +

+ Backed by Studio's built-in github_repo seed reader. Large + repos can take minutes, so start with small limits for previews. +

+
+ ); +} + function getErrorMessage(error: unknown, fallback: string): string { if (error instanceof Error && error.message) { return error.message; @@ -85,7 +464,8 @@ function getErrorMessage(error: unknown, fallback: string): string { function stringifyCell(value: unknown): string { if (value === null || value === undefined) return ""; if (typeof value === "string") return value; - if (typeof value === "number" || typeof value === "boolean") return String(value); + if (typeof value === "number" || typeof value === "boolean") + return String(value); try { return JSON.stringify(value); } catch { @@ -111,7 +491,8 @@ function getPreviewEmptyStateCopy(mode: SeedConfig["seed_source_type"]): { if (mode === "local") { return { title: "No preview yet", - description: "Upload a CSV, JSON, or JSONL file and click Load to see a sample.", + description: + "Upload a CSV, JSON, or JSONL file and click Load to see a sample.", }; } if (mode === "unstructured") { @@ -121,9 +502,17 @@ function getPreviewEmptyStateCopy(mode: SeedConfig["seed_source_type"]): { "Upload your documents and the preview will appear once processing is done.", }; } + if (mode === "github_repo") { + return { + title: "GitHub data loads during Check or Run", + description: + "Configure repos, item types, and limits above. GitHub crawling can take minutes on large repos; watch logs for page and rate-limit updates.", + }; + } return { title: "No preview yet", - description: "Select a Hugging Face dataset and click Load to see a sample.", + description: + "Select a Hugging Face dataset and click Load to see a sample.", }; } @@ -175,24 +564,32 @@ async function fileToBase64Payload(file: File): Promise { }); } -export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactElement { +export function SeedDialog({ + config, + onUpdate, + open, +}: SeedDialogProps): ReactElement { const [inspectError, setInspectError] = useState(null); const [isInspecting, setIsInspecting] = useState(false); const advancedOpen = config.advancedOpen === true; const [previewRows, setPreviewRows] = useState[]>([]); - const [expandedPreviewRows, setExpandedPreviewRows] = useState>({}); + const [expandedPreviewRows, setExpandedPreviewRows] = useState< + Record + >({}); const [localFile, setLocalFile] = useState(null); - const [unstructuredFiles, setUnstructuredFiles] = useState(() => { - if (config.unstructured_file_ids?.length) { - return config.unstructured_file_ids.map((id, i) => ({ - id, - name: config.unstructured_file_names?.[i] ?? "Unknown", - size: config.unstructured_file_sizes?.[i] ?? 0, - status: "ok" as const, - })); - } - return []; - }); + const [unstructuredFiles, setUnstructuredFiles] = useState( + () => { + if (config.unstructured_file_ids?.length) { + return config.unstructured_file_ids.map((id, i) => ({ + id, + name: config.unstructured_file_names?.[i] ?? "Unknown", + size: config.unstructured_file_sizes?.[i] ?? 0, + status: "ok" as const, + })); + } + return []; + }, + ); const mode = config.seed_source_type ?? "hf"; const previewEmpty = getPreviewEmptyStateCopy(mode); @@ -241,7 +638,14 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl status: "ok" as const, })), ); - }, [open, mode, unstructuredFiles.length, config.unstructured_file_ids, config.unstructured_file_names, config.unstructured_file_sizes]); + }, [ + open, + mode, + unstructuredFiles.length, + config.unstructured_file_ids, + config.unstructured_file_names, + config.unstructured_file_sizes, + ]); const handleUnstructuredFilesChange = useCallback( (updater: FileEntry[] | ((prev: FileEntry[]) => FileEntry[])) => { @@ -291,142 +695,135 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl const { chunkSize, chunkOverlap } = resolveChunking(config); const fileKey = okFiles.map((f) => `${f.id}|${f.name}`).join(","); return `unstructured:${fileKey}|${chunkSize}|${chunkOverlap}`; - }, [ - config, - localFile, - mode, - unstructuredFiles, - ]); + }, [config, localFile, mode, unstructuredFiles]); - const loadSeedMetadata = useCallback(async (opts?: { silent?: boolean }): Promise => { - const loadKey = getCurrentLoadKey(); - if (!opts?.silent) { - setInspectError(null); - } - setIsInspecting(true); - try { - if (mode === "hf") { - const datasetName = config.hf_repo_id.trim(); - if (!datasetName) { - throw new Error("Dataset repo is required."); - } - const response = await inspectSeedDataset({ - dataset_name: datasetName, - hf_token: config.hf_token?.trim() || undefined, - split: config.hf_split?.trim() || undefined, - subset: config.hf_subset?.trim() || undefined, - preview_size: 10, - }); - onUpdate({ - hf_path: response.resolved_path, - seed_columns: response.columns, - seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => - response.columns.includes(name), - ), - seed_preview_rows: response.preview_rows ?? [], - hf_split: response.split ?? "", - hf_subset: response.subset ?? "", - local_file_name: "", - unstructured_file_ids: [], - unstructured_file_names: [], - unstructured_file_sizes: [], - }); - setPreviewRows(response.preview_rows ?? []); - setLastLoadedKey(loadKey); - return true; - } - - if (mode === "local") { - if (!localFile) { - throw new Error("Select a local CSV/JSON/JSONL file first."); - } - if (localFile.size > MAX_UPLOAD_BYTES) { - throw new Error("File too large (max 50MB)."); - } - const payload = await fileToBase64Payload(localFile); - const response = await inspectSeedUpload({ - filename: localFile.name, - content_base64: payload, - preview_size: 10, - }); - onUpdate({ - hf_path: response.resolved_path, - seed_columns: response.columns, - seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => - response.columns.includes(name), - ), - seed_preview_rows: response.preview_rows ?? [], - hf_repo_id: "", - hf_subset: "", - hf_split: "", - local_file_name: localFile.name, - unstructured_file_ids: [], - unstructured_file_names: [], - unstructured_file_sizes: [], - }); - setPreviewRows(response.preview_rows ?? []); - setLastLoadedKey(loadKey); - return true; - } - - if (mode === "unstructured") { - const fileIds = unstructuredFiles - .filter((f) => f.status === "ok") - .map((f) => f.id); - const fileNames = unstructuredFiles - .filter((f) => f.status === "ok") - .map((f) => f.name); - - if (fileIds.length === 0) { - setInspectError("No files uploaded"); - return false; - } - - const { chunkSize, chunkOverlap } = resolveChunking(config); - const response = await inspectSeedUpload({ - block_id: config.id, - file_ids: fileIds, - file_names: fileNames, - preview_size: 10, - seed_source_type: "unstructured", - unstructured_chunk_size: chunkSize, - unstructured_chunk_overlap: chunkOverlap, - }); - - onUpdate({ - hf_path: response.resolved_path, - resolved_paths: response.resolved_paths ?? [], - seed_columns: response.columns, - seed_preview_rows: response.preview_rows ?? [], - unstructured_file_ids: fileIds, - unstructured_file_names: fileNames, - unstructured_file_sizes: unstructuredFiles - .filter((f) => f.status === "ok") - .map((f) => f.size), - }); - setPreviewRows(response.preview_rows ?? []); - setLastLoadedKey(loadKey); - return true; - } - - return false; - } catch (error) { + const loadSeedMetadata = useCallback( + async (opts?: { silent?: boolean }): Promise => { + const loadKey = getCurrentLoadKey(); if (!opts?.silent) { - setInspectError(getErrorMessage(error, "Failed to load seed metadata.")); + setInspectError(null); } - setPreviewRows([]); - return false; - } finally { - setIsInspecting(false); - } - }, [ - config, - getCurrentLoadKey, - localFile, - mode, - onUpdate, - unstructuredFiles, - ]); + setIsInspecting(true); + try { + if (mode === "hf") { + const datasetName = config.hf_repo_id.trim(); + if (!datasetName) { + throw new Error("Dataset repo is required."); + } + const response = await inspectSeedDataset({ + dataset_name: datasetName, + hf_token: config.hf_token?.trim() || undefined, + split: config.hf_split?.trim() || undefined, + subset: config.hf_subset?.trim() || undefined, + preview_size: 10, + }); + onUpdate({ + hf_path: response.resolved_path, + seed_columns: response.columns, + seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => + response.columns.includes(name), + ), + seed_preview_rows: response.preview_rows ?? [], + hf_split: response.split ?? "", + hf_subset: response.subset ?? "", + local_file_name: "", + unstructured_file_ids: [], + unstructured_file_names: [], + unstructured_file_sizes: [], + }); + setPreviewRows(response.preview_rows ?? []); + setLastLoadedKey(loadKey); + return true; + } + + if (mode === "local") { + if (!localFile) { + throw new Error("Select a local CSV/JSON/JSONL file first."); + } + if (localFile.size > MAX_UPLOAD_BYTES) { + throw new Error("File too large (max 50MB)."); + } + const payload = await fileToBase64Payload(localFile); + const response = await inspectSeedUpload({ + filename: localFile.name, + content_base64: payload, + preview_size: 10, + }); + onUpdate({ + hf_path: response.resolved_path, + seed_columns: response.columns, + seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => + response.columns.includes(name), + ), + seed_preview_rows: response.preview_rows ?? [], + hf_repo_id: "", + hf_subset: "", + hf_split: "", + local_file_name: localFile.name, + unstructured_file_ids: [], + unstructured_file_names: [], + unstructured_file_sizes: [], + }); + setPreviewRows(response.preview_rows ?? []); + setLastLoadedKey(loadKey); + return true; + } + + if (mode === "unstructured") { + const fileIds = unstructuredFiles + .filter((f) => f.status === "ok") + .map((f) => f.id); + const fileNames = unstructuredFiles + .filter((f) => f.status === "ok") + .map((f) => f.name); + + if (fileIds.length === 0) { + setInspectError("No files uploaded"); + return false; + } + + const { chunkSize, chunkOverlap } = resolveChunking(config); + const response = await inspectSeedUpload({ + block_id: config.id, + file_ids: fileIds, + file_names: fileNames, + preview_size: 10, + seed_source_type: "unstructured", + unstructured_chunk_size: chunkSize, + unstructured_chunk_overlap: chunkOverlap, + }); + + onUpdate({ + hf_path: response.resolved_path, + resolved_paths: response.resolved_paths ?? [], + seed_columns: response.columns, + seed_preview_rows: response.preview_rows ?? [], + unstructured_file_ids: fileIds, + unstructured_file_names: fileNames, + unstructured_file_sizes: unstructuredFiles + .filter((f) => f.status === "ok") + .map((f) => f.size), + }); + setPreviewRows(response.preview_rows ?? []); + setLastLoadedKey(loadKey); + return true; + } + + return false; + } catch (error) { + if (!opts?.silent) { + setInspectError( + getErrorMessage(error, "Failed to load seed metadata."), + ); + } + setPreviewRows([]); + return false; + } finally { + setIsInspecting(false); + } + }, + [config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles], + ); useEffect(() => { const wasOpen = wasOpenRef.current; @@ -463,7 +860,8 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl return []; }, [config.seed_columns, previewRows]); const selectedSeedDropColumns = useMemo( - () => (config.seed_drop_columns ?? []).filter((name) => name.trim().length > 0), + () => + (config.seed_drop_columns ?? []).filter((name) => name.trim().length > 0), [config.seed_drop_columns], ); const selectedSeedDropSet = useMemo( @@ -540,10 +938,11 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl className="nodrag" placeholder="hf_..." value={config.hf_token ?? ""} - onChange={(event) => onUpdate({ hf_token: event.target.value })} + onChange={(event) => + onUpdate({ hf_token: event.target.value }) + } />
- )} @@ -600,9 +999,15 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl /> )} - {inspectError &&

{inspectError}

} + {mode === "github_repo" && ( + + )} - {mode !== "unstructured" && ( + {inspectError && ( +

{inspectError}

+ )} + + {mode !== "unstructured" && mode !== "github_repo" && (
{ const isChecked = value === true; const next = isChecked - ? Array.from(new Set([...selectedSeedDropColumns, columnName])) - : selectedSeedDropColumns.filter((name) => name !== columnName); + ? Array.from( + new Set([ + ...selectedSeedDropColumns, + columnName, + ]), + ) + : selectedSeedDropColumns.filter( + (name) => name !== columnName, + ); onUpdate({ seed_drop_columns: next }); }} /> @@ -660,7 +1072,9 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl onUpdate({ selection_start: event.target.value })} + onChange={(event) => + onUpdate({ selection_start: event.target.value }) + } />
- + onUpdate({ selection_end: event.target.value })} + onChange={(event) => + onUpdate({ selection_end: event.target.value }) + } />
@@ -772,17 +1203,24 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl className="nodrag" inputMode="numeric" value={config.selection_index ?? ""} - onChange={(event) => onUpdate({ selection_index: event.target.value })} + onChange={(event) => + onUpdate({ selection_index: event.target.value }) + } />
- + - onUpdate({ selection_num_partitions: event.target.value }) + onUpdate({ + selection_num_partitions: event.target.value, + }) } />
@@ -830,7 +1268,8 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl { @@ -850,7 +1289,9 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl className="max-w-[260px] whitespace-pre-wrap break-words text-xs" > {(() => { - const imagePreview = resolveImagePreview(row[col]); + const imagePreview = resolveImagePreview( + row[col], + ); if (imagePreview?.kind === "ready") { return ( ; + rows: number; + setRows: (rows: number) => void; + updateConfig: (id: string, patch: Partial) => void; + onRun: () => void; + runLoading: boolean; + runErrors: string[]; + onSwitchToAdvanced: () => void; +}; + +export function GithubCrawlerEasyView({ + configs, + rows, + setRows, + updateConfig, + onRun, + runLoading, + runErrors, + onSwitchToAdvanced, +}: GithubCrawlerEasyViewProps): ReactElement { + const seedConfig = useMemo( + () => + Object.values(configs).find((c): c is SeedConfig => c.kind === "seed") ?? + null, + [configs], + ); + const modelConfig = useMemo( + () => + Object.values(configs).find( + (c): c is ModelConfig => c.kind === "model_config", + ) ?? null, + [configs], + ); + + // Local buffer for the Rows input so the user can hold transient invalid + // state (empty while backspacing, partial digits, etc.) without the parent + // snapping them back to 1 on every keystroke. The canonical ``rows`` value + // only advances when the buffer parses to a valid positive integer; on + // blur we clamp back to a sane default if the user left it empty. + const [rowsText, setRowsText] = useState(String(rows)); + useEffect(() => { + setRowsText(String(rows)); + }, [rows]); + + const handleSeedUpdate = (patch: Partial): void => { + if (!seedConfig) return; + updateConfig(seedConfig.id, patch); + }; + + const handleModelChange = (value: string): void => { + if (!modelConfig) return; + updateConfig(modelConfig.id, { model: value }); + }; + + if (!seedConfig) { + return ( +
+

+ This recipe has no seed node. Switch to{" "} + {" "} + to configure it. +

+
+ ); + } + + return ( +
+
+
+ +
+
+

GitHub Crawler

+

+ Crawl real GitHub issues and PRs and turn each thread into a{" "} + {"{User, Assistant}"} training pair. + Defaults use the server's GH_TOKEN env var and the + bundled local model. +

+
+
+ + + +
+

+ Run settings +

+
+
+ + { + // Allow empty / partial strings while the user is editing. + // type="text" avoids the browser's number spinner and the + // related backspace quirks; we still parse + clamp below. + const raw = event.target.value.replace(/[^0-9]/g, ""); + setRowsText(raw); + const next = Number.parseInt(raw, 10); + if (Number.isFinite(next) && next > 0 && next <= 10000) { + setRows(next); + } + }} + onBlur={() => { + const next = Number.parseInt(rowsText, 10); + if (!Number.isFinite(next) || next < 1) { + setRows(1); + setRowsText("1"); + } else if (next > 10000) { + setRows(10000); + setRowsText("10000"); + } else { + setRowsText(String(next)); + } + }} + /> +
+
+ + handleModelChange(event.target.value)} + placeholder="unsloth/gemma-4-E2B-it-GGUF" + disabled={!modelConfig} + /> +
+
+
+ + {runErrors.length > 0 && ( +
+

Cannot run:

+
    + {runErrors.slice(0, 4).map((err) => ( +
  • {err}
    • + ))} +
+
+ )} + +
+ + +
+
+ ); +} diff --git a/studio/frontend/src/features/recipe-studio/execution-types.ts b/studio/frontend/src/features/recipe-studio/execution-types.ts index c1976ab3f9..92dd88c4f5 100644 --- a/studio/frontend/src/features/recipe-studio/execution-types.ts +++ b/studio/frontend/src/features/recipe-studio/execution-types.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -export type RecipeStudioView = "editor" | "executions"; +export type RecipeStudioView = "easy" | "editor" | "executions"; export type RecipeExecutionKind = "preview" | "full"; @@ -30,6 +30,28 @@ export type RecipeExecutionBatch = { total?: number | null; }; +export type RecipeSourceProgress = { + source?: string | null; + status?: string | null; + repo?: string | null; + resource?: string | null; + page?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + page_items?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + fetched_items?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + estimated_total?: number | null; + percent?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + rate_remaining?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + retry_after_sec?: number | null; + message?: string | null; + // biome-ignore lint/style/useNamingConvention: backend schema + updated_at?: number | null; +}; + export type RecipeExecutionAnalysis = { num_records?: number; target_num_records?: number; @@ -65,6 +87,8 @@ export type RecipeExecutionRecord = { column_progress: RecipeExecutionProgress | null; batch: RecipeExecutionBatch | null; // biome-ignore lint/style/useNamingConvention: backend schema + source_progress: RecipeSourceProgress | null; + // biome-ignore lint/style/useNamingConvention: backend schema model_usage: Record | null; // biome-ignore lint/style/useNamingConvention: backend schema lastEventId: number | null; diff --git a/studio/frontend/src/features/recipe-studio/executions/runtime.ts b/studio/frontend/src/features/recipe-studio/executions/runtime.ts index bdc3e653fc..e37a1c9929 100644 --- a/studio/frontend/src/features/recipe-studio/executions/runtime.ts +++ b/studio/frontend/src/features/recipe-studio/executions/runtime.ts @@ -6,6 +6,7 @@ import type { RecipeExecutionBatch, RecipeExecutionKind, RecipeExecutionRecord, + RecipeSourceProgress, } from "../execution-types"; import { DATASET_PAGE_SIZE, @@ -72,6 +73,14 @@ export function toExecutionLogLine(event: JobEvent): string | null { return null; } +function normalizeSourceProgress(input: unknown): RecipeSourceProgress | null { + const raw = normalizeObject(input); + if (!raw) { + return null; + } + return raw as RecipeSourceProgress; +} + export function applyExecutionStatusSnapshot( execution: RecipeExecutionRecord, status: JobStatusResponse, @@ -100,6 +109,7 @@ export function applyExecutionStatusSnapshot( (normalizeObject(status.column_progress) as RecipeExecutionRecord["column_progress"]) ?? null, batch, + source_progress: normalizeSourceProgress(status.source_progress), model_usage: normalizeObject(status.model_usage), artifact_path: status.artifact_path ?? execution.artifact_path, error: status.error ?? null, @@ -137,6 +147,7 @@ export function createBaseExecutionRecord(input: { progress: null, column_progress: null, batch: null, + source_progress: null, model_usage: null, lastEventId: null, artifact_path: null, diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts index 24d02e0f6c..1040025a5d 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts @@ -3,7 +3,12 @@ import { useCallback, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; +import { toast } from "sonner"; import { toastError } from "@/shared/toast"; +import { + getInferenceStatus, + loadModel, +} from "@/features/chat/api/chat-api"; import { cancelRecipeJob, createRecipeJob, @@ -39,10 +44,84 @@ import { } from "../stores/recipe-executions"; import type { RecipePayload, RecipePayloadResult } from "../utils/payload/types"; +/** + * Auto-load the local model before running a recipe that uses it. + * + * Looks at payload.recipe.model_providers for any provider with is_local=true, + * finds the bound model_configs and asks the backend to load whichever model + * the first local-bound model_config points at. Skips when the inference + * server already has that exact model active. This removes the "open /chat + * first" prerequisite that users kept tripping on. + */ +async function ensureLocalModelLoaded( + payload: RecipePayload, +): Promise { + const providers = Array.isArray(payload.recipe.model_providers) + ? (payload.recipe.model_providers as Array>) + : []; + const localProviderNames = new Set(); + for (const p of providers) { + if (p.is_local === true && typeof p.name === "string") { + localProviderNames.add(p.name); + } + } + if (localProviderNames.size === 0) { + return null; + } + const modelConfigs = Array.isArray(payload.recipe.model_configs) + ? (payload.recipe.model_configs as Array>) + : []; + const boundConfig = modelConfigs.find( + (c) => typeof c.provider === "string" && localProviderNames.has(c.provider), + ); + const target = + typeof boundConfig?.model === "string" ? boundConfig.model.trim() : ""; + if (!target) { + return null; + } + + try { + const status = await getInferenceStatus(); + if ( + status.active_model && + status.active_model.toLowerCase() === target.toLowerCase() + ) { + return null; + } + } catch { + // Fall through to load attempt; the backend will re-error if needed. + } + + const toastId = toast.loading(`Loading ${target}…`, { + description: "Starting the local inference server for this recipe.", + }); + try { + const isGguf = /gguf/i.test(target); + await loadModel({ + model_path: target, + hf_token: null, + max_seq_length: isGguf ? 0 : 4096, + load_in_4bit: true, + is_lora: false, + gguf_variant: null, + trust_remote_code: false, + chat_template_override: null, + cache_type_kv: null, + speculative_type: null, + }); + toast.success(`Loaded ${target}`, { id: toastId, duration: 2000 }); + return null; + } catch (error) { + toast.dismiss(toastId); + return error instanceof Error ? error.message : String(error); + } +} + type UseRecipeExecutionsParams = { recipeId: string; currentSignature: string; payloadResult: RecipePayloadResult; + initialRunRows?: number | null; onExecutionStart?: () => void; onPreviewSuccess?: () => void; }; @@ -101,6 +180,7 @@ export function useRecipeExecutions({ recipeId, currentSignature, payloadResult, + initialRunRows, onExecutionStart, onPreviewSuccess, }: UseRecipeExecutionsParams): UseRecipeExecutionsResult { @@ -181,6 +261,19 @@ export function useRecipeExecutions({ resetForRecipe(); + // Seed previewRows from the recipe's original run.rows (read from the + // loaded JSON, not the rebuilt payload (the builder hardcodes 5). + // Templates ship their own suggested preview size (e.g. GitHub Support + // Bot: 10); we honor it so users don't see a surprise 5. + if ( + typeof initialRunRows === "number" && + Number.isFinite(initialRunRows) && + initialRunRows > 0 && + initialRunRows !== 5 + ) { + setPreviewRows(Math.floor(initialRunRows)); + } + async function hydrate(): Promise { try { const records = await loadSortedRecipeExecutions(recipeId); @@ -216,10 +309,12 @@ export function useRecipeExecutions({ cancelled = true; }; }, [ + initialRunRows, onPreviewSuccess, recipeId, resetForRecipe, setExecutions, + setPreviewRows, setRunErrors, upsertAndPersist, ]); @@ -340,6 +435,20 @@ export function useRecipeExecutions({ return false; } + // Flip to the Runs pane BEFORE we run ensureLocalModelLoaded + validate. + // Validation re-crawls the seed (multiple seconds for the github_repo + // reader) and the user otherwise stares at a "Running..." button with + // nothing else changing. runExecution() later no-ops this callback if + // the view has already been flipped, so we fire it once here. + onExecutionStart?.(); + + const localLoadError = await ensureLocalModelLoaded(payload); + if (localLoadError) { + setRunErrors([localLoadError]); + toastError("Local model failed to load", localLoadError); + return false; + } + const normalizedRows = sanitizeExecutionRows(rows, kind); const executionPayload = buildExecutionPayload({ payload, @@ -374,7 +483,13 @@ export function useRecipeExecutions({ runName, }); }, - [readExecutablePayload, runExecution, runSettings, setRunErrors], + [ + onExecutionStart, + readExecutablePayload, + runExecution, + runSettings, + setRunErrors, + ], ); const runPreview = useCallback(async (): Promise => { diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts index f6f7390cf7..79c75c1885 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts @@ -112,6 +112,10 @@ export function useRecipeStudioActions({ recipeId, currentSignature: persistence.currentSignature, payloadResult, + initialRunRows: + typeof initialPayload?.run?.rows === "number" + ? initialPayload.run.rows + : null, onExecutionStart, onPreviewSuccess, }); diff --git a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx index 586d5d8ad8..5dbc004ffd 100644 --- a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx +++ b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx @@ -47,6 +47,7 @@ import { ConfigDialog } from "./dialogs/config-dialog"; import { ImportDialog } from "./dialogs/import-dialog"; import { RunDialog } from "./dialogs/preview-dialog"; import { ProcessorsDialog } from "./dialogs/processors-dialog"; +import { GithubCrawlerEasyView } from "./easy/github-crawler-easy-view"; import type { RecipeExecutionRecord, RecipeStudioView, @@ -201,7 +202,40 @@ export function RecipeStudioPage({ null, ); const flowContainerRef = useRef(null); - const [activeView, setActiveView] = useState("editor"); + const supportsEasyMode = + initialPayload?.ui?.seed_source_type === "github_repo" || + (initialPayload?.recipe?.seed_config as { source?: { seed_type?: string } } | undefined) + ?.source?.seed_type === "github_repo"; + const viewModeStorageKey = `recipe-studio:view-mode:${recipeId}`; + const [activeView, setActiveViewState] = useState(() => { + if (typeof window !== "undefined") { + const stored = window.localStorage.getItem(viewModeStorageKey); + if (stored === "easy" && supportsEasyMode) return "easy"; + if (stored === "editor" || stored === "executions") return stored; + } + return supportsEasyMode ? "easy" : "editor"; + }); + const setActiveView = useCallback( + (next: RecipeStudioView | ((prev: RecipeStudioView) => RecipeStudioView)) => { + setActiveViewState((prev) => { + const resolved = typeof next === "function" ? next(prev) : next; + if (typeof window !== "undefined") { + window.localStorage.setItem(viewModeStorageKey, resolved); + } + return resolved; + }); + }, + [viewModeStorageKey], + ); + // Easy mode has no canvas overlay/progress island, so once a run starts the + // user sees the Run button stuck on "Running..." with nothing else changing. + // Flip to the Runs pane so they land where progress is actually rendered. + // Advanced (editor) keeps its island and stays put. + const handleExecutionStart = useCallback(() => { + setActiveView((currentView) => + currentView === "easy" ? "executions" : currentView, + ); + }, [setActiveView]); const [processorsOpen, setProcessorsOpen] = useState(false); const [interactive, setInteractive] = useState(true); const [runtimeIslandMinimized, setRuntimeIslandMinimized] = useState(false); @@ -326,6 +360,8 @@ export function RecipeStudioPage({ validateResult, cancelExecution, loadExecutionDatasetPage, + runPreview, + runFull, copyRecipe, importRecipe, } = useRecipeStudioActions({ @@ -338,6 +374,7 @@ export function RecipeStudioPage({ resetRecipe, loadRecipe, getCurrentPayloadFromStore, + onExecutionStart: handleExecutionStart, }); const { activeExecution, @@ -359,6 +396,19 @@ export function RecipeStudioPage({ const runBusy = previewLoading || fullLoading || executionLocked; const islandExecution = activeExecution ?? recentCompletedExecution; + // Easy mode runs a full run (artifact persisted, tracked in Runs pane) + // using runFull. runFull requires a non-empty fullRunName but the Easy form + // has no run-name input, so seed a default here as soon as Easy is active. + // User can still rename it from Advanced/Runs dialogs before clicking Run. + useEffect(() => { + if (!supportsEasyMode) return; + if (activeView !== "easy") return; + if (fullRunName.trim()) return; + const stamp = new Date().toISOString().replace(/[:.]/g, "-").slice(0, 16); + const base = workflowName.trim() || "Easy run"; + setFullRunName(`${base} ${stamp}`); + }, [supportsEasyMode, activeView, fullRunName, workflowName, setFullRunName]); + const toggleInteractive = useCallback(() => { if (executionLocked) { return; @@ -739,6 +789,7 @@ export function RecipeStudioPage({ savedAtLabel={savedAtLabel} workflowName={workflowName} warnings={getGraphWarnings(configs, edges)} + supportsEasyMode={supportsEasyMode} onWorkflowNameChange={setWorkflowName} onViewChange={setActiveView} onSaveRecipe={() => { @@ -749,7 +800,25 @@ export function RecipeStudioPage({ className="h-[75vh] w-full rounded-t-none" ref={flowContainerRef} > - {activeView === "editor" ? ( + {activeView === "easy" ? ( + { + // Easy mode is a full run (artifact persisted, tracked in + // the Runs pane) capped at the user's row count. runFull + // requires a non-empty fullRunName; the effect below + // populates one on mount so the closure in runFull is + // already up to date by the time the user clicks Run. + void runFull(); + }} + runLoading={fullLoading || executionLocked} + runErrors={runErrors} + onSwitchToAdvanced={() => setActiveView("editor")} + /> + ) : activeView === "editor" ? ( editorContent ) : ( ((set, get) => ({ nextSourceType = "local"; } else if (type === "seed_unstructured") { nextSourceType = "unstructured"; + } else if (type === "seed_github") { + nextSourceType = "github_repo"; } - const nextConfig = { + const nextConfig: typeof existing = { ...existing, seed_source_type: nextSourceType, hf_repo_id: "", @@ -415,6 +417,12 @@ export const useRecipeStudioStore = create((set, get) => ({ seed_preview_rows: [], unstructured_chunk_size: "1200", unstructured_chunk_overlap: "200", + github_repo_slug: "", + github_token: "", + github_limit: "100", + github_item_types: ["issues", "pulls"], + github_include_comments: true, + github_max_comments_per_item: "30", }; return { configs: { diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index cd1127a4e2..9c720a06d3 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -38,7 +38,10 @@ export type LayoutDirection = "LR" | "TB"; export type SeedSamplingStrategy = "ordered" | "shuffle"; export type SeedSelectionType = "none" | "index_range" | "partition_block"; -export type SeedSourceType = "hf" | "local" | "unstructured"; +export type SeedSourceType = "hf" | "local" | "unstructured" | "github_repo"; + +export type GithubItemType = "issues" | "pulls" | "commits"; +export type GithubStateFilter = "all" | "open" | "closed"; export const INFRA_NODE_KINDS = new Set([ "model_provider", "model_config", @@ -338,6 +341,20 @@ export type SeedConfig = { unstructured_file_ids?: string[]; unstructured_file_names?: string[]; unstructured_file_sizes?: number[]; + // biome-ignore lint/style/useNamingConvention: api schema + github_repo_slug?: string; + // biome-ignore lint/style/useNamingConvention: api schema + github_token?: string; + // biome-ignore lint/style/useNamingConvention: api schema + github_limit?: string; + // biome-ignore lint/style/useNamingConvention: api schema + github_item_types?: GithubItemType[]; + // biome-ignore lint/style/useNamingConvention: api schema + github_state?: GithubStateFilter; + // biome-ignore lint/style/useNamingConvention: api schema + github_include_comments?: boolean; + // biome-ignore lint/style/useNamingConvention: api schema + github_max_comments_per_item?: string; resolved_paths?: string[]; // ui-only seed_preview_rows?: Record[]; diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts index 0d92129232..21eadb3195 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts @@ -80,6 +80,12 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { let resolved_paths: string[] = []; let unstructured_chunk_size = "1200"; let unstructured_chunk_overlap = "200"; + let github_repo_slug = ""; + let github_token = ""; + let github_limit = "100"; + let github_item_types: ("issues" | "pulls" | "commits")[] = ["issues", "pulls"]; + let github_include_comments = true; + let github_max_comments_per_item = "30"; const sourceRaw = seedConfigRaw.source; if (isRecord(sourceRaw)) { const seedType = readString(sourceRaw.seed_type); @@ -107,6 +113,26 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { unstructuredFileNames = []; unstructured_chunk_size = readNumberString(sourceRaw.chunk_size) || "1200"; unstructured_chunk_overlap = readNumberString(sourceRaw.chunk_overlap) || "200"; + } else if (seedType === "github_repo") { + seed_source_type = "github_repo"; + const rawRepos = Array.isArray(sourceRaw.repos) ? sourceRaw.repos : []; + const repos = rawRepos.filter((r): r is string => typeof r === "string"); + github_repo_slug = repos.join("\n"); + github_token = readString(sourceRaw.token) ?? ""; + github_limit = readNumberString(sourceRaw.limit) || "100"; + const rawItems = Array.isArray(sourceRaw.item_types) ? sourceRaw.item_types : []; + const validItems = rawItems.filter( + (t): t is "issues" | "pulls" | "commits" => + t === "issues" || t === "pulls" || t === "commits", + ); + if (validItems.length > 0) { + github_item_types = validItems; + } + if (typeof sourceRaw.include_comments === "boolean") { + github_include_comments = sourceRaw.include_comments; + } + github_max_comments_per_item = + readNumberString(sourceRaw.max_comments_per_item) || "30"; } } @@ -147,6 +173,12 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { resolved_paths, unstructured_chunk_size, unstructured_chunk_overlap, + github_repo_slug, + github_token, + github_limit, + github_item_types, + github_include_comments, + github_max_comments_per_item, sampling_strategy, selection_type, selection_start, diff --git a/studio/frontend/src/features/recipe-studio/utils/node-data.ts b/studio/frontend/src/features/recipe-studio/utils/node-data.ts index 2fc5205db8..96042b5919 100644 --- a/studio/frontend/src/features/recipe-studio/utils/node-data.ts +++ b/studio/frontend/src/features/recipe-studio/utils/node-data.ts @@ -70,7 +70,9 @@ export function nodeDataFromConfig( ? "Hugging Face dataset" : seedSourceType === "local" ? "CSV or JSON file" - : "Document file"; + : seedSourceType === "github_repo" + ? "GitHub repositories" + : "Document file"; return { title: "Source data", kind: "seed", diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts index f73a2b46da..bb48b43857 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts @@ -6,6 +6,7 @@ import type { NodeConfig, SeedConfig } from "../../types"; const DEFAULT_CHUNK_SIZE = 1200; const DEFAULT_CHUNK_OVERLAP = 200; const MAX_CHUNK_SIZE = 20000; +const GITHUB_ITEM_TYPES = new Set(["issues", "pulls", "commits"]); function parseIntStrict(value: string | undefined): number | null { const trimmed = value?.trim(); @@ -56,33 +57,86 @@ export function buildSeedConfig( selectionStrategy = { index, num_partitions: numPartitions }; } - const source = - seedSourceType === "hf" - ? { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "hf", - path, - token, - endpoint, - } - : seedSourceType === "unstructured" - ? (() => { - const { chunkSize, chunkOverlap } = resolveChunking(config); - return { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "unstructured", - paths: config.resolved_paths?.length ? config.resolved_paths : [config.hf_path], - // biome-ignore lint/style/useNamingConvention: api schema - chunk_size: chunkSize, - // biome-ignore lint/style/useNamingConvention: api schema - chunk_overlap: chunkOverlap, - }; - })() - : { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "local", - path, - }; + let source: Record; + if (seedSourceType === "hf") { + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "hf", + path, + token, + endpoint, + }; + } else if (seedSourceType === "unstructured") { + const { chunkSize, chunkOverlap } = resolveChunking(config); + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "unstructured", + paths: config.resolved_paths?.length ? config.resolved_paths : [config.hf_path], + // biome-ignore lint/style/useNamingConvention: api schema + chunk_size: chunkSize, + // biome-ignore lint/style/useNamingConvention: api schema + chunk_overlap: chunkOverlap, + }; + } else if (seedSourceType === "github_repo") { + const repos = (config.github_repo_slug ?? "") + .split(/[\n,]/) + .map((r) => r.trim()) + .filter(Boolean); + if (repos.length === 0) { + errors.push(`Seed ${config.name}: at least one repo is required.`); + return undefined; + } + const invalidRepo = repos.find((repo) => { + const parts = repo.split("/"); + return parts.length !== 2 || parts.some((part) => !part); + }); + if (invalidRepo) { + errors.push( + `Seed ${config.name}: GitHub repositories must use owner/name format.`, + ); + return undefined; + } + const itemTypes = config.github_item_types?.length + ? config.github_item_types + : ["issues", "pulls"]; + if (itemTypes.some((itemType) => !GITHUB_ITEM_TYPES.has(itemType))) { + errors.push(`Seed ${config.name}: GitHub item types invalid.`); + return undefined; + } + const limitNum = parseIntStrict(config.github_limit ?? "100"); + if (limitNum === null || limitNum < 1 || limitNum > 5000) { + errors.push( + `Seed ${config.name}: GitHub items per repo must be an integer from 1 to 5000.`, + ); + return undefined; + } + const maxCommentsNum = parseIntStrict(config.github_max_comments_per_item ?? "30"); + if (maxCommentsNum === null || maxCommentsNum < 0 || maxCommentsNum > 200) { + errors.push( + `Seed ${config.name}: GitHub max comments per item must be an integer from 0 to 200.`, + ); + return undefined; + } + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "github_repo", + repos, + token: (config.github_token ?? "").trim(), + // biome-ignore lint/style/useNamingConvention: api schema + item_types: itemTypes, + limit: limitNum, + // biome-ignore lint/style/useNamingConvention: api schema + include_comments: config.github_include_comments ?? true, + // biome-ignore lint/style/useNamingConvention: api schema + max_comments_per_item: maxCommentsNum, + }; + } else { + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "local", + path, + }; + } return { source, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts index b8b1aba93b..902ea796b4 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts @@ -57,7 +57,7 @@ export type RecipePayload = { // ui-only: graph orientation layout_direction?: "LR" | "TB"; // ui-only, used to preserve seed block mode across imports/refresh - seed_source_type?: "hf" | "local" | "unstructured"; + seed_source_type?: "hf" | "local" | "unstructured" | "github_repo"; // ui-only, persisted aux node positions by llm name + aux key aux_nodes?: Array<{ llm: string; diff --git a/studio/frontend/src/features/recipe-studio/utils/validation.ts b/studio/frontend/src/features/recipe-studio/utils/validation.ts index 5f48fd44a5..d5200eeed0 100644 --- a/studio/frontend/src/features/recipe-studio/utils/validation.ts +++ b/studio/frontend/src/features/recipe-studio/utils/validation.ts @@ -8,6 +8,7 @@ import { isOxcCodeShape } from "./validators/oxc-code-shape"; import { isOxcValidationMode } from "./validators/oxc-mode"; const TRACE_MODES = new Set(["none", "last_message", "all_messages"]); +const GITHUB_ITEM_TYPES = new Set(["issues", "pulls", "commits"]); // biome-ignore lint/complexity/noExcessiveCognitiveComplexity: validation rules export function getConfigErrors(config: NodeConfig | null): string[] { @@ -268,15 +269,49 @@ export function getConfigErrors(config: NodeConfig | null): string[] { } if (config.kind === "seed") { const seedSourceType = config.seed_source_type ?? "hf"; - if (seedSourceType === "hf" && !config.hf_repo_id.trim()) { + if (seedSourceType === "github_repo") { + const repos = (config.github_repo_slug ?? "") + .split(/[\n,]/) + .map((repo) => repo.trim()) + .filter(Boolean); + if (repos.length === 0) { + errors.push("Add at least one GitHub repository."); + } + if ( + repos.some((repo) => { + const parts = repo.split("/"); + return parts.length !== 2 || parts.some((part) => !part); + }) + ) { + errors.push("GitHub repositories must use owner/name format."); + } + const itemTypes = config.github_item_types?.length + ? config.github_item_types + : ["issues", "pulls"]; + if (itemTypes.length === 0) { + errors.push("Choose at least one GitHub item type."); + } else if (itemTypes.some((itemType) => !GITHUB_ITEM_TYPES.has(itemType))) { + errors.push("GitHub item types must be issues, pulls, or commits."); + } + const limit = parseIntNumber(config.github_limit ?? "100"); + if (limit === null || limit < 1 || limit > 5000) { + errors.push("Items per repo must be an integer from 1 to 5000."); + } + const maxComments = parseIntNumber(config.github_max_comments_per_item ?? "30"); + if (maxComments === null || maxComments < 0 || maxComments > 200) { + errors.push("Max comments per item must be an integer from 0 to 200."); + } + } else if (seedSourceType === "hf" && !config.hf_repo_id.trim()) { errors.push("Choose a Hugging Face dataset."); } - const hasPath = - seedSourceType === "unstructured" - ? (config.resolved_paths?.length ?? 0) > 0 - : Boolean(config.hf_path.trim()); - if (!hasPath) { - errors.push("Load the source-data preview first."); + if (seedSourceType !== "github_repo") { + const hasPath = + seedSourceType === "unstructured" + ? (config.resolved_paths?.length ?? 0) > 0 + : Boolean(config.hf_path.trim()); + if (!hasPath) { + errors.push("Load the source-data preview first."); + } } if ( seedSourceType === "hf" && @@ -304,7 +339,7 @@ export function getConfigErrors(config: NodeConfig | null): string[] { ) { errors.push("Chunk overlap must be less than chunk size."); } - } else { + } else if (seedSourceType !== "github_repo") { const selectedDropColumns = (config.seed_drop_columns ?? []) .map((value) => value.trim()) .filter(Boolean); diff --git a/studio/frontend/src/lib/api-base.ts b/studio/frontend/src/lib/api-base.ts index 664b9e58bd..c8be442653 100644 --- a/studio/frontend/src/lib/api-base.ts +++ b/studio/frontend/src/lib/api-base.ts @@ -8,6 +8,12 @@ if (isTauri && !isViteDev) { apiBase = 'http://127.0.0.1:8888' } +const initialApiBase = apiBase + +export function resetApiBase() { + apiBase = initialApiBase +} + export function setApiBase(port: number) { apiBase = `http://127.0.0.1:${port}` } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index b56b737cfa..5fe4415970 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -418,6 +418,9 @@ CONSTRAINTS = SINGLE_ENV / "constraints.txt" LOCAL_DD_UNSTRUCTURED_PLUGIN = ( SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed" ) +LOCAL_DD_GITHUB_PLUGIN = ( + SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed" +) # -- Unicode-safe printing --------------------------------------------- # On Windows the default console encoding can be a legacy code page @@ -1135,22 +1138,28 @@ def install_python_stack() -> int: req = SINGLE_ENV / "data-designer.txt", ) - # 11. Local Data Designer seed plugin - if not LOCAL_DD_UNSTRUCTURED_PLUGIN.is_dir(): - _safe_print( - _red( - f"❌ Missing local plugin directory: {LOCAL_DD_UNSTRUCTURED_PLUGIN}", - ), - ) - return 1 + # 11. Local Data Designer seed plugins + local_dd_plugins = [ + ("unstructured", LOCAL_DD_UNSTRUCTURED_PLUGIN), + ("github", LOCAL_DD_GITHUB_PLUGIN), + ] + for _plugin_name, plugin_dir in local_dd_plugins: + if not plugin_dir.is_dir(): + _safe_print( + _red( + f"❌ Missing local plugin directory: {plugin_dir}", + ), + ) + return 1 _progress("local plugin") - pip_install( - "Installing local data-designer unstructured plugin", - "--no-cache-dir", - "--no-deps", - str(LOCAL_DD_UNSTRUCTURED_PLUGIN), - constrain = False, - ) + for plugin_name, plugin_dir in local_dd_plugins: + pip_install( + f"Installing local data-designer {plugin_name} plugin", + "--no-cache-dir", + "--no-deps", + str(plugin_dir), + constrain = False, + ) # 12. Patch metadata for single-env compatibility _progress("finalizing") From efed5c37394a144349cd9b1ea525e132e04584e5 Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Sat, 25 Apr 2026 16:49:05 -0700 Subject: [PATCH 10/54] fix(studio): use endswith for mmproj F16 variant selection (#5184) "f16" in filename matched BF16 files because "bf16" contains "f16" as a substring. Switch to endswith("-f16.gguf") for an exact match. --- studio/backend/core/inference/llama_cpp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 33722f103c..c1a37e2fa4 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1339,7 +1339,7 @@ class LlamaCppBackend: # Prefer F16 variant target = None for f in mmproj_files: - if "f16" in f.lower(): + if f.lower().endswith("-f16.gguf"): target = f break if target is None: From df3a205726d231bd12402960a3cd3a90b85c7fc7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Apr 2026 14:17:03 -0700 Subject: [PATCH 11/54] [pre-commit.ci] pre-commit autoupdate (#5204) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit updates: - [github.com/astral-sh/ruff-pre-commit: v0.15.11 → v0.15.12](https://github.com/astral-sh/ruff-pre-commit/compare/v0.15.11...v0.15.12) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1d3137b373..a2a4995d62 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.11 + rev: v0.15.12 hooks: - id: ruff args: From daf0889804aa3f33bbf2bde74034cad926f4a3bf Mon Sep 17 00:00:00 2001 From: Etherll <61019402+Etherll@users.noreply.github.com> Date: Tue, 28 Apr 2026 11:10:47 +0300 Subject: [PATCH 12/54] Fix Windows install when paths contain spaces or Python 3.14 is on PATH (#5201) * fix(studio): use py.exe to detect supported Python on Windows Description: The previous detection looked at `python --version` on PATH and hard-failed if the resolved Python wasn't 3.11-3.13. On systems where Python 3.14 sits ahead of 3.13 in PATH order, this aborted the installer even though a supported interpreter was installed. Prefer the py.exe launcher and probe `py -3.13`, `py -3.12`, `py -3.11` in turn. Fall back to `python --version` only when py.exe is absent, and surface a clearer error when no supported version can be found via either path. * Studio: consolidate Windows studio overlay into single Tauri-gated block Replace the in-file sentinel hotfix and the unconditional file-copy overlay with a single block gated on $TauriMode. Hash-compare makes re-runs no-ops, removing the sentinel-clobbering bug that occurred when the second copy path overwrote the marker without re-adding it. Non-Tauri --local installs no longer need a copy overlay: the editable install above (uv pip install -e $RepoRoot --no-deps) makes _PACKAGE_ROOT in unsloth_cli/commands/studio.py resolve to the repo source tree via PEP 660 __file__-relative resolution, so `unsloth studio setup` finds the local setup.ps1 and install_python_stack.py without any file copying. Plain PyPI installs invoked from a checked-out repo directory are also no longer silently overlaid from cwd. * fix(studio): work around uv space-in-path truncation on Windows uv 0.11.x truncates `-c ` and `-r ` arguments at the first space, breaking installs on Windows when the venv or repo sits under a path containing spaces (e.g. C:\Users\First Last\...). Pass paths through GetShortPathNameW to convert to 8.3 short form before handing them to uv. Plain pip is unaffected and keeps the original long path. No-op on Linux/Mac (gated on IS_WINDOWS and on the path actually containing a space). * Refactor Python stack overlay logic in install.ps1 Refactor overlay logic for Python stack installation and improve handling of missing target directories. * Update Python installation logic in setup.ps1 --- install.ps1 | 60 ++++++++++++++++++++------------- studio/install_python_stack.py | 47 ++++++++++++++++++++------ studio/setup.ps1 | 61 +++++++++++++++++++++++++++------- 3 files changed, 123 insertions(+), 45 deletions(-) diff --git a/install.ps1 b/install.ps1 index 77c0034125..2544eaf78a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1027,32 +1027,46 @@ shell.Run cmd, 0, False } } - # Hotfix: patch install_python_stack.py for Windows GUI stdout - # The PyPI version crashes with OSError when stdout is piped from a GUI app. - # Copy our fixed version (bundled by Tauri) over the installed one. - # Remove this block once PyPI ships the fix from commit 18c5aae7. + # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped + # for --local: the editable install above already makes _PACKAGE_ROOT in + # unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__). + # Source paths match the Tauri bundle layout in studio/src-tauri/tauri.conf.json, + # which bundles install_python_stack.py at the bundle root next to install.ps1. if ($TauriMode) { $rawPath = if ($PSCommandPath) { $PSCommandPath } else { $MyInvocation.ScriptName } - $scriptDir = Split-Path -Parent ($rawPath -replace '^\\\\\?\\', '') - $fixedPy = Join-Path $scriptDir "install_python_stack.py" - $target = Join-Path $VenvDir "Lib\site-packages\studio\install_python_stack.py" - $sentinel = "# UNSLOTH_DESKTOP_HOTFIX_APPLIED_v1" - $sentinelPattern = [regex]::Escape($sentinel) - if ((Test-Path $fixedPy) -and (Test-Path $target)) { - $installed = Get-Content $target -Raw - if ($installed -notmatch $sentinelPattern) { - Copy-Item $fixedPy $target -Force - Add-Content -Path $target -Value "`n$sentinel" - substep "patched install_python_stack.py (stdout fix)" - } else { - substep "install_python_stack.py already has stdout fix" + if ($rawPath) { + # Strip leading \\?\ extended-length prefix if the launcher passed one. + $scriptDir = Split-Path -Parent ($rawPath -replace '^\\\\\?\\', '') + $overlayMap = [ordered]@{ + "install_python_stack.py" = "Lib\site-packages\studio\install_python_stack.py" + } + foreach ($rel in $overlayMap.Keys) { + $src = Join-Path $scriptDir $rel + $dst = Join-Path $VenvDir $overlayMap[$rel] + if (-not (Test-Path $src)) { continue } + $dstParent = Split-Path -Parent $dst + if (-not (Test-Path $dstParent)) { + Write-Host "[WARN] Overlay target dir missing: $dstParent; studio setup may use stale bundled file" -ForegroundColor Yellow + continue + } + try { + if (-not (Test-Path $dst)) { + # Backfill: target file missing but parent dir exists. + Copy-Item $src $dst -Force + substep ("backfilled bundled " + (Split-Path -Leaf $rel)) + } else { + # Hash-compare so re-runs are no-ops when files already match. + $srcHash = (Get-FileHash $src -Algorithm SHA256).Hash + $dstHash = (Get-FileHash $dst -Algorithm SHA256).Hash + if ($srcHash -ne $dstHash) { + Copy-Item $src $dst -Force + substep ("applied bundled " + (Split-Path -Leaf $rel)) + } + } + } catch { + Write-Host "[WARN] Could not overlay $($rel): $($_.Exception.Message); studio setup may use stale bundled file" -ForegroundColor Yellow + } } - } elseif ((Test-Path $fixedPy) -and (Test-Path (Split-Path $target))) { - Copy-Item $fixedPy $target -Force - Add-Content -Path $target -Value "`n$sentinel" - substep "patched install_python_stack.py (stdout fix)" - } else { - Write-Host "[WARN] Could not patch install_python_stack.py (bundled file or target dir missing)" -ForegroundColor Yellow } } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 5fe4415970..8ac59d389d 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -359,6 +359,27 @@ def _ensure_rocm_torch() -> None: ) +def _uv_safe_path(path: object) -> str: + # uv 0.11.x: `-c ` truncates at the space; use 8.3 short form. + s = str(path) + if not IS_WINDOWS or " " not in s: + return s + try: + import ctypes + from ctypes import wintypes + + get_short = ctypes.windll.kernel32.GetShortPathNameW + get_short.argtypes = [wintypes.LPCWSTR, wintypes.LPWSTR, wintypes.DWORD] + get_short.restype = wintypes.DWORD + buf = ctypes.create_unicode_buffer(32768) + rc = get_short(s, buf, 32768) + if 0 < rc < 32768 and " " not in buf.value: + return buf.value + except Exception: + pass + return s + + def _windows_hidden_subprocess_kwargs() -> dict[str, object]: """Return Windows-only subprocess kwargs that suppress console windows.""" if not IS_WINDOWS: @@ -751,14 +772,16 @@ def pip_install_try( """Like pip_install but returns False on failure instead of exiting. For optional installs with a follow-up fallback. """ - constraint_args: list[str] = [] + constraint_args_pip: list[str] = [] + constraint_args_uv: list[str] = [] if constrain and CONSTRAINTS.is_file(): - constraint_args = ["-c", str(CONSTRAINTS)] + constraint_args_pip = ["-c", str(CONSTRAINTS)] + constraint_args_uv = ["-c", _uv_safe_path(CONSTRAINTS)] if USE_UV: - cmd = _build_uv_cmd(args) + constraint_args + cmd = _build_uv_cmd(args) + constraint_args_uv else: - cmd = _build_pip_cmd(args) + constraint_args + cmd = _build_pip_cmd(args) + constraint_args_pip if VERBOSE: _step(_LABEL, f"{label}...", _dim) @@ -781,9 +804,11 @@ def pip_install( constrain: bool = True, ) -> None: """Build and run a pip install command (uses uv when available, falls back to pip).""" - constraint_args: list[str] = [] + constraint_args_pip: list[str] = [] + constraint_args_uv: list[str] = [] if constrain and CONSTRAINTS.is_file(): - constraint_args = ["-c", str(CONSTRAINTS)] + constraint_args_pip = ["-c", str(CONSTRAINTS)] + constraint_args_uv = ["-c", _uv_safe_path(CONSTRAINTS)] actual_req = req temp_reqs: list[Path] = [] @@ -793,13 +818,15 @@ def pip_install( if actual_req is not None and NO_TORCH and NO_TORCH_SKIP_PACKAGES: actual_req = _filter_requirements(actual_req, NO_TORCH_SKIP_PACKAGES) temp_reqs.append(actual_req) - req_args: list[str] = [] + req_args_pip: list[str] = [] + req_args_uv: list[str] = [] if actual_req is not None: - req_args = ["-r", str(actual_req)] + req_args_pip = ["-r", str(actual_req)] + req_args_uv = ["-r", _uv_safe_path(actual_req)] try: if USE_UV: - uv_cmd = _build_uv_cmd(args) + constraint_args + req_args + uv_cmd = _build_uv_cmd(args) + constraint_args_uv + req_args_uv if VERBOSE: print(f" {label}...") result = subprocess.run( @@ -814,7 +841,7 @@ def pip_install( if result.stdout: print(result.stdout.decode(errors = "replace")) - pip_cmd = _build_pip_cmd(args) + constraint_args + req_args + pip_cmd = _build_pip_cmd(args) + constraint_args_pip + req_args_pip run(f"{label} (pip)" if USE_UV else label, pip_cmd) finally: for temp_req in temp_reqs: diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 24ef3ef1eb..0da30c9229 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1145,28 +1145,58 @@ if ($IsPipInstall) { } } -# ============================================ -# 1g. Python (>= 3.11 and < 3.14, matching setup.sh) -# ============================================ +# 1g. Python (>= 3.11 and < 3.14). Prefer py.exe so a 3.14 ahead of 3.13 on PATH does not trip the gate. $HasPython = $null -ne (Get-Command python -ErrorAction SilentlyContinue) +$PyLauncher = Get-Command py -CommandType Application -ErrorAction SilentlyContinue $PythonOk = $false +$DetectedPyVer = $null -if ($HasPython) { +if ($PyLauncher) { + foreach ($minor in @("3.13", "3.12", "3.11")) { + try { + $out = & $PyLauncher.Source "-$minor" --version 2>&1 | Out-String + if ($out -match 'Python (3\.\d+\.\d+)') { + $DetectedPyVer = $Matches[1] + # Make `python` resolvable for the rest of setup. Without this, + # py-launcher-only installs (no python.exe on PATH) pass the gate + # and then crash on the first bare `python` call below. + try { + $resolvedExe = (& $PyLauncher.Source "-$minor" -c "import sys; print(sys.executable)" 2>$null | Select-Object -First 1) + if ($resolvedExe -and (Test-Path $resolvedExe)) { + $resolvedDir = Split-Path -Parent $resolvedExe + $alreadyOnPath = ($env:PATH -split ';' | Where-Object { $_.TrimEnd('\') -ieq $resolvedDir.TrimEnd('\') }).Count -gt 0 + if (-not $alreadyOnPath) { + $env:PATH = "$resolvedDir;$env:PATH" + } + $HasPython = $true + } + } catch { } + $PythonOk = $true + break + } + } catch { } + } +} + +if (-not $PythonOk -and $HasPython) { $PyVer = python --version 2>&1 if ($PyVer -match "(\d+)\.(\d+)") { $PyMajor = [int]$Matches[1]; $PyMinor = [int]$Matches[2] if ($PyMajor -eq 3 -and $PyMinor -ge 11 -and $PyMinor -lt 14) { - substep "Python $PyVer" + $DetectedPyVer = "$PyMajor.$PyMinor" $PythonOk = $true - } else { - Write-Host "[ERROR] Python $PyVer is outside supported range (need >= 3.11 and < 3.14)." -ForegroundColor Red - Write-Host " Install Python 3.12 from https://python.org/downloads/" -ForegroundColor Yellow - exit 1 } } -} else { - # No Python at all -- install 3.12 - Write-Host "Python not found -- installing Python 3.12 via winget..." -ForegroundColor Yellow +} + +if ($PythonOk) { + substep "Python $DetectedPyVer" +} elseif (-not $HasPython) { + # No `python` on PATH (and py.exe either absent or only had unsupported + # minors). Try winget as before -- gating on $HasPython alone, not also + # on $PyLauncher, so a launcher-only install with just 3.14 still gets + # an automatic 3.12 install instead of a hard error. + Write-Host "Python 3.11-3.13 not found -- installing Python 3.12 via winget..." -ForegroundColor Yellow $HasWinget = $null -ne (Get-Command winget -ErrorAction SilentlyContinue) if ($HasWinget) { winget install -e --id Python.Python.3.12 --source winget --accept-package-agreements --accept-source-agreements @@ -1180,6 +1210,13 @@ if ($HasPython) { } step "python" "$(python --version 2>&1)" $PythonOk = $true +} else { + # python.exe is on PATH but its version is unsupported, and py.exe (if + # present) had no supported minor either. + Write-Host "[ERROR] No supported Python (3.11-3.13) found on this system." -ForegroundColor Red + Write-Host " py.exe could not locate -3.11/-3.12/-3.13 and `python` on PATH is unsupported." -ForegroundColor Yellow + Write-Host " Install Python 3.12 from https://python.org/downloads/" -ForegroundColor Yellow + exit 1 } # Add user-scheme Python Scripts dir to PATH (nt_user only, no venv fallback). From 230d58872d0fa44279f6358cd9c5116bcc2cd3fe Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Tue, 28 Apr 2026 10:39:48 +0100 Subject: [PATCH 13/54] Studio: Preserve transparency in uploaded profile avatars (#5200) * fix: preserve transparency in uploaded profile avatars * fix: guard unsupported canvas mime fallback --- .../profile/utils/resize-image-file.ts | 50 +++++++++++++++---- 1 file changed, 40 insertions(+), 10 deletions(-) diff --git a/studio/frontend/src/features/profile/utils/resize-image-file.ts b/studio/frontend/src/features/profile/utils/resize-image-file.ts index 3f829ba975..c420497395 100644 --- a/studio/frontend/src/features/profile/utils/resize-image-file.ts +++ b/studio/frontend/src/features/profile/utils/resize-image-file.ts @@ -3,6 +3,11 @@ const MAX_EDGE = 256; const MAX_BYTES = 380_000; +const MAX_DATA_URL_LENGTH = Math.floor(MAX_BYTES * 1.35); +const JPEG_QUALITY_START = 0.88; +const WEBP_QUALITY_START = 0.9; +const MIN_QUALITY = 0.45; +const QUALITY_STEP = 0.08; function loadImage(file: File): Promise { return new Promise((resolve, reject) => { @@ -20,8 +25,33 @@ function loadImage(file: File): Promise { }); } +function canvasHasTransparency(ctx: CanvasRenderingContext2D, width: number, height: number): boolean { + const { data } = ctx.getImageData(0, 0, width, height); + for (let i = 3; i < data.length; i += 4) { + if (data[i] < 255) return true; + } + return false; +} + +function encodeCanvasWithinLimit( + canvas: HTMLCanvasElement, + mimeType: string, + startQuality: number, +): string | null { + let quality = startQuality; + let dataUrl = canvas.toDataURL(mimeType, quality); + if (!dataUrl.startsWith(`data:${mimeType}`)) return null; + + while (dataUrl.length > MAX_DATA_URL_LENGTH && quality > MIN_QUALITY) { + quality -= QUALITY_STEP; + dataUrl = canvas.toDataURL(mimeType, quality); + } + + return dataUrl.length <= MAX_DATA_URL_LENGTH ? dataUrl : null; +} + /** - * Downscale and re-encode as JPEG so localStorage stays within reasonable size. + * Downscale the image and keep transparency when present while staying localStorage-friendly. */ export async function resizeImageFileToDataUrl(file: File): Promise { const img = await loadImage(file); @@ -40,14 +70,14 @@ export async function resizeImageFileToDataUrl(file: File): Promise { if (!ctx) throw new Error("Canvas not available"); ctx.drawImage(img, 0, 0, cw, ch); - let quality = 0.88; - let dataUrl = canvas.toDataURL("image/jpeg", quality); - while (dataUrl.length > MAX_BYTES * 1.35 && quality > 0.45) { - quality -= 0.08; - dataUrl = canvas.toDataURL("image/jpeg", quality); + const hasTransparency = canvasHasTransparency(ctx, cw, ch); + if (hasTransparency) { + const webpDataUrl = encodeCanvasWithinLimit(canvas, "image/webp", WEBP_QUALITY_START); + if (webpDataUrl) return webpDataUrl; } - if (dataUrl.length > MAX_BYTES * 1.35) { - throw new Error("Image is still too large after compression. Try a smaller file."); - } - return dataUrl; + + const jpegDataUrl = encodeCanvasWithinLimit(canvas, "image/jpeg", JPEG_QUALITY_START); + if (jpegDataUrl) return jpegDataUrl; + + throw new Error("Image is still too large after compression. Try a smaller file."); } From 2469ac885bb0450532d6c76f7ff1a24d484a9b81 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Tue, 28 Apr 2026 10:39:59 +0100 Subject: [PATCH 14/54] UX: single chat header error placement and selector alignment (#5173) * UX: single chat header error placement and selector alignment * fix: centre model dropdown chevron * fix: revert view mode single * fix: allow model selector label to truncate in narrow headers --- .../assistant-ui/model-selector.tsx | 28 +++++++++++-------- .../frontend/src/features/chat/chat-page.tsx | 17 +++++++---- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 3628252177..679034c9fa 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -66,7 +66,7 @@ function ModelSelectorTrigger({ type="button" data-tour={dataTour} className={cn( - "flex items-center gap-2 transition-colors", + "flex min-w-0 items-center gap-2 transition-colors", variant === "outline" && "rounded-[8px] border border-border/60 hover:bg-[#ececec] dark:hover:bg-[#2e3035]", variant === "ghost" && "rounded-[8px] hover:bg-[#ececec] dark:hover:bg-[#2e3035]", @@ -80,17 +80,23 @@ function ModelSelectorTrigger({ {isLoaded && ( )} - - {currentModel?.name ?? "Select model"} + + + {currentModel?.name ?? "Select model"} + + {currentModel?.description && ( + + {currentModel.description} + + )} + + + - {currentModel?.description && ( - {currentModel.description} - )} - ); diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 9987af230b..e3239f7cc1 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -886,7 +886,7 @@ export function ChatPage(): ReactElement { onOpenChange={handleModelSelectorOpenChange} triggerDataTour="chat-model-selector" contentDataTour="chat-model-selector-popover" - className="max-w-[62vw] sm:max-w-none !h-[34px]" + className="max-w-[62vw] !pr-3 sm:max-w-none !h-[34px]" /> )} {loadingModel && loadToastDismissed ? ( @@ -910,12 +910,17 @@ export function ChatPage(): ReactElement { onStop={cancelLoading} /> ) : null} + {!loadingModel && modelsError ? ( +
+ {modelsError} +
+ ) : null}
- {modelsError && ( -
- {modelsError} -
- )}
{view.mode === "single" && ggufContextLength && contextUsage ? ( Date: Tue, 28 Apr 2026 10:40:15 +0100 Subject: [PATCH 15/54] Studio: Refine chat preset and group built-in presets (#5159) * UX: Refine chat preset and group built-in presets * fix: reuse built-in preset names and unify GGUF state reads * fix: built-in chat preset save and refresh behavior * Add chat preset invariant tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: decouple chat presets from model-specific settings Limit chat preset compare/apply/save behavior to temperature, topP, topK, minP, repetitionPenalty, presencePenalty, maxTokens, and systemPrompt. Preserve legacy stored preset data on load for backwards compatibility, but stop treating model-specific settings such as checkpoint, trustRemoteCode, and maxSeqLength as part of preset identity. Also align legacy prompt migration dedupe with the new preset semantics and add invariant coverage for preset-owned config comparisons. * fix: detect built-in preset edits from param changes * fix: correct built-in preset dirty state and speculative select values * fix: preserve default preset sync and keep qwen think pristine --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../frontend/src/features/chat/chat-page.tsx | 4 - .../src/features/chat/chat-settings-sheet.tsx | 738 ++++++++---------- .../chat/hooks/use-chat-model-runtime.ts | 111 +-- .../features/chat/presets/preset-policy.ts | 351 +++++++++ .../chat/stores/chat-runtime-store.ts | 29 + .../test_chat_preset_builtin_invariants.py | 272 +++++++ 6 files changed, 1040 insertions(+), 465 deletions(-) create mode 100644 studio/frontend/src/features/chat/presets/preset-policy.ts create mode 100644 tests/studio/test_chat_preset_builtin_invariants.py diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index e3239f7cc1..39683661c7 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -484,10 +484,6 @@ export function ChatPage(): ReactElement { const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen); const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen); - useEffect(() => { - return () => setSettingsOpen(false); - }, [setSettingsOpen]); - useEffect(() => { const threadId = search.thread; if (!threadId) return; diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index fc5f097969..08c5ef4080 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -19,6 +19,7 @@ import { DropdownMenu, DropdownMenuContent, DropdownMenuItem, + DropdownMenuSeparator, DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; import { Input } from "@/components/ui/input"; @@ -64,54 +65,36 @@ import { } from "@/components/ui/tooltip"; import { Tooltip as TooltipPrimitive } from "radix-ui"; import { AnimatePresence, motion } from "motion/react"; -import type { ReactNode } from "react"; +import { Fragment, type ReactNode } from "react"; import { useEffect, useLayoutEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { - DEFAULT_INFERENCE_PARAMS, - type InferenceParams, -} from "./types/runtime"; + applyPresetParams, + BUILTIN_PRESET_NAMES, + BUILTIN_PRESETS, + defaultInferenceParams, + getBuiltinVariantName, + getOrderedPresets, + getPresetOwnedConfigKey, + getPresetSaveState, + getPresetSource, + getUniquePresetName, + isSamePresetConfig, + normalizeCustomPresets, + toPresetParams, + type Preset, +} from "./presets/preset-policy"; +import type { InferenceParams } from "./types/runtime"; -export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS; +export { defaultInferenceParams, type Preset } from "./presets/preset-policy"; export type { InferenceParams } from "./types/runtime"; -export interface Preset { - name: string; - params: InferenceParams; -} - interface LegacySystemPromptTemplate { name: string; content: string; } -const BUILTIN_PRESETS: Preset[] = [ - { name: "Default", params: { ...defaultInferenceParams } }, - { - name: "Creative", - params: { - ...defaultInferenceParams, - temperature: 1.5, - topP: 1.0, - topK: 0, - minP: 0.1, - repetitionPenalty: 1.0, - }, - }, - { - name: "Precise", - params: { - ...defaultInferenceParams, - temperature: 0.1, - topP: 0.95, - topK: 80, - minP: 0.01, - repetitionPenalty: 1.0, - }, - }, -]; - const CHAT_PRESETS_KEY = "unsloth_chat_custom_presets"; const CHAT_ACTIVE_PRESET_KEY = "unsloth_chat_active_preset"; const LEGACY_CHAT_SYSTEM_PROMPTS_KEY = "unsloth_chat_system_prompts"; @@ -122,16 +105,13 @@ function canUseStorage(): boolean { return typeof window !== "undefined"; } -function getUniquePresetName(baseName: string, usedNames: Set): string { - const normalizedBase = baseName.trim() || "Imported Prompt"; - let nextName = normalizedBase; - let suffix = 2; - while (usedNames.has(nextName)) { - nextName = `${normalizedBase} ${suffix}`; - suffix += 1; +function saveCustomPresets(presets: Preset[]): void { + if (!canUseStorage()) return; + try { + localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(presets)); + } catch { + // ignore } - usedNames.add(nextName); - return nextName; } function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] { @@ -161,18 +141,7 @@ function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] { ]); const seenImportedConfigKeys = new Set( [...BUILTIN_PRESETS, ...presets].map((preset) => - JSON.stringify({ - temperature: preset.params.temperature, - topP: preset.params.topP, - topK: preset.params.topK, - minP: preset.params.minP, - repetitionPenalty: preset.params.repetitionPenalty, - presencePenalty: preset.params.presencePenalty, - maxSeqLength: preset.params.maxSeqLength, - maxTokens: preset.params.maxTokens, - systemPrompt: preset.params.systemPrompt, - trustRemoteCode: preset.params.trustRemoteCode ?? false, - }), + getPresetOwnedConfigKey(preset.params), ), ); const importedPresets = parsed @@ -191,18 +160,7 @@ function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] { }, })) .filter(({ importedParams }) => { - const configKey = JSON.stringify({ - temperature: importedParams.temperature, - topP: importedParams.topP, - topK: importedParams.topK, - minP: importedParams.minP, - repetitionPenalty: importedParams.repetitionPenalty, - presencePenalty: importedParams.presencePenalty, - maxSeqLength: importedParams.maxSeqLength, - maxTokens: importedParams.maxTokens, - systemPrompt: importedParams.systemPrompt, - trustRemoteCode: importedParams.trustRemoteCode ?? false, - }); + const configKey = getPresetOwnedConfigKey(importedParams); if (seenImportedConfigKeys.has(configKey)) return false; seenImportedConfigKeys.add(configKey); return true; @@ -216,8 +174,8 @@ function migrateLegacySystemPromptTemplates(presets: Preset[]): Preset[] { localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw); return presets; } - const mergedPresets = [...presets, ...importedPresets]; - localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(mergedPresets)); + const mergedPresets = normalizeCustomPresets([...presets, ...importedPresets]); + saveCustomPresets(mergedPresets); try { localStorage.setItem(LEGACY_CHAT_SYSTEM_PROMPTS_MIGRATED_KEY, raw); localStorage.removeItem(LEGACY_CHAT_SYSTEM_PROMPTS_KEY); @@ -255,7 +213,11 @@ function loadSavedCustomPresets(): Preset[] { }, })) .filter((preset) => preset.name.length > 0); - return migrateLegacySystemPromptTemplates(presets); + const normalized = normalizeCustomPresets(presets); + if (JSON.stringify(normalized) !== JSON.stringify(presets)) { + saveCustomPresets(normalized); + } + return migrateLegacySystemPromptTemplates(normalized); } catch { return migrateLegacySystemPromptTemplates([]); } @@ -270,82 +232,6 @@ function loadSavedActivePreset(): string { } } -type PresetSaveMode = - | "disabled" - | "overwrite-active" - | "overwrite-other" - | "create"; - -interface PresetSaveState { - mode: PresetSaveMode; - canSubmit: boolean; - isSaveReady: boolean; - buttonLabel: string; - title: string; -} - -function isSamePresetConfig(a: InferenceParams, b: InferenceParams): boolean { - return ( - a.temperature === b.temperature && - a.topP === b.topP && - a.topK === b.topK && - a.minP === b.minP && - a.repetitionPenalty === b.repetitionPenalty && - a.presencePenalty === b.presencePenalty && - a.maxSeqLength === b.maxSeqLength && - a.maxTokens === b.maxTokens && - a.systemPrompt === b.systemPrompt && - (a.trustRemoteCode ?? false) === (b.trustRemoteCode ?? false) - ); -} - -function getPresetSaveState({ - rawName, - activePreset, - presets, - activePresetDirty, -}: { - rawName: string; - activePreset: string; - presets: Preset[]; - activePresetDirty: boolean; -}): PresetSaveState { - const trimmedName = rawName.trim(); - if (!trimmedName) { - return { - mode: "disabled", - canSubmit: false, - isSaveReady: false, - buttonLabel: "Save", - title: "Enter a preset name", - }; - } - - const matchingPreset = presets.find((preset) => preset.name === trimmedName); - if (matchingPreset) { - const isActiveMatch = matchingPreset.name === activePreset; - return { - mode: isActiveMatch ? "overwrite-active" : "overwrite-other", - canSubmit: !isActiveMatch || activePresetDirty, - isSaveReady: !isActiveMatch || activePresetDirty, - buttonLabel: isActiveMatch && !activePresetDirty ? "Saved" : "Overwrite", - title: isActiveMatch - ? activePresetDirty - ? "Save current settings to this preset" - : "No unsaved changes" - : `Overwrite preset "${trimmedName}"`, - }; - } - - return { - mode: "create", - canSubmit: true, - isSaveReady: true, - buttonLabel: "Save as New", - title: `Save current settings as "${trimmedName}"`, - }; -} - function ParamSlider({ label, value, @@ -518,6 +404,10 @@ export function ChatSettingsPanel({ const setCustomContextLength = useChatRuntimeStore( (s) => s.setCustomContextLength, ); + const setActivePresetSource = useChatRuntimeStore( + (s) => s.setActivePresetSource, + ); + const activePresetSource = useChatRuntimeStore((s) => s.activePresetSource); const ctxDisplayValue = customContextLength ?? ggufContextLength ?? ""; const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null; @@ -540,12 +430,9 @@ export function ChatSettingsPanel({ >(undefined); const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false); const [systemPromptDraft, setSystemPromptDraft] = useState(""); + const [activePresetBaseline, setActivePresetBaseline] = useState(params); const presets = useMemo(() => { - const overrides = new Set(customPresets.map((preset) => preset.name)); - return [ - ...BUILTIN_PRESETS.filter((preset) => !overrides.has(preset.name)), - ...customPresets, - ]; + return getOrderedPresets(customPresets); }, [customPresets]); const activePresetDefinition = useMemo( () => presets.find((preset) => preset.name === activePreset) ?? null, @@ -555,17 +442,23 @@ export function ChatSettingsPanel({ () => customPresets.find((preset) => preset.name === activePreset) ?? null, [activePreset, customPresets], ); - const activeBuiltinPreset = useMemo( - () => - BUILTIN_PRESETS.find((preset) => preset.name === activePreset) ?? null, - [activePreset], - ); - const activePresetDirty = useMemo( - () => - activePresetDefinition == null - ? false - : !isSamePresetConfig(activePresetDefinition.params, params), - [activePresetDefinition, params], + const hasUnsavedPresetChanges = useMemo( + () => { + if (activePresetDefinition == null) { + return false; + } + if (BUILTIN_PRESET_NAMES.has(activePresetDefinition.name)) { + if (activePresetDefinition.name === "Default") { + return activePresetSource === "modified"; + } + return ( + activePresetSource === "modified" || + !isSamePresetConfig(activePresetDefinition.params, params) + ); + } + return !isSamePresetConfig(activePresetDefinition.params, params); + }, + [activePresetDefinition, activePresetSource, params], ); const presetSaveState = useMemo( () => @@ -573,9 +466,9 @@ export function ChatSettingsPanel({ rawName: presetNameInput, activePreset, presets, - activePresetDirty, + hasUnsavedPresetChanges, }), - [activePreset, activePresetDirty, presetNameInput, presets], + [activePreset, hasUnsavedPresetChanges, presetNameInput, presets], ); const systemPromptEditorDirty = systemPromptDraft !== params.systemPrompt; const trustRemoteCodeMissing = @@ -584,27 +477,24 @@ export function ChatSettingsPanel({ !(params.trustRemoteCode ?? false); function set(key: K) { - return (v: InferenceParams[K]) => onParamsChange({ ...params, [key]: v }); + return (v: InferenceParams[K]) => { + const nextParams = { ...params, [key]: v }; + const nextSource = isSamePresetConfig(activePresetBaseline, nextParams) + ? getPresetSource(activePreset) + : "modified"; + setActivePresetSource(nextSource); + onParamsChange(nextParams); + }; } function applyPreset(name: string) { const p = presets.find((pr) => pr.name === name); if (p) { - if ( - modelRequiresTrustRemoteCode && - !(p.params.trustRemoteCode ?? false) - ) { - toast.warning("This configuration turns custom code off", { - description: - "The current model needs custom code enabled to load. Keep it on for this model.", - }); - return; - } onParamsChange({ - ...p.params, - checkpoint: params.checkpoint, + ...applyPresetParams(params, p.params), }); setActivePreset(name); + setActivePresetSource(getPresetSource(name)); if (canUseStorage()) { try { localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, name); @@ -621,27 +511,29 @@ export function ChatSettingsPanel({ toast.error("Enter a preset name"); return; } + const usedNames = new Set([ + ...BUILTIN_PRESET_NAMES, + ...customPresets.map((preset) => preset.name), + ]); + const saveName = BUILTIN_PRESET_NAMES.has(trimmed) + ? getBuiltinVariantName(trimmed, usedNames) + : trimmed; setCustomPresets((prev) => { - const next = prev.filter((p) => p.name !== trimmed); - const merged = [...next, { name: trimmed, params: { ...params } }]; - if (canUseStorage()) { - try { - localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(merged)); - } catch { - // ignore - } - } + const next = prev.filter((p) => p.name !== saveName); + const merged = [...next, { name: saveName, params: toPresetParams(params) }]; + saveCustomPresets(merged); return merged; }); if (canUseStorage()) { try { - localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, trimmed); + localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, saveName); } catch { // ignore } } - setActivePreset(trimmed); - setPresetNameInput(trimmed); + setActivePreset(saveName); + setActivePresetSource("custom"); + setPresetNameInput(saveName); } function deletePreset(name: string) { @@ -651,41 +543,20 @@ export function ChatSettingsPanel({ if (!hasCustomPreset) { return; } - const builtinPreset = BUILTIN_PRESETS.find((preset) => preset.name === name); const fallbackPreset = - builtinPreset ?? - BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? - null; - if ( - activePreset === name && - fallbackPreset && - modelRequiresTrustRemoteCode && - !(fallbackPreset.params.trustRemoteCode ?? false) - ) { - toast.warning("Reset would turn custom code off", { - description: - "The current model needs custom code enabled to load. Keep it on for this model.", - }); - return; - } + BUILTIN_PRESETS.find((preset) => preset.name === "Default") ?? null; setCustomPresets((prev) => { const next = prev.filter((preset) => preset.name !== name); - if (canUseStorage()) { - try { - localStorage.setItem(CHAT_PRESETS_KEY, JSON.stringify(next)); - } catch { - // ignore - } - } + saveCustomPresets(next); return next; }); if (activePreset === name) { if (fallbackPreset) { onParamsChange({ - ...fallbackPreset.params, - checkpoint: params.checkpoint, + ...applyPresetParams(params, fallbackPreset.params), }); setActivePreset(fallbackPreset.name); + setActivePresetSource("builtin-default"); if (canUseStorage()) { try { localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, fallbackPreset.name); @@ -708,8 +579,46 @@ export function ChatSettingsPanel({ } useEffect(() => { - if (presets.some((preset) => preset.name === activePreset)) return; + if (activePresetSource !== "modified") { + setActivePresetBaseline(params); + } + }, [activePresetSource, params]); + + useEffect(() => { + if (presets.some((preset) => preset.name === activePreset)) { + const expectedSource = getPresetSource(activePreset); + if (activePresetDefinition != null) { + if (BUILTIN_PRESET_NAMES.has(activePresetDefinition.name)) { + if (activePresetDefinition.name === "Default") { + if ( + activePresetSource !== "modified" && + activePresetSource !== expectedSource + ) { + setActivePresetSource(expectedSource); + } + return; + } + const matchesActivePreset = isSamePresetConfig( + activePresetDefinition.params, + params, + ); + const nextSource = matchesActivePreset ? expectedSource : "modified"; + if (activePresetSource !== nextSource) { + setActivePresetSource(nextSource); + } + return; + } + } + if ( + activePresetSource !== "modified" && + activePresetSource !== expectedSource + ) { + setActivePresetSource(expectedSource); + } + return; + } setActivePreset("Default"); + setActivePresetSource("builtin-default"); if (canUseStorage()) { try { localStorage.setItem(CHAT_ACTIVE_PRESET_KEY, "Default"); @@ -717,7 +626,14 @@ export function ChatSettingsPanel({ // ignore } } - }, [activePreset, presets]); + }, [ + activePreset, + activePresetDefinition, + activePresetSource, + params, + presets, + setActivePresetSource, + ]); useEffect(() => { setPresetNameInput(activePreset); @@ -741,6 +657,187 @@ export function ChatSettingsPanel({ return () => ro.disconnect(); }, [open]); + const modelSection = ( + +
+ {isGguf && ( + <> +
+
+ Context Length + { + const raw = e.target.value; + if (raw === "") { + setCustomContextLength(null); + return; + } + const v = Number.parseInt(raw, 10); + if (!Number.isNaN(v) && v >= 0) { + const maxCtx = ctxMaxValue ?? Number.POSITIVE_INFINITY; + const clamped = Math.min(v, maxCtx); + setCustomContextLength( + clamped === (ggufContextLength ?? 0) ? null : clamped, + ); + } + }} + /> +
+ { + setCustomContextLength( + v === (ggufContextLength ?? 0) ? null : v, + ); + }} + /> + {ggufMaxContextLength != null && + typeof ctxDisplayValue === "number" && + ctxDisplayValue > ggufMaxContextLength && ( +

+ Exceeds estimated VRAM capacity ( + {ggufMaxContextLength.toLocaleString()} tokens). The model + may use system RAM. +

+ )} +
+
+
+
KV Cache Dtype
+
+ Quantize KV cache to reduce VRAM. +
+
+
+ +
+
+ {!currentModelIsVision && ( +
+
+
+ Speculative Decoding +
+
+ Speed up generation with no VRAM cost. +
+
+
+ +
+
+ )} + {modelSettingsDirty && ( +
+ + +
+ )} + + )} + {!isGguf && params.checkpoint && ( + <> +
+
+
Enable custom code
+
+ Allow models with custom code (e.g. Nemotron). Only enable if + sure. +
+
+ +
+ {trustRemoteCodeMissing && ( + + + Keep custom code enabled for this model + + + This model requires custom code to load. You can edit the + toggle, but loading will stay blocked until it is turned back + on. + + + )} + + )} +
+ + ); + const settingsContent = ( <>
@@ -834,13 +931,16 @@ export function ChatSettingsPanel({ : undefined } > - {presets.map((p) => ( - applyPreset(p.name)} - > - {p.name} - + {presets.map((p, index) => ( + + applyPreset(p.name)}> + {p.name} + + {index === BUILTIN_PRESETS.length - 1 && + presets.length > BUILTIN_PRESETS.length && ( + + )} + ))} @@ -874,9 +974,7 @@ export function ChatSettingsPanel({ className="h-8 w-full text-xs text-muted-foreground" title={ activeCustomPreset - ? activeBuiltinPreset - ? "Reset selected preset to built-in defaults" - : "Delete selected preset" + ? "Delete selected preset" : "No saved override to delete" } > @@ -918,188 +1016,6 @@ export function ChatSettingsPanel({ />
- -
- {isGguf && ( - <> -
-
- Context Length - { - const raw = e.target.value; - if (raw === "") { - setCustomContextLength(null); - return; - } - const v = Number.parseInt(raw, 10); - if (!Number.isNaN(v) && v >= 0) { - const maxCtx = - ctxMaxValue ?? Number.POSITIVE_INFINITY; - const clamped = Math.min(v, maxCtx); - setCustomContextLength( - clamped === (ggufContextLength ?? 0) - ? null - : clamped, - ); - } - }} - /> -
- { - setCustomContextLength( - v === (ggufContextLength ?? 0) ? null : v, - ); - }} - /> - {ggufMaxContextLength != null && - typeof ctxDisplayValue === "number" && - ctxDisplayValue > ggufMaxContextLength && ( -

- Exceeds estimated VRAM capacity ( - {ggufMaxContextLength.toLocaleString()} tokens). The - model may use system RAM. -

- )} -
-
-
-
KV Cache Dtype
-
- Quantize KV cache to reduce VRAM. -
-
-
- -
-
- {!currentModelIsVision && ( -
-
-
- Speculative Decoding -
-
- Speed up generation with no VRAM cost. -
-
-
- -
-
- )} - {modelSettingsDirty && ( -
- - -
- )} - - )} - {!isGguf && params.checkpoint && ( - <> -
-
-
Enable custom code
-
- Allow models with custom code (e.g. Nemotron). Only - enable if sure. -
-
- -
- {trustRemoteCodeMissing && ( - - - Keep custom code enabled for this model - - - This model requires custom code to load. You can edit the - toggle, but loading will stay blocked until it is turned - back on. - - - )} - - )} -
- - + {modelSection} +
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e90cab73fa..cfac6dc8cf 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -17,6 +17,10 @@ import { } from "../api/chat-api"; import { formatEta, formatRate } from "../utils/format-transfer"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import { + mergeBackendRecommendedInference, + resolveLoadMaxSeqLength, +} from "../presets/preset-policy"; import type { InferenceStatusResponse, LoadModelResponse } from "../types/api"; import type { ChatLoraSummary, @@ -132,46 +136,10 @@ function toLoraSummary(lora: { }; } -function toFiniteNumber(value: unknown): number | undefined { - if (typeof value !== "number" || !Number.isFinite(value)) { - return undefined; - } - return value; -} - function getTrustRemoteCodeRequiredMessage(modelName: string): string { return `${modelName} needs custom code enabled to load. Turn on "Enable custom code" in Chat Settings, then try again.`; } -function mergeRecommendedInference( - current: InferenceParams, - response: LoadModelResponse | InferenceStatusResponse, - modelId: string, -): InferenceParams { - const inference = response.inference; - // GGUF: use actual context length from GGUF metadata, fallback to 131072 - // Non-GGUF: 4096 - const defaultMaxTokens = response.is_gguf - ? (response.context_length ?? 131072) - : 4096; - return { - ...current, - checkpoint: modelId, - maxTokens: defaultMaxTokens, - temperature: - toFiniteNumber(inference?.temperature) ?? current.temperature, - topP: toFiniteNumber(inference?.top_p) ?? current.topP, - topK: toFiniteNumber(inference?.top_k) ?? current.topK, - minP: toFiniteNumber(inference?.min_p) ?? current.minP, - presencePenalty: - toFiniteNumber(inference?.presence_penalty) ?? current.presencePenalty, - trustRemoteCode: - typeof inference?.trust_remote_code === "boolean" - ? inference.trust_remote_code - : current.trustRemoteCode, - }; -} - export function useChatModelRuntime() { const params = useChatRuntimeStore((state) => state.params); const models = useChatRuntimeStore((state) => state.models); @@ -273,7 +241,12 @@ export function useChatModelRuntime() { speculative_type: statusRes.speculative_type, }; setParams( - mergeRecommendedInference(currentParams, statusRes, statusRes.active_model), + mergeBackendRecommendedInference({ + current: currentParams, + response: statusRes, + modelId: statusRes.active_model, + presetSource: useChatRuntimeStore.getState().activePresetSource, + }), ); } @@ -465,13 +438,25 @@ export function useChatModelRuntime() { previousWasUnloaded = true; } - const { chatTemplateOverride, kvCacheDtype, customContextLength, ggufContextLength, speculativeType } = useChatRuntimeStore.getState(); - // GGUF: use custom context length, or 0 = model's native context - // Non-GGUF: use the Max Seq Length slider value - const isDirectGgufFile = modelId.toLowerCase().endsWith(".gguf"); - const effectiveMaxSeqLength = customContextLength != null - ? customContextLength - : (ggufVariant != null || isDirectGgufFile) ? (ggufContextLength ?? 0) : maxSeqLength; + const { + chatTemplateOverride, + kvCacheDtype, + customContextLength, + ggufContextLength, + speculativeType, + activePresetSource, + activeGgufVariant, + } = useChatRuntimeStore.getState(); + const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ + modelId, + ggufVariant, + customContextLength, + ggufContextLength, + currentCheckpoint, + activeGgufVariant, + maxSeqLength, + presetSource: activePresetSource, + }); const loadResponse = await loadModel({ model_path: modelId, hf_token: hfToken, @@ -491,7 +476,12 @@ export function useChatModelRuntime() { const currentParams = useChatRuntimeStore.getState().params; setParams( - mergeRecommendedInference(currentParams, loadResponse, modelId), + mergeBackendRecommendedInference({ + current: currentParams, + response: loadResponse, + modelId, + presetSource: useChatRuntimeStore.getState().activePresetSource, + }), ); // Qwen3.5/3.6 small models (0.8B, 2B, 4B, 9B) disable thinking by default let reasoningDefault = loadResponse.supports_reasoning ?? false; @@ -545,12 +535,31 @@ export function useChatModelRuntime() { // Qwen3/3.5/3.6: apply thinking-mode-specific params after load if (modelId.toLowerCase().includes("qwen3") && (loadResponse.supports_reasoning ?? false)) { const store = useChatRuntimeStore.getState(); - const mid = modelId.toLowerCase(); - const needsPresencePenalty = mid.includes("qwen3.5") || mid.includes("qwen3.6"); - const p = reasoningDefault - ? { temperature: 0.6, topP: 0.95, topK: 20, minP: 0.0, ...(needsPresencePenalty ? { presencePenalty: 1.5 } : {}) } - : { temperature: 0.7, topP: 0.8, topK: 20, minP: 0.0, ...(needsPresencePenalty ? { presencePenalty: 1.5 } : {}) }; - store.setParams({ ...store.params, ...p }); + if (store.activePresetSource === "builtin-default") { + const mid = modelId.toLowerCase(); + const needsPresencePenalty = + mid.includes("qwen3.5") || mid.includes("qwen3.6"); + const p = reasoningDefault + ? { + temperature: 0.6, + topP: 0.95, + topK: 20, + minP: 0.0, + ...(needsPresencePenalty + ? { presencePenalty: 1.5 } + : {}), + } + : { + temperature: 0.7, + topP: 0.8, + topK: 20, + minP: 0.0, + ...(needsPresencePenalty + ? { presencePenalty: 1.5 } + : {}), + }; + store.setParams({ ...store.params, ...p }); + } } await refresh(); } catch (error) { diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts new file mode 100644 index 0000000000..8c5d2573c9 --- /dev/null +++ b/studio/frontend/src/features/chat/presets/preset-policy.ts @@ -0,0 +1,351 @@ +// 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 { + DEFAULT_INFERENCE_PARAMS, + type InferenceParams, +} from "../types/runtime"; + +export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS; + +export interface Preset { + name: string; + params: InferenceParams; +} + +export type PresetOwnedParams = Pick< + InferenceParams, + | "temperature" + | "topP" + | "topK" + | "minP" + | "repetitionPenalty" + | "presencePenalty" + | "maxTokens" + | "systemPrompt" +>; + +export const BUILTIN_PRESETS: Preset[] = [ + { name: "Default", params: { ...defaultInferenceParams } }, + { + name: "Creative", + params: { + ...defaultInferenceParams, + temperature: 1.5, + topP: 1.0, + topK: 0, + minP: 0.1, + repetitionPenalty: 1.0, + }, + }, + { + name: "Precise", + params: { + ...defaultInferenceParams, + temperature: 0.1, + topP: 0.95, + topK: 80, + minP: 0.01, + repetitionPenalty: 1.0, + }, + }, +]; + +export const BUILTIN_PRESET_NAMES = new Set( + BUILTIN_PRESETS.map((preset) => preset.name), +); + +export type ChatPresetSource = + | "builtin-default" + | "builtin-fixed" + | "custom" + | "modified"; + +export function getPresetSource(name: string): ChatPresetSource { + if (name === "Default") return "builtin-default"; + if (BUILTIN_PRESET_NAMES.has(name)) return "builtin-fixed"; + return "custom"; +} + +export function getUniquePresetName( + baseName: string, + usedNames: Set, +): string { + const normalizedBase = baseName.trim() || "Imported Prompt"; + let nextName = normalizedBase; + let suffix = 2; + while (usedNames.has(nextName)) { + nextName = `${normalizedBase} ${suffix}`; + suffix += 1; + } + usedNames.add(nextName); + return nextName; +} + +export function getBuiltinVariantName( + baseName: string, + usedNames: Set, +): string { + const normalizedBase = baseName.trim() || "Imported Prompt"; + let suffix = 1; + let nextName = `${normalizedBase} ${suffix}`; + while (usedNames.has(nextName)) { + suffix += 1; + nextName = `${normalizedBase} ${suffix}`; + } + usedNames.add(nextName); + return nextName; +} + +export function normalizeCustomPresets(presets: Preset[]): Preset[] { + const usedNames = new Set(BUILTIN_PRESET_NAMES); + return presets + .map((preset): Preset | null => { + const trimmedName = preset.name.trim(); + if (!trimmedName) return null; + const name = usedNames.has(trimmedName) + ? getBuiltinVariantName(trimmedName, usedNames) + : trimmedName; + usedNames.add(name); + return { + name, + params: preset.params, + }; + }) + .filter((preset): preset is Preset => preset !== null); +} + +export function getOrderedPresets(customPresets: Preset[]): Preset[] { + return [...BUILTIN_PRESETS, ...normalizeCustomPresets(customPresets)]; +} + +export function isSamePresetConfig( + a: InferenceParams, + b: InferenceParams, +): boolean { + const left = getPresetOwnedParams(a); + const right = getPresetOwnedParams(b); + return ( + left.temperature === right.temperature && + left.topP === right.topP && + left.topK === right.topK && + left.minP === right.minP && + left.repetitionPenalty === right.repetitionPenalty && + left.presencePenalty === right.presencePenalty && + left.maxTokens === right.maxTokens && + left.systemPrompt === right.systemPrompt + ); +} + +export function getPresetOwnedParams( + params: InferenceParams, +): PresetOwnedParams { + return { + temperature: params.temperature, + topP: params.topP, + topK: params.topK, + minP: params.minP, + repetitionPenalty: params.repetitionPenalty, + presencePenalty: params.presencePenalty, + maxTokens: params.maxTokens, + systemPrompt: params.systemPrompt, + }; +} + +export function getPresetOwnedConfigKey(params: InferenceParams): string { + return JSON.stringify(getPresetOwnedParams(params)); +} + +export function toPresetParams(params: InferenceParams): InferenceParams { + return { + ...defaultInferenceParams, + ...getPresetOwnedParams(params), + }; +} + +export function applyPresetParams( + current: InferenceParams, + preset: InferenceParams, +): InferenceParams { + return { + ...current, + ...getPresetOwnedParams(preset), + }; +} + +export type PresetSaveMode = + | "disabled" + | "overwrite-active" + | "overwrite-other" + | "copy-builtin" + | "create"; + +export interface PresetSaveState { + mode: PresetSaveMode; + canSubmit: boolean; + isSaveReady: boolean; + buttonLabel: string; + title: string; +} + +export function getPresetSaveState({ + rawName, + activePreset, + presets, + hasUnsavedPresetChanges, +}: { + rawName: string; + activePreset: string; + presets: Preset[]; + hasUnsavedPresetChanges: boolean; +}): PresetSaveState { + const trimmedName = rawName.trim(); + if (!trimmedName) { + return { + mode: "disabled", + canSubmit: false, + isSaveReady: false, + buttonLabel: "Save", + title: "Enter a preset name", + }; + } + + if (BUILTIN_PRESET_NAMES.has(trimmedName)) { + const variantName = getBuiltinVariantName(trimmedName, new Set(presets.map((preset) => preset.name))); + return { + mode: "copy-builtin", + canSubmit: activePreset !== trimmedName || hasUnsavedPresetChanges, + isSaveReady: activePreset !== trimmedName || hasUnsavedPresetChanges, + buttonLabel: + activePreset === trimmedName && !hasUnsavedPresetChanges + ? "Saved" + : "Save", + title: + activePreset === trimmedName && !hasUnsavedPresetChanges + ? "No unsaved changes" + : `Save current settings as "${variantName}"`, + }; + } + + const matchingPreset = presets.find((preset) => preset.name === trimmedName); + if (matchingPreset) { + const isActiveMatch = matchingPreset.name === activePreset; + return { + mode: isActiveMatch ? "overwrite-active" : "overwrite-other", + canSubmit: !isActiveMatch || hasUnsavedPresetChanges, + isSaveReady: !isActiveMatch || hasUnsavedPresetChanges, + buttonLabel: + isActiveMatch && !hasUnsavedPresetChanges ? "Saved" : "Save", + title: isActiveMatch + ? hasUnsavedPresetChanges + ? "Save current settings to this preset" + : "No unsaved changes" + : `Overwrite preset "${trimmedName}"`, + }; + } + + return { + mode: "create", + canSubmit: true, + isSaveReady: true, + buttonLabel: "Save", + title: `Save current settings as "${trimmedName}"`, + }; +} + +function toFiniteNumber(value: unknown): number | undefined { + if (typeof value !== "number" || !Number.isFinite(value)) { + return undefined; + } + return value; +} + +interface BackendInferenceDefaults { + temperature?: number; + top_p?: number; + top_k?: number; + min_p?: number; + presence_penalty?: number; + trust_remote_code?: boolean; +} + +export interface BackendInferenceEnvelope { + is_gguf?: boolean; + context_length?: number | null; + inference?: BackendInferenceDefaults; +} + +export function mergeBackendRecommendedInference({ + current, + response, + modelId, + presetSource, +}: { + current: InferenceParams; + response: BackendInferenceEnvelope; + modelId: string; + presetSource: ChatPresetSource; +}): InferenceParams { + const inference = response.inference; + const next: InferenceParams = { + ...current, + checkpoint: modelId, + trustRemoteCode: + typeof inference?.trust_remote_code === "boolean" + ? inference.trust_remote_code + : current.trustRemoteCode, + }; + + if (presetSource !== "builtin-default") { + return next; + } + + const defaultMaxTokens = response.is_gguf + ? (response.context_length ?? current.maxTokens) + : 4096; + return { + ...next, + maxTokens: defaultMaxTokens, + temperature: + toFiniteNumber(inference?.temperature) ?? defaultInferenceParams.temperature, + topP: toFiniteNumber(inference?.top_p) ?? defaultInferenceParams.topP, + topK: toFiniteNumber(inference?.top_k) ?? defaultInferenceParams.topK, + minP: toFiniteNumber(inference?.min_p) ?? defaultInferenceParams.minP, + presencePenalty: + toFiniteNumber(inference?.presence_penalty) ?? + defaultInferenceParams.presencePenalty, + }; +} + +export function resolveLoadMaxSeqLength({ + modelId, + ggufVariant, + customContextLength, + ggufContextLength, + currentCheckpoint, + activeGgufVariant, + maxSeqLength, + presetSource, +}: { + modelId: string; + ggufVariant?: string | null; + customContextLength: number | null; + ggufContextLength: number | null; + currentCheckpoint: string; + activeGgufVariant?: string | null; + maxSeqLength: number; + presetSource: ChatPresetSource; +}): number { + const isDirectGgufFile = modelId.toLowerCase().endsWith(".gguf"); + const isGgufLoad = ggufVariant != null || isDirectGgufFile; + const isReloadingCurrentGguf = + isGgufLoad && + currentCheckpoint === modelId && + (ggufVariant ?? null) === (activeGgufVariant ?? null); + + if (customContextLength != null) return customContextLength; + if (isGgufLoad && presetSource === "builtin-default") return 0; + if (isReloadingCurrentGguf) return ggufContextLength ?? 0; + if (isGgufLoad) return 0; + return maxSeqLength; +} diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index ce69d8f6dc..13f2a23c36 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -9,6 +9,10 @@ import { type ChatModelSummary, type InferenceParams, } from "../types/runtime"; +import { + getPresetSource, + type ChatPresetSource, +} from "../presets/preset-policy"; const AUTO_TITLE_KEY = "unsloth_chat_auto_title"; const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls"; @@ -16,6 +20,8 @@ const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message"; const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout"; const HF_TOKEN_KEY = "unsloth_hf_token"; const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params"; +const CHAT_ACTIVE_PRESET_KEY = "unsloth_chat_active_preset"; +const CHAT_ACTIVE_PRESET_SOURCE_KEY = "unsloth_chat_active_preset_source"; const REASONING_EFFORT_KEY = "unsloth_reasoning_effort"; const PRESERVE_THINKING_KEY = "unsloth_preserve_thinking"; @@ -157,8 +163,24 @@ function saveInferenceParams(params: InferenceParams): boolean { } } +function loadPresetSource(): ChatPresetSource { + const activePreset = loadString(CHAT_ACTIVE_PRESET_KEY, "Default"); + if (canUseStorage()) { + try { + const raw = localStorage.getItem(CHAT_ACTIVE_PRESET_SOURCE_KEY); + if (raw === "modified") { + return "modified"; + } + } catch { + // ignore + } + } + return getPresetSource(activePreset); +} + type ChatRuntimeStore = { params: InferenceParams; + activePresetSource: ChatPresetSource; models: ChatModelSummary[]; loras: ChatLoraSummary[]; runningByThreadId: Record; @@ -207,6 +229,7 @@ type ChatRuntimeStore = { setModelLoading: (loading: boolean) => void; setModelRequiresTrustRemoteCode: (required: boolean) => void; setParams: (params: InferenceParams) => void; + setActivePresetSource: (source: ChatPresetSource) => void; setModels: (models: ChatModelSummary[]) => void; setLoras: (loras: ChatLoraSummary[]) => void; setThreadRunning: (threadId: string, running: boolean) => void; @@ -241,6 +264,7 @@ type ChatRuntimeStore = { export const useChatRuntimeStore = create((set) => ({ params: loadInferenceParams(), + activePresetSource: loadPresetSource(), models: [], loras: [], runningByThreadId: {}, @@ -296,6 +320,11 @@ export const useChatRuntimeStore = create((set) => ({ } return { params }; }), + setActivePresetSource: (activePresetSource) => + set(() => { + saveString(CHAT_ACTIVE_PRESET_SOURCE_KEY, activePresetSource); + return { activePresetSource }; + }), setModels: (models) => set({ models }), setLoras: (loras) => set({ loras }), setThreadRunning: (threadId, running) => diff --git a/tests/studio/test_chat_preset_builtin_invariants.py b/tests/studio/test_chat_preset_builtin_invariants.py new file mode 100644 index 0000000000..da5f21099e --- /dev/null +++ b/tests/studio/test_chat_preset_builtin_invariants.py @@ -0,0 +1,272 @@ +import json +import os +import shutil +import subprocess +import textwrap +from pathlib import Path + +import pytest + +WORKDIR = Path(__file__).resolve().parents[2] +PRESET_POLICY = ( + WORKDIR / "unsloth_repo/studio/frontend/src/features/chat/presets/preset-policy.ts" +) +RUNTIME_TYPES = ( + WORKDIR / "unsloth_repo/studio/frontend/src/features/chat/types/runtime.ts" +) +TEMP = WORKDIR / "temp" / "chat_preset_builtin_invariants" + + +def _require_node(): + if shutil.which("node") is None: + pytest.skip("node not available") + if not PRESET_POLICY.exists() or not RUNTIME_TYPES.exists(): + pytest.skip("studio chat sources not present") + + +def _ensure_harness(): + TEMP.mkdir(parents = True, exist_ok = True) + (TEMP / "register.mjs").write_text( + "import { register } from 'node:module';\n" + "register('./loader.mjs', import.meta.url);\n" + ) + (TEMP / "loader.mjs").write_text( + "export function resolve(specifier, context, next) {\n" + " if (specifier.endsWith('/types/runtime')) return next(specifier + '.ts', context);\n" + " return next(specifier, context);\n" + "}\n" + ) + + +def _run(script: str): + _require_node() + _ensure_harness() + script_path = TEMP / "run.mts" + script_path.write_text(script) + env = dict(os.environ, NODE_NO_WARNINGS = "1") + result = subprocess.run( + [ + "node", + "--experimental-strip-types", + "--import=./register.mjs", + "--no-warnings", + "run.mts", + ], + cwd = str(TEMP), + capture_output = True, + text = True, + timeout = 30, + env = env, + ) + assert result.returncode == 0, f"stderr: {result.stderr}\nstdout: {result.stdout}" + last = [line for line in result.stdout.strip().splitlines() if line.strip()][-1] + return json.loads(last) + + +def _policy_path(): + return os.path.relpath(PRESET_POLICY, TEMP).replace("\\", "/") + + +def _runtime_path(): + return os.path.relpath(RUNTIME_TYPES, TEMP).replace("\\", "/") + + +def test_default_builtin_matches_default_inference_params(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, isSamePresetConfig }} from "{_policy_path()}"; + import {{ DEFAULT_INFERENCE_PARAMS }} from "{_runtime_path()}"; + const def = BUILTIN_PRESETS.find((p) => p.name === "Default"); + console.log(JSON.stringify({{ + found: !!def, + matches: def ? isSamePresetConfig(def.params, DEFAULT_INFERENCE_PARAMS) : null, + }})); + """ + ) + ) + assert out["found"] is True + assert out["matches"] is True + + +def test_is_same_preset_config_detects_temperature_edit(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, isSamePresetConfig }} from "{_policy_path()}"; + const def = BUILTIN_PRESETS.find((p) => p.name === "Default"); + const edited = {{ ...def.params, temperature: def.params.temperature + 0.1 }}; + console.log(JSON.stringify({{ same: isSamePresetConfig(def.params, edited) }})); + """ + ) + ) + assert out["same"] is False + + +def test_is_same_preset_config_detects_system_prompt_edit(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, isSamePresetConfig }} from "{_policy_path()}"; + const def = BUILTIN_PRESETS.find((p) => p.name === "Default"); + const edited = {{ ...def.params, systemPrompt: "you are a pirate" }}; + console.log(JSON.stringify({{ same: isSamePresetConfig(def.params, edited) }})); + """ + ) + ) + assert out["same"] is False + + +def test_is_same_preset_config_ignores_checkpoint_difference(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, isSamePresetConfig }} from "{_policy_path()}"; + const def = BUILTIN_PRESETS.find((p) => p.name === "Default"); + const withCheckpoint = {{ ...def.params, checkpoint: "meta-llama/Llama-3-8B" }}; + console.log(JSON.stringify({{ same: isSamePresetConfig(def.params, withCheckpoint) }})); + """ + ) + ) + assert out["same"] is True + + +def test_is_same_preset_config_ignores_model_owned_fields(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, isSamePresetConfig }} from "{_policy_path()}"; + const def = BUILTIN_PRESETS.find((p) => p.name === "Default"); + const edited = {{ + ...def.params, + maxSeqLength: def.params.maxSeqLength + 1024, + trustRemoteCode: !def.params.trustRemoteCode, + }}; + console.log(JSON.stringify({{ same: isSamePresetConfig(def.params, edited) }})); + """ + ) + ) + assert out["same"] is True + + +def test_preset_owned_config_key_ignores_model_owned_fields(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, getPresetOwnedConfigKey }} from "{_policy_path()}"; + const def = BUILTIN_PRESETS.find((p) => p.name === "Default"); + const edited = {{ + ...def.params, + checkpoint: "foo/bar", + maxSeqLength: def.params.maxSeqLength + 1024, + trustRemoteCode: !def.params.trustRemoteCode, + }}; + console.log(JSON.stringify({{ + same: getPresetOwnedConfigKey(def.params) === getPresetOwnedConfigKey(edited), + }})); + """ + ) + ) + assert out["same"] is True + + +def test_to_preset_params_strips_model_owned_fields(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ toPresetParams }} from "{_policy_path()}"; + const sanitized = toPresetParams({{ + temperature: 0.9, + topP: 0.8, + topK: 40, + minP: 0.05, + repetitionPenalty: 1.1, + presencePenalty: 0.4, + maxSeqLength: 16384, + maxTokens: 2048, + systemPrompt: "hello", + checkpoint: "foo/bar", + trustRemoteCode: true, + }}); + console.log(JSON.stringify({{ + checkpoint: sanitized.checkpoint, + trustRemoteCode: sanitized.trustRemoteCode, + maxSeqLength: sanitized.maxSeqLength, + maxTokens: sanitized.maxTokens, + systemPrompt: sanitized.systemPrompt, + }})); + """ + ) + ) + assert out["checkpoint"] == "" + assert out["trustRemoteCode"] is False + assert out["maxSeqLength"] == 4096 + assert out["maxTokens"] == 2048 + assert out["systemPrompt"] == "hello" + + +def test_apply_preset_params_preserves_model_owned_fields(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, applyPresetParams }} from "{_policy_path()}"; + const creative = BUILTIN_PRESETS.find((p) => p.name === "Creative"); + const applied = applyPresetParams( + {{ + temperature: 0.6, + topP: 0.95, + topK: 20, + minP: 0.01, + repetitionPenalty: 1.0, + presencePenalty: 0.0, + maxSeqLength: 16384, + maxTokens: 8192, + systemPrompt: "keep me?", + checkpoint: "foo/bar", + trustRemoteCode: true, + }}, + creative.params, + ); + console.log(JSON.stringify({{ + checkpoint: applied.checkpoint, + trustRemoteCode: applied.trustRemoteCode, + maxSeqLength: applied.maxSeqLength, + temperature: applied.temperature, + topK: applied.topK, + }})); + """ + ) + ) + assert out["checkpoint"] == "foo/bar" + assert out["trustRemoteCode"] is True + assert out["maxSeqLength"] == 16384 + assert out["temperature"] == 1.5 + assert out["topK"] == 0 + + +def test_creative_and_precise_builtins_differ_from_default(): + out = _run( + textwrap.dedent( + f""" + // @ts-nocheck + import {{ BUILTIN_PRESETS, isSamePresetConfig }} from "{_policy_path()}"; + const def = BUILTIN_PRESETS.find((p) => p.name === "Default"); + const creative = BUILTIN_PRESETS.find((p) => p.name === "Creative"); + const precise = BUILTIN_PRESETS.find((p) => p.name === "Precise"); + console.log(JSON.stringify({{ + creativeDiffers: !isSamePresetConfig(def.params, creative.params), + preciseDiffers: !isSamePresetConfig(def.params, precise.params), + }})); + """ + ) + ) + assert out["creativeDiffers"] is True + assert out["preciseDiffers"] is True From ff759ba7e4d852c288723dd96e0a1fb7e935db37 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Tue, 28 Apr 2026 22:49:13 +0100 Subject: [PATCH 16/54] Studio: Fix image-only chat requests failing validation (#5212) * fix: allow image-only chat messages * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: deduplicate empty content validation coverage --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/models/inference.py | 11 ++++++----- .../tests/test_openai_tool_passthrough.py | 19 ++++++++++++++----- 2 files changed, 20 insertions(+), 10 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index bf0177efbf..eb2bb9c5ce 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -396,7 +396,10 @@ class ChatMessage(BaseModel): if self.name is not None and self.role != "tool": raise ValueError('"name" is only valid on role="tool" messages.') - # Per-role content requirements. + # Per-role content requirements. OpenAI-compatible clients may send + # ``content=""`` for image-only turns when the image travels in a + # companion field such as Studio's ``image_base64`` extension, so treat + # empty strings as present content for user/system messages. if self.role == "tool": if not self.tool_call_id: raise ValueError( @@ -411,10 +414,8 @@ class ChatMessage(BaseModel): 'role="assistant" messages require either "content" or "tool_calls".' ) else: # "user" | "system" - if not self.content: - raise ValueError( - f'role="{self.role}" messages require non-empty "content".' - ) + if self.content is None or self.content == []: + raise ValueError(f'role="{self.role}" messages require "content".') return self diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index ccb0dba325..cdb7f5d270 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -144,13 +144,14 @@ class TestChatMessageToolRoles: # ── Role-aware content requirements ──────────────────────────── - def test_user_empty_content_rejected(self): - with pytest.raises(ValidationError): - ChatMessage(role = "user", content = "") + @pytest.mark.parametrize("role", ["user", "system"]) + def test_empty_string_content_allowed(self, role): + msg = ChatMessage(role = role, content = "") + assert msg.content == "" - def test_system_empty_content_rejected(self): + def test_user_missing_content_rejected(self): with pytest.raises(ValidationError): - ChatMessage(role = "system", content = "") + ChatMessage(role = "user") def test_user_empty_list_content_rejected(self): with pytest.raises(ValidationError): @@ -226,6 +227,14 @@ class TestChatCompletionRequestToolFields: assert len(req.tools) == 1 assert req.tools[0]["function"]["name"] == "get_weather" + def test_image_base64_allows_empty_user_text(self): + req = ChatCompletionRequest( + messages = [{"role": "user", "content": ""}], + image_base64 = "aW1hZ2U=", + ) + assert req.messages[0].content == "" + assert req.image_base64 == "aW1hZ2U=" + def test_tool_choice_string_auto(self): assert self._make(tool_choice = "auto").tool_choice == "auto" From a5615426a56857983909b05016b3d985710e0015 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 28 Apr 2026 22:43:44 -0700 Subject: [PATCH 17/54] Studio: fix 7 failing studio_unit_tests on main (#5216) * Studio: fix 4 failing studio_unit_tests on main Three of the failing tests had drifted from production: 1. test_health_response_reports_desktop_capability_fields stubbed `routes` with a SimpleNamespace that omitted `inference_studio_router`, so importing studio.backend.main raised ImportError. Add the missing router stub. 2. test_local_recipe_token_preserves_desktop_marker and test_local_recipe_token_keeps_web_marker_absent decoded the local provider's api_key as a JWT, but _inject_local_providers now mints a unified sk-unsloth-* internal API key (not a forwarded JWT), so jwt.decode raised "Not enough segments". Renamed and rewrote both tests to validate the API-key contract: starts with storage.API_KEY_PREFIX and authenticates via get_current_subject as the real admin user. The web vs desktop distinction is irrelevant at this layer because the unified API-key path does not carry session flags. The fourth failure was a real production bug: 3. test_github_validate_skips_live_access_with_honest_note expected github-seed validation to return valid=True per _GITHUB_VALIDATE_NOTE ("GitHub access and rate limits are checked when the run starts"). The validate route called build_config_builder which lazy-imports the optional data_designer module; when it is missing, the bare except blocked the recipe. Catch ImportError specifically and treat it as a deferred check, matching the documented intent. Verified all 4 tests pass and the rest of studio/backend/tests still pass (608 total, with the only remaining failures being environment specific: 4 GPU-aware tests on a no-GPU host and 1 Anthropic-API smoke test, both unrelated). * Studio: fix 3 test_gpu_selection route tests after load_model signature change `routes/inference.load_model` gained a `fastapi_request: Request` positional argument (used to read `app.state.llama_parallel_slots` inside the GGUF path), but the three TestRouteErrors cases that exercise the early validation path were not updated and failed with `TypeError: load_model() missing 1 required positional argument: 'fastapi_request'`. Pass a SimpleNamespace mock that satisfies the attribute path the production code reads. The validation under test fires before the mock is consumed, but supplying the realistic shape protects against regressions if the validation order changes. Affected tests: - test_inference_route_rejects_gpu_ids_for_gguf - test_inference_route_returns_400_for_invalid_gpu_ids - test_inference_route_returns_400_for_uuid_parent_visibility_gpu_ids * Studio: address review feedback on validate.py ImportError handling Two reviewers flagged the ImportError bypass added in b0d33cf: - chatgpt-codex-connector[bot]: catching bare ImportError marks recipes as valid even when build_config_builder fails for unrelated import problems (broken internal imports, missing transitive deps after a version bump), hiding real regressions until run start. - gemini-code-assist[bot]: silent pass discourages troubleshooting; the deferred-validation case should be logged at debug level. Tighten the bypass to ModuleNotFoundError where the missing module name starts with "data_designer". Other ImportErrors propagate to the outer handler and surface as validation failures, restoring the visibility the reviewers asked for. Add a debug-level log entry that names the missing module so operators can trace why validation deferred. --- studio/backend/routes/data_recipe/validate.py | 18 ++++++++++ studio/backend/tests/test_desktop_auth.py | 33 ++++++++++--------- studio/backend/tests/test_gpu_selection.py | 30 +++++++++++++++-- 3 files changed, 62 insertions(+), 19 deletions(-) diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py index 87eef939b4..e794d68e54 100644 --- a/studio/backend/routes/data_recipe/validate.py +++ b/studio/backend/routes/data_recipe/validate.py @@ -14,8 +14,10 @@ from core.data_recipe.service import ( create_data_designer, validate_recipe, ) +from loggers import get_logger from models.data_recipe import RecipePayload, ValidateError, ValidateResponse +logger = get_logger(__name__) router = APIRouter() _GITHUB_VALIDATE_NOTE = "Recipe shape is valid. GitHub access and rate limits are checked when the run starts." @@ -151,6 +153,22 @@ def validate(payload: RecipePayload) -> ValidateResponse: return ValidateResponse(valid = False, errors = static_errors) try: build_config_builder(recipe) + except ModuleNotFoundError as exc: + # data_designer is an optional runtime dep. Static validation + # already passed; live access + full config validation are + # deferred to run start (per _GITHUB_VALIDATE_NOTE), so a missing + # optional import at validate time should not block the recipe. + # Restrict the bypass to the data_designer module specifically so + # other ImportErrors (e.g. broken internal imports or missing + # transitive deps after a package upgrade) still surface as + # validation failures instead of being silently swallowed. + if not (exc.name or "").startswith("data_designer"): + raise + logger.debug( + "data_designer not installed; deferring full config " + "validation to run start", + missing_module = exc.name, + ) except Exception as exc: detail = str(exc).strip() or "Validation failed." return ValidateResponse( diff --git a/studio/backend/tests/test_desktop_auth.py b/studio/backend/tests/test_desktop_auth.py index c8cf1c7081..a5508c1c8b 100644 --- a/studio/backend/tests/test_desktop_auth.py +++ b/studio/backend/tests/test_desktop_auth.py @@ -246,7 +246,10 @@ def test_desktop_session_uses_real_admin_identity_for_api_keys(): assert [row["name"] for row in rows] == ["desktop"] -def test_local_recipe_token_preserves_desktop_marker(loaded_local_model): +def test_local_recipe_token_authenticates_as_admin_for_desktop_user(loaded_local_model): + # _inject_local_providers mints an internal sk-unsloth-* API key (not a + # forwarded JWT). The unified API-key path validates as the real admin + # user regardless of whether the incoming session was desktop or web. from auth.authentication import create_access_token, get_current_subject seed_user(must_change_password = True) @@ -260,13 +263,7 @@ def test_local_recipe_token_preserves_desktop_marker(loaded_local_model): jobs_route._inject_local_providers(recipe, local_recipe_request(incoming_token)) local_token = recipe["model_providers"][0]["api_key"] - payload = jwt.decode( - local_token, - storage.get_jwt_secret(storage.DEFAULT_ADMIN_USERNAME), - algorithms = ["HS256"], - ) - assert payload["sub"] == storage.DEFAULT_ADMIN_USERNAME - assert payload["desktop"] is True + assert local_token.startswith(storage.API_KEY_PREFIX) credentials = HTTPAuthorizationCredentials( scheme = "Bearer", credentials = local_token, @@ -276,8 +273,10 @@ def test_local_recipe_token_preserves_desktop_marker(loaded_local_model): ) -def test_local_recipe_token_keeps_web_marker_absent(loaded_local_model): - from auth.authentication import create_access_token +def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_model): + # Mirror of the desktop variant: API-key issuance is identical for web + # and desktop incoming tokens; auth via get_current_subject works the same. + from auth.authentication import create_access_token, get_current_subject seed_user(must_change_password = False) jobs_route = data_recipe_jobs_module() @@ -287,13 +286,14 @@ def test_local_recipe_token_keeps_web_marker_absent(loaded_local_model): jobs_route._inject_local_providers(recipe, local_recipe_request(incoming_token)) local_token = recipe["model_providers"][0]["api_key"] - payload = jwt.decode( - local_token, - storage.get_jwt_secret(storage.DEFAULT_ADMIN_USERNAME), - algorithms = ["HS256"], + assert local_token.startswith(storage.API_KEY_PREFIX) + credentials = HTTPAuthorizationCredentials( + scheme = "Bearer", + credentials = local_token, + ) + assert ( + asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME ) - assert payload["sub"] == storage.DEFAULT_ADMIN_USERNAME - assert "desktop" not in payload def test_desktop_login_rejects_invalid_secret(): @@ -381,6 +381,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch): datasets_router = APIRouter(), export_router = APIRouter(), inference_router = APIRouter(), + inference_studio_router = APIRouter(), models_router = APIRouter(), training_history_router = APIRouter(), training_router = APIRouter(), diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py index c6f26037af..fb4f38565b 100644 --- a/studio/backend/tests/test_gpu_selection.py +++ b/studio/backend/tests/test_gpu_selection.py @@ -746,7 +746,15 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model(request, current_subject = "test-user") + inference_route.load_model( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) ) self.assertEqual(exc_info.exception.status_code, 400) @@ -886,7 +894,15 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model(request, current_subject = "test-user") + inference_route.load_model( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) ) self.assertEqual(exc_info.exception.status_code, 400) @@ -942,7 +958,15 @@ class TestRouteErrors(unittest.TestCase): ): with self.assertRaises(HTTPException) as exc_info: asyncio.run( - inference_route.load_model(request, current_subject = "test-user") + inference_route.load_model( + request, + SimpleNamespace( + app = SimpleNamespace( + state = SimpleNamespace(llama_parallel_slots = 1), + ), + ), + current_subject = "test-user", + ) ) self.assertEqual(exc_info.exception.status_code, 400) From c4597298be34d2af88c05c1a6e2dae3070819326 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 29 Apr 2026 15:20:49 +0530 Subject: [PATCH 18/54] Patch checkpoint reload init functions to strip unsupported args (#5167) * Patch checkpoint reload init functions to strip unsupported args * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Try adding attrs back if possible * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * import_fixes: harden peft weight converter compatibility shim Three small fixes to patch_peft_weight_converter_compatibility: - Restore peft_config=None default and @functools.wraps on the build_peft_weight_mapping wrapper so the upstream coordinated signature (weight_conversions, adapter_name, peft_config=None) is preserved. Without the default, callers using the documented two-argument form raise TypeError after import unsloth. - Serialize the temporary class-init patch/restore behind a threading.RLock. The previous unsynchronized window let two concurrent build_peft_weight_mapping calls (e.g. dynamic LoRA serving) re-expose the original distributed_operation TypeError when one thread restored a class while another was still inside original_build. - Hand _patch_weight_converter_ctors a caller-owned accumulator list and append in place. If signature inspection ever raises mid-loop, the finally block now sees the partial list and restores already-patched classes instead of leaving them with the compat init permanently installed. * Add tests for peft weight converter compatibility shim * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han --- tests/test_peft_weight_converter_compat.py | 259 +++++++++++++++++++++ unsloth/__init__.py | 3 + unsloth/import_fixes.py | 83 +++++++ 3 files changed, 345 insertions(+) create mode 100644 tests/test_peft_weight_converter_compat.py diff --git a/tests/test_peft_weight_converter_compat.py b/tests/test_peft_weight_converter_compat.py new file mode 100644 index 0000000000..62f17f0f45 --- /dev/null +++ b/tests/test_peft_weight_converter_compat.py @@ -0,0 +1,259 @@ +import importlib.util +import inspect +import sys +import threading +import types +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +IMPORT_FIXES = REPO_ROOT / "unsloth" / "import_fixes.py" + + +def _load_patch_function(): + spec = importlib.util.spec_from_file_location( + "_unsloth_import_fixes_under_test", IMPORT_FIXES + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.patch_peft_weight_converter_compatibility + + +def _install_fake_peft(twc_namespace): + peft_pkg = types.ModuleType("peft") + peft_pkg.__path__ = [] + peft_utils = types.ModuleType("peft.utils") + peft_utils.__path__ = [] + twc = types.ModuleType("peft.utils.transformers_weight_conversion") + for k, v in twc_namespace.items(): + setattr(twc, k, v) + peft_utils.transformers_weight_conversion = twc + sys.modules["peft"] = peft_pkg + sys.modules["peft.utils"] = peft_utils + sys.modules["peft.utils.transformers_weight_conversion"] = twc + return twc + + +@pytest.fixture(autouse = True) +def _restore_peft_modules(): + saved = { + k: sys.modules.get(k) + for k in ( + "peft", + "peft.utils", + "peft.utils.transformers_weight_conversion", + ) + } + yield + for k, v in saved.items(): + if v is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = v + + +class _LegacyConverter: + def __init__(self, source_patterns, target_patterns, operations): + self.source_patterns = source_patterns + self.target_patterns = target_patterns + self.operations = operations + self.distributed_operation = None + self.quantization_operation = None + + +class _ModernConverter: + def __init__( + self, + source_patterns, + target_patterns, + operations, + distributed_operation = None, + quantization_operation = None, + ): + self.source_patterns = source_patterns + self.target_patterns = target_patterns + self.operations = operations + self.distributed_operation = distributed_operation + self.quantization_operation = quantization_operation + + +def _make_legacy_converter(): + return _LegacyConverter(["src.*"], ["tgt.*"], []) + + +def _make_modern_converter(): + return _ModernConverter(["src.*"], ["tgt.*"], []) + + +def _build_that_calls_init(weight_conversions, adapter_name, peft_config = None): + out = [] + for c in weight_conversions or []: + out.append( + c.__class__( + source_patterns = c.source_patterns, + target_patterns = c.target_patterns, + operations = c.operations, + distributed_operation = "dist-x", + quantization_operation = "quant-y", + ) + ) + return out + + +def test_two_arg_call_preserves_upstream_signature(): + twc = _install_fake_peft({"build_peft_weight_mapping": _build_that_calls_init}) + patch = _load_patch_function() + patch() + + sig = inspect.signature(twc.build_peft_weight_mapping) + assert "peft_config" in sig.parameters + assert sig.parameters["peft_config"].default is None + + out = twc.build_peft_weight_mapping([_make_legacy_converter()], "default") + assert len(out) == 1 + assert out[0].distributed_operation == "dist-x" + assert out[0].quantization_operation == "quant-y" + + +def test_legacy_init_succeeds_after_patch(): + twc = _install_fake_peft({"build_peft_weight_mapping": _build_that_calls_init}) + patch = _load_patch_function() + patch() + + out = twc.build_peft_weight_mapping([_make_legacy_converter()], "default", None) + assert len(out) == 1 + assert out[0].distributed_operation == "dist-x" + assert out[0].quantization_operation == "quant-y" + + +def test_modern_init_not_patched(): + twc = _install_fake_peft({"build_peft_weight_mapping": _build_that_calls_init}) + pre_init = _ModernConverter.__init__ + patch = _load_patch_function() + patch() + + twc.build_peft_weight_mapping([_make_modern_converter()], "default", None) + assert _ModernConverter.__init__ is pre_init + + +def test_class_init_restored_after_call(): + twc = _install_fake_peft({"build_peft_weight_mapping": _build_that_calls_init}) + pre_init = _LegacyConverter.__init__ + patch = _load_patch_function() + patch() + + twc.build_peft_weight_mapping([_make_legacy_converter()], "default", None) + assert _LegacyConverter.__init__ is pre_init + + +def test_class_init_restored_after_original_build_raises(): + def _raise(weight_conversions, adapter_name, peft_config = None): + raise RuntimeError("simulated PEFT failure") + + twc = _install_fake_peft({"build_peft_weight_mapping": _raise}) + pre_init = _LegacyConverter.__init__ + patch = _load_patch_function() + patch() + + with pytest.raises(RuntimeError): + twc.build_peft_weight_mapping([_make_legacy_converter()], "default", None) + assert _LegacyConverter.__init__ is pre_init + + +def test_partial_patch_restored_when_inspect_signature_raises_mid_loop(): + twc = _install_fake_peft({"build_peft_weight_mapping": _build_that_calls_init}) + pre_legacy = _LegacyConverter.__init__ + + class _BadInitConverter: + def __init__(self, source_patterns, target_patterns, operations): + self.source_patterns = source_patterns + self.target_patterns = target_patterns + self.operations = operations + + pre_bad = _BadInitConverter.__init__ + patch = _load_patch_function() + patch() + + real_signature = inspect.signature + + def _fake_signature(callable_): + if callable_ is _BadInitConverter.__init__: + raise ValueError("inspect.signature failed mid-loop") + return real_signature(callable_) + + inspect.signature = _fake_signature + try: + legacy = _LegacyConverter(["src.*"], ["tgt.*"], []) + bad = _BadInitConverter.__new__(_BadInitConverter) + bad.source_patterns = ["src.*"] + bad.target_patterns = ["tgt.*"] + bad.operations = [] + with pytest.raises(ValueError): + twc.build_peft_weight_mapping([legacy, bad], "default", None) + finally: + inspect.signature = real_signature + + assert _LegacyConverter.__init__ is pre_legacy + assert _BadInitConverter.__init__ is pre_bad + + +def test_idempotent_install_does_not_double_wrap(): + twc = _install_fake_peft({"build_peft_weight_mapping": _build_that_calls_init}) + patch = _load_patch_function() + patch() + first_wrapped = twc.build_peft_weight_mapping + patch() + assert twc.build_peft_weight_mapping is first_wrapped + + +def test_concurrent_legacy_calls_no_typeerror(): + import time + + def _slow_build(weight_conversions, adapter_name, peft_config = None): + time.sleep(0.05) + return _build_that_calls_init(weight_conversions, adapter_name, peft_config) + + twc = _install_fake_peft({"build_peft_weight_mapping": _slow_build}) + patch = _load_patch_function() + patch() + + errors = [] + results = [] + start = threading.Event() + + def _worker(): + start.wait(timeout = 10) + try: + out = twc.build_peft_weight_mapping( + [_make_legacy_converter()], "default", None + ) + results.append(out) + except Exception as e: + errors.append(e) + + threads = [threading.Thread(target = _worker) for _ in range(8)] + for t in threads: + t.start() + start.set() + for t in threads: + t.join(timeout = 15) + + assert errors == [] + assert len(results) == 8 + for out in results: + assert out[0].distributed_operation == "dist-x" + assert out[0].quantization_operation == "quant-y" + assert _LegacyConverter.__init__.__qualname__.startswith("_LegacyConverter") + + +def test_empty_conversions_short_circuits_without_patching(): + twc = _install_fake_peft({"build_peft_weight_mapping": _build_that_calls_init}) + pre_init = _LegacyConverter.__init__ + patch = _load_patch_function() + patch() + + out = twc.build_peft_weight_mapping([], "default", None) + assert out == [] + assert _LegacyConverter.__init__ is pre_init diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 52114bb544..9db9ae0a32 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -153,6 +153,7 @@ from .import_fixes import ( patch_torchcodec_audio_decoder, disable_torchcodec_if_broken, disable_broken_wandb, + patch_peft_weight_converter_compatibility, ) fix_xformers_performance_issue() @@ -176,6 +177,7 @@ patch_vllm_for_notebooks() patch_torchcodec_audio_decoder() disable_torchcodec_if_broken() disable_broken_wandb() +patch_peft_weight_converter_compatibility() del fix_xformers_performance_issue del fix_vllm_aimv2_issue @@ -197,6 +199,7 @@ del patch_vllm_for_notebooks del patch_torchcodec_audio_decoder del disable_torchcodec_if_broken del disable_broken_wandb +del patch_peft_weight_converter_compatibility # Torch 2.4 has including_emulation if DEVICE_TYPE == "cuda": diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 6cebad7939..f1fba0ce91 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -25,6 +25,7 @@ import textwrap import warnings import sys import functools +import inspect # We cannot do from unsloth_zoo.log import logger since FBGEMM might cause seg faults. UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") in ( @@ -1371,6 +1372,88 @@ def disable_broken_wandb(): os.environ["WANDB_DISABLED"] = "true" +def patch_peft_weight_converter_compatibility(): + """Allow PEFT converter rebuilds on legacy converter constructors.""" + try: + from peft.utils import transformers_weight_conversion as twc + except (ImportError, AttributeError): + return + + if getattr(twc, "_unsloth_weight_converter_compat_patch", False): + return + + import threading + + original_build = twc.build_peft_weight_mapping + patch_lock = threading.RLock() + + def _patch_weight_converter_ctors(weight_conversions, patched): + seen_classes = set() + + for conversion in weight_conversions: + conversion_cls = conversion.__class__ + if conversion_cls in seen_classes: + continue + seen_classes.add(conversion_cls) + + original_init = conversion_cls.__init__ + params = inspect.signature(original_init).parameters + supports_kwargs = any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values() + ) + supports_distributed = "distributed_operation" in params + supports_quantization = "quantization_operation" in params + if supports_kwargs or (supports_distributed and supports_quantization): + continue + + def _compat_init( + self, + *args, + __original_init = original_init, + __supports_distributed = supports_distributed, + __supports_quantization = supports_quantization, + **kwargs, + ): + unsupported = {} + if not __supports_distributed and "distributed_operation" in kwargs: + unsupported["distributed_operation"] = kwargs.pop( + "distributed_operation" + ) + if not __supports_quantization and "quantization_operation" in kwargs: + unsupported["quantization_operation"] = kwargs.pop( + "quantization_operation" + ) + result = __original_init(self, *args, **kwargs) + for name, value in unsupported.items(): + if hasattr(self, name): + setattr(self, name, value) + return result + + conversion_cls.__init__ = _compat_init + patched.append((conversion_cls, original_init)) + + @functools.wraps(original_build) + def _build_peft_weight_mapping_compat( + weight_conversions, + adapter_name, + peft_config = None, + ): + if not weight_conversions: + return original_build(weight_conversions, adapter_name, peft_config) + + patched_classes = [] + with patch_lock: + try: + _patch_weight_converter_ctors(weight_conversions, patched_classes) + return original_build(weight_conversions, adapter_name, peft_config) + finally: + for conversion_cls, original_init in patched_classes: + conversion_cls.__init__ = original_init + + twc.build_peft_weight_mapping = _build_peft_weight_mapping_compat + twc._unsloth_weight_converter_compat_patch = True + + CAUSAL_CONV1D_BROKEN = False _CAUSAL_CONV1D_PREFIX = "causal_conv1d" _CAUSAL_CONV1D_BLOCKER_SENTINEL = "_unsloth_causal_conv1d_blocker" From 146295eeca7d6129bb527f20f1f911ad6316a3b5 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 29 Apr 2026 10:51:25 +0100 Subject: [PATCH 19/54] Studio: Fix clipped model selector text descenders (#5210) * fix: clipped model selector text descenders * Studio: Fix image-only chat requests failing validation (#5212) * fix: allow image-only chat messages * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: deduplicate empty content validation coverage --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> * Fix descender clipping in sidebar user account section Replace `leading-none` with `leading-tight` on the parent div wrapping `displayTitle` and the "Studio" label inside `SidebarMenuButton`. The child spans use `truncate` (overflow: hidden), so `line-height: 1` clipped descenders (g, p, q, y, j) on user names. Same root cause and fix as the model selector trigger. * Add tests for studio text descender clipping * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../frontend/src/components/app-sidebar.tsx | 2 +- .../assistant-ui/model-selector.tsx | 2 +- .../test_studio_text_descender_clipping.py | 69 +++++++++++++++++++ 3 files changed, 71 insertions(+), 2 deletions(-) create mode 100644 tests/studio/test_studio_text_descender_clipping.py diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 6d029a59bd..b1dc286df5 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -527,7 +527,7 @@ export function AppSidebar() { className="!size-8" />
-
+
{displayTitle} Studio
diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 679034c9fa..4fa39beff4 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -81,7 +81,7 @@ function ModelSelectorTrigger({ )} - + {currentModel?.name ?? "Select model"} {currentModel?.description && ( diff --git a/tests/studio/test_studio_text_descender_clipping.py b/tests/studio/test_studio_text_descender_clipping.py new file mode 100644 index 0000000000..dd0df320bc --- /dev/null +++ b/tests/studio/test_studio_text_descender_clipping.py @@ -0,0 +1,69 @@ +""" +Regression guard: descender-prone text spans in Studio must not pair +`leading-none` with `truncate` (overflow: hidden), which clips glyph +descenders (g, p, q, y, j) in real user-visible labels. +""" + +from __future__ import annotations + +import re +from pathlib import Path + + +WORKDIR = Path(__file__).resolve().parents[2] +MODEL_SELECTOR = ( + WORKDIR + / "studio" + / "frontend" + / "src" + / "components" + / "assistant-ui" + / "model-selector.tsx" +) +APP_SIDEBAR = WORKDIR / "studio" / "frontend" / "src" / "components" / "app-sidebar.tsx" + + +def _read(path: Path) -> str: + assert path.exists(), f"missing source file: {path}" + return path.read_text() + + +def test_model_selector_trigger_label_uses_leading_tight(): + src = _read(MODEL_SELECTOR) + pattern = re.compile( + r'', + ) + matches = pattern.findall(src) + assert matches, "could not find sidebar account-block parent div" + leading_classes = [m for m in matches if m.startswith("leading-")] + assert ( + leading_classes + ), f"no leading-* class on sidebar account-block parent: {matches}" + for cls in leading_classes: + assert ( + cls == "leading-tight" + ), f"sidebar account-block must use leading-tight, got: {cls}" + + +def test_no_truncate_plus_leading_none_in_changed_files(): + for path in (MODEL_SELECTOR, APP_SIDEBAR): + src = _read(path) + for line in src.splitlines(): + if "truncate" in line and "leading-none" in line: + raise AssertionError( + f"{path.name}: same line uses truncate + leading-none, descenders will clip: {line.strip()}" + ) From 4f9c8321a2136e62fd86fe722a544afd534334a5 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Wed, 29 Apr 2026 16:45:34 +0530 Subject: [PATCH 20/54] Fix DPO trainer multi process hang (#5199) * Fix DPO trainer multi process hang * Fix datacollator error * further dpo vision changes * cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden DPO vision row processing and source rewrites - dpo_trainer_vision_signature_columns: also match TRL 0.22.x layout (image_sizes followed by ref_chosen_logps), so vision keys are not stripped via remove_unused_columns on the originally-affected version. - dpo_trainer_concatenated_inputs: fall back to inserting after the image_sizes block when no token_type_ids anchor follows it. - Apply the same vision model_kwargs forwarding rewrite to _compute_loss_liger via dpo_trainer_compute_loss_liger so the Liger DPO path does not drop pixel_position_ids/image_position_ids/ mm_token_type_ids when args.use_liger_loss is true. - dpo_trainer_vision_process_row: - guard chosen/rejected EOS append with tokenizer.eos_token_id is not None - use features.get("images") and features.get("prompt") to match the existing get on line 164 and avoid KeyError on rows without those keys - drop the torch.is_tensor gate so list-form pixel_position_ids/ image_position_ids returned without return_tensors are still aliased - skip the loop entry for image_position_ids when it was already promoted to pixel_position_ids, so the output dict no longer carries both keys with identical data - dpo_trainer_data_collator_vision_keys: switch from pad_sequence to trl.trainer.utils.pad with padding_side='left' (matches the DPO collator's prompt left-pad) and padding_value=-1 for *_position_ids keys (sentinel for padded patches), 0 otherwise. Skip the key when not every example carries it. Falls back to pad_sequence if trl.pad is unavailable or the tensor rank is too high. - dpo_trainer_prepare_dataset: keep TRL's writer_batch_size=10 when popping num_proc; removing it defaults to 1000 and reintroduces the vision OOM risk that writer_batch_size=10 was set to avoid. * DPO vision row: keep upstream-facing keys and fix patch padding - dpo_trainer_vision_process_row: no longer aliases image_position_ids to pixel_position_ids. Each upstream-emitted vision key is forwarded under its own name. Gemma4 ForConditionalGeneration.forward accepts image_position_ids directly and renames it to pixel_position_ids only at the vision-tower call site, so aliasing in the row helper hid the kwarg the model actually consumes. - dpo_trainer_vision_process_row: extract pixel_values via "in" membership instead of unconditional indexing. With the missing-images path returning [] to the processor, modern processors no longer emit a pixel_values key, and the previous indexing raised KeyError. - dpo_trainer_data_collator_vision_keys: pick padding_side per key family. *_position_ids tensors are patch-aligned to pixel_values (TRL's DataCollatorForPreference right-pads pixel_values), so pad them right with the -1 sentinel; mm_token_type_ids is token-aligned to prompt_input_ids (left-padded by TRL), so pad it left with 0. * DPO vision: handle multi-image prompts and arbitrary-rank collator pad - dpo_trainer_vision_process_row: when a prompt is missing vision placeholders, insert one placeholder per missing image instead of always inserting a single token. Multi-image rows now satisfy the processor's token-vs-image count check rather than under-inserting and tripping the placeholder/feature mismatch. - dpo_trainer_data_collator_vision_keys: drop the dim()<=2 gate around trl.trainer.utils.pad. trl.pad handles arbitrary rank correctly, while the previous fallback to torch.nn.utils.rnn.pad_sequence raised RuntimeError on rank-3 patch-position tensors with mismatched non-leading dimensions. The pad_sequence path remains as a degraded fallback only when trl.pad is unavailable or raises. * DPO vision row: support scalar images and align prompt-aligned aux ids - dpo_trainer_vision_process_row: type-aware normalization of the features['images'] column instead of a truthiness/len check that raised on single image objects (PIL.Image has no __len__) and on numpy ndarrays (truthiness ambiguous). Lists/tuples count as their length, scalar image objects count as one, None counts as zero, and the original value is forwarded to the processor. - dpo_trainer_vision_process_row: when max_prompt_length truncates prompt_input_ids, also slice token_type_ids and mm_token_type_ids by the same [-max_prompt_length:] suffix. Those keys are 1:1 token aligned to prompt_input_ids (Gemma 4 vision attention keys off mm_token_type_ids per modular_gemma4.py), so leaving them at the original length silently misaligned the multimodal mask. * DPO vision row: stop synthesizing vision-token placeholders Pass features['prompt'] and features['images'] straight to the processor without inserting any extra placeholder tokens. The previous helper used processing_class.image_token, which is the right prompt placeholder for Gemma 4 but the wrong one for Gemma 3 (whose prompt placeholder is boi_token while image_token is the inner expansion target). Synthesizing that token also broke multi-image rows: text ended up with N placeholders while the row helper only forwarded the first image's pixel_values via the standard [0] indexing that mirrors upstream TRL process_row, so token vs image-feature counts diverged. Removing the synthesis matches stock TRL behavior; users provide the correct placeholders for their processor in the prompt. * Add tests for DPO vision row processor passthrough * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han --- .../test_dpo_vision_processor_passthrough.py | 149 ++++++++++ unsloth/models/rl_replacements.py | 272 ++++++++++++++++++ 2 files changed, 421 insertions(+) create mode 100644 tests/python/test_dpo_vision_processor_passthrough.py diff --git a/tests/python/test_dpo_vision_processor_passthrough.py b/tests/python/test_dpo_vision_processor_passthrough.py new file mode 100644 index 0000000000..a4f2e2e12a --- /dev/null +++ b/tests/python/test_dpo_vision_processor_passthrough.py @@ -0,0 +1,149 @@ +"""Verify dpo_trainer_vision_process_row forwards prompt and images verbatim.""" + +import ast +import os + +import numpy as np + + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +RL_PATH = os.path.join(REPO_ROOT, "unsloth", "models", "rl_replacements.py") + + +def _load_helpers(): + src = open(RL_PATH).read() + tree = ast.parse(src) + import torch as _torch + + ns = {"torch": _torch} + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "_DPO_VISION_KEYS" for t in node.targets + ): + exec(ast.get_source_segment(src, node), ns) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name.startswith( + ("dpo_trainer_", "_dpo_trainer_") + ): + exec(ast.get_source_segment(src, node), ns) + return ns + + +class _Tok: + eos_token_id = 99 + bos_token_id = None + + def __call__(self, t, add_special_tokens = False): + return {"input_ids": [10]} + + +class _Capture: + image_token = "" + boi_token = "" + + def __init__(self): + self.tokenizer = _Tok() + self.last_text = None + self.last_images = "__sentinel__" + + def __call__(self, images = None, text = None, add_special_tokens = False): + self.last_text = text + self.last_images = images + out = {"input_ids": [[1, 2]]} + if images is not None: + out["pixel_values"] = [object()] + return out + + +def test_prompt_passes_through_without_image_token_synthesis(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": "describe", "chosen": "c", "rejected": "r", "images": ["i"]}, + proc, + ) + assert proc.last_text == "describe" + + +def test_prompt_with_existing_image_token_unchanged(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": " describe", "chosen": "c", "rejected": "r", "images": ["i"]}, + proc, + ) + assert proc.last_text == " describe" + + +def test_gemma3_style_boi_token_prompt_not_corrupted(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": " describe", "chosen": "c", "rejected": "r", "images": ["i"]}, + proc, + ) + assert proc.last_text == " describe" + assert "" not in proc.last_text + + +def test_multi_image_prompt_unchanged_no_extra_placeholders(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + { + "prompt": "compare", + "chosen": "c", + "rejected": "r", + "images": ["a", "b", "c"], + }, + proc, + ) + assert proc.last_text == "compare" + + +def test_list_images_forwarded_verbatim(): + ns = _load_helpers() + proc = _Capture() + payload = ["a", "b"] + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r", "images": payload}, + proc, + ) + assert proc.last_images is payload + + +def test_single_pil_like_image_forwarded_verbatim(): + ns = _load_helpers() + + class PIL: + def __bool__(self): + return True + + proc = _Capture() + pil = PIL() + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r", "images": pil}, + proc, + ) + assert proc.last_images is pil + + +def test_numpy_ndarray_image_forwarded_verbatim(): + ns = _load_helpers() + proc = _Capture() + arr = np.zeros((2, 3, 3), dtype = np.uint8) + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r", "images": arr}, + proc, + ) + assert proc.last_images is arr + + +def test_missing_images_key_passes_none_to_processor(): + ns = _load_helpers() + proc = _Capture() + ns["dpo_trainer_vision_process_row"]( + {"prompt": "p", "chosen": "c", "rejected": "r"}, + proc, + ) + assert proc.last_images is None diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 4d36af62cc..5d2c4151cf 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -55,6 +55,12 @@ RL_CONFIG_CHANGES = defaultdict(list) RL_METRICS_CHANGES = defaultdict(list) RL_ADDITIONAL_FUNCTIONS = defaultdict(list) +_DPO_VISION_KEYS = ( + "pixel_position_ids", + "image_position_ids", + "mm_token_type_ids", +) + torch_compile_options = { "epilogue_fusion": True, "max_autotune": False, # I saw speedups, but not sure if this has issues in collab @@ -120,6 +126,272 @@ def dpo_trainer_fix_columns(call_args, extra_args): RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_columns) +def dpo_trainer_fix_data_collator(call_args, extra_args): + if ( + "data_collator" in call_args + and "train_dataset" in call_args + and "processing_class" in call_args + ): + fix_collator = ( + "if hasattr(train_dataset, 'column_names'):\n" + " column_names = set(train_dataset.column_names)\n" + " is_dpo_dataset = ({'chosen', 'rejected'}.issubset(column_names) or\n" + " {'prompt_input_ids', 'chosen_input_ids', 'rejected_input_ids'}.issubset(column_names))\n" + " if is_dpo_dataset and isinstance(data_collator, TransformersDataCollatorForLanguageModeling):\n" + " data_collator = None\n" + " del is_dpo_dataset, column_names\n" + ) + return fix_collator + return "" + + +RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_fix_data_collator) + + +def dpo_trainer_vision_process_row( + features, + processing_class, + max_prompt_length = None, + max_completion_length = None, + add_special_tokens = True, + is_chat = False, +): + text = features.get("prompt", "") + images = features.get("images") + processor, tokenizer = processing_class, processing_class.tokenizer + processed_features = processor( + images = images, + text = text, + add_special_tokens = False, + ) + + prompt_input_ids = processed_features["input_ids"][0] + chosen_input_ids = tokenizer(features["chosen"], add_special_tokens = False)[ + "input_ids" + ] + rejected_input_ids = tokenizer(features["rejected"], add_special_tokens = False)[ + "input_ids" + ] + + if add_special_tokens: + if tokenizer.bos_token_id is not None: + prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids + if tokenizer.eos_token_id is not None: + prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] + if not is_chat and tokenizer.eos_token_id is not None: + chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] + rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] + + if max_prompt_length is not None: + prompt_input_ids = prompt_input_ids[-max_prompt_length:] + if max_completion_length is not None: + chosen_input_ids = chosen_input_ids[:max_completion_length] + rejected_input_ids = rejected_input_ids[:max_completion_length] + + output = { + "prompt_input_ids": prompt_input_ids, + "chosen_input_ids": chosen_input_ids, + "rejected_input_ids": rejected_input_ids, + } + if "pixel_values" in processed_features: + output["pixel_values"] = processed_features["pixel_values"][0] + if "pixel_attention_mask" in processed_features: + output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0] + if "image_sizes" in processed_features: + output["image_sizes"] = processed_features["image_sizes"][0] + if "token_type_ids" in processed_features: + token_type_ids = processed_features["token_type_ids"][0] + if max_prompt_length is not None: + token_type_ids = token_type_ids[-max_prompt_length:] + output["token_type_ids"] = token_type_ids + if "pixel_position_ids" in processed_features: + output["pixel_position_ids"] = processed_features["pixel_position_ids"][0] + if "image_position_ids" in processed_features: + output["image_position_ids"] = processed_features["image_position_ids"][0] + if "mm_token_type_ids" in processed_features: + mm_token_type_ids = processed_features["mm_token_type_ids"][0] + if max_prompt_length is not None: + mm_token_type_ids = mm_token_type_ids[-max_prompt_length:] + output["mm_token_type_ids"] = mm_token_type_ids + + return output + + +def dpo_trainer_vision_signature_columns(function_name, function): + if function_name != "_set_signature_columns_if_needed": + return function + + if all(_k in function for _k in _DPO_VISION_KEYS): + return function + + _extra_columns = "".join(f' "{_k}",\n' for _k in _DPO_VISION_KEYS) + new_function = function.replace( + ' "image_sizes",\n' ' "token_type_ids",\n', + f' "image_sizes",\n' + f"{_extra_columns}" + f' "token_type_ids",\n', + ) + if new_function != function: + return new_function + return function.replace( + ' "image_sizes",\n' ' "ref_chosen_logps",\n', + f' "image_sizes",\n' + f"{_extra_columns}" + f' "ref_chosen_logps",\n', + ) + + +def dpo_trainer_concatenated_inputs(function_name, function): + if function_name != "concatenated_inputs": + return function + + if all(_k in function for _k in _DPO_VISION_KEYS): + return function + + _extra_inputs = "".join( + f' if "{_k}" in batch:\n' + f' output["{_k}"] = torch.cat((batch["{_k}"], batch["{_k}"]), dim=0)\n' + for _k in _DPO_VISION_KEYS + ) + + image_sizes_block = ( + ' if "image_sizes" in batch:\n' + ' output["image_sizes"] = torch.cat([batch["image_sizes"], batch["image_sizes"]], dim=0)\n' + ) + new_function = function.replace( + image_sizes_block + ' if "token_type_ids" in batch:\n', + image_sizes_block + _extra_inputs + ' if "token_type_ids" in batch:\n', + ) + if new_function != function: + return new_function + if image_sizes_block in function: + return function.replace(image_sizes_block, image_sizes_block + _extra_inputs, 1) + return function + + +def _dpo_trainer_extend_vision_model_kwargs(function): + if all(_k in function for _k in _DPO_VISION_KEYS): + return function + + _extra_forward = "".join( + f' if "{_k}" in concatenated_batch:\n' + f' model_kwargs["{_k}"] = concatenated_batch["{_k}"]\n' + for _k in ( + "pixel_values", + "pixel_attention_mask", + "image_sizes", + *_DPO_VISION_KEYS, + ) + ) + + return function.replace( + ' if "pixel_values" in concatenated_batch:\n' + ' model_kwargs["pixel_values"] = concatenated_batch["pixel_values"]\n' + ' if "pixel_attention_mask" in concatenated_batch:\n' + ' model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"]\n' + ' if "image_sizes" in concatenated_batch:\n' + ' model_kwargs["image_sizes"] = concatenated_batch["image_sizes"]\n', + f"{_extra_forward}", + ) + + +def dpo_trainer_concatenated_forward(function_name, function): + if function_name != "concatenated_forward": + return function + return _dpo_trainer_extend_vision_model_kwargs(function) + + +def dpo_trainer_compute_loss_liger(function_name, function): + if function_name != "_compute_loss_liger": + return function + return _dpo_trainer_extend_vision_model_kwargs(function) + + +def dpo_trainer_data_collator_vision_keys(call_args, extra_args): + if "data_collator" not in call_args: + return "" + + _vision_keys = str(_DPO_VISION_KEYS) + return ( + "from trl.trainer.dpo_trainer import DataCollatorForPreference\n" + "if not hasattr(DataCollatorForPreference, '_unsloth_vision_keys_patch'):\n" + " _old_dpo_collator_torch_call = DataCollatorForPreference.torch_call\n" + "\n" + " def _unsloth_dpo_torch_call(self, examples):\n" + " output = _old_dpo_collator_torch_call(self, examples)\n" + " import torch as _unsloth_torch\n" + " try:\n" + " from trl.trainer.utils import pad as _unsloth_trl_pad\n" + " except Exception:\n" + " _unsloth_trl_pad = None\n" + " for _k in " + _vision_keys + ":\n" + " if not all(_k in example for example in examples):\n" + " continue\n" + " _is_position_key = _k.endswith('position_ids')\n" + " _padding_value = -1 if _is_position_key else 0\n" + " _padding_side = 'right' if _is_position_key else 'left'\n" + " _values = [_unsloth_torch.as_tensor(example[_k]) for example in examples]\n" + " try:\n" + " if _unsloth_trl_pad is not None:\n" + " output[_k] = _unsloth_trl_pad(_values, padding_value=_padding_value, padding_side=_padding_side)\n" + " else:\n" + " from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence\n" + " output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value)\n" + " except Exception:\n" + " from torch.nn.utils.rnn import pad_sequence as _unsloth_pad_sequence\n" + " output[_k] = _unsloth_pad_sequence(_values, batch_first=True, padding_value=_padding_value)\n" + " return output\n" + "\n" + " DataCollatorForPreference.torch_call = _unsloth_dpo_torch_call\n" + " DataCollatorForPreference._unsloth_vision_keys_patch = True\n" + ) + + +def dpo_trainer_prepare_dataset(function_name, function): + if function_name != "_prepare_dataset": + return function + + legacy_call = "self.tokenize_row if not self.is_vision_model else self.process_row" + if legacy_call not in function: + return function + + function = function.replace( + legacy_call, + "self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row", + ) + + legacy_tokenize_block = ( + " # Tokenize the dataset\n" + " if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`\n" + ' map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"\n' + "\n" + " dataset = dataset.map(\n" + " self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row,\n" + ) + patched_tokenize_block = ( + " # Tokenize the dataset\n" + " if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc`\n" + ' map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"\n' + " if self.is_vision_model:\n" + ' map_kwargs.pop("num_proc", None)\n' + "\n" + " dataset = dataset.map(\n" + " self.tokenize_row if not self.is_vision_model else dpo_trainer_vision_process_row,\n" + ) + if legacy_tokenize_block in function: + function = function.replace(legacy_tokenize_block, patched_tokenize_block, 1) + return function + + +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_prepare_dataset) +RL_PRE_ITEMS["dpo_trainer"].append(inspect.getsource(dpo_trainer_vision_process_row)) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_vision_signature_columns) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_concatenated_inputs) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_concatenated_forward) +RL_FUNCTIONS["dpo_trainer"].append(dpo_trainer_compute_loss_liger) +RL_EXTRA_ARGS["dpo_trainer"].append(dpo_trainer_data_collator_vision_keys) + + # Fix tokenizer double BOS def sft_trainer_prepare_dataset(function_name, function): if ( From 4ab5378d28f8be14b0f6aaa93b5d07e8b82c00f1 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Thu, 30 Apr 2026 12:50:23 +0100 Subject: [PATCH 21/54] Studio: Pin assistant-ui core for fresh installs (#5229) * fix(studio): pin assistant-ui core for fresh installs * fix(studio): use assistant-ui internal export --- studio/frontend/package.json | 1 + .../frontend/src/features/chat/utils/delete-thread-message.ts | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/studio/frontend/package.json b/studio/frontend/package.json index c5cb949ccd..7bea824071 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -16,6 +16,7 @@ "biome:fix": "biome check . --write" }, "dependencies": { + "@assistant-ui/core": "0.1.17", "@assistant-ui/react": "^0.12.19", "@assistant-ui/react-markdown": "^0.12.3", "@assistant-ui/react-streamdown": "^0.1.2", diff --git a/studio/frontend/src/features/chat/utils/delete-thread-message.ts b/studio/frontend/src/features/chat/utils/delete-thread-message.ts index a6d6556a87..9901ccf80e 100644 --- a/studio/frontend/src/features/chat/utils/delete-thread-message.ts +++ b/studio/frontend/src/features/chat/utils/delete-thread-message.ts @@ -16,7 +16,7 @@ import type { * delete + reload smoke tests; the path or API may change without a semver signal on “public” * surface area. */ -import { MessageRepository } from "@assistant-ui/core/runtime/utils/message-repository"; +import { MessageRepository } from "@assistant-ui/core/internal"; import { db } from "@/features/chat/db"; import type { MessageRecord } from "@/features/chat/types"; From 11c04ed6320e5a46d96684a8bfd925189a6458a9 Mon Sep 17 00:00:00 2001 From: Anish Umale Date: Thu, 30 Apr 2026 19:43:20 +0530 Subject: [PATCH 22/54] Fix local model scanner to handle ollama cloud models (#5220) * fix _scan_ollama_dir to handle ollama cloud models correctly * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/routes/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index db27ce1907..c312b0ee1d 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -623,7 +623,7 @@ def _scan_ollama_dir( gguf_link_path: Optional[str] = None quant = f"-{file_type}" if file_type else "" safe_name = repo_name.replace("/", "-") - for layer in manifest.get("layers", []): + for layer in manifest.get("layers") or []: media = layer.get("mediaType", "") digest = layer.get("digest", "") if not digest: From 507417579f010c6b730ad097459906f075db64d6 Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Thu, 30 Apr 2026 17:40:39 +0200 Subject: [PATCH 23/54] Fix Studio desktop tray installer and titlebar and bux fixes (#5179) * fix(tauri): dedupe tray and brand nsis installer * feat(tauri): add linux windows custom titlebar * Fix desktop auth gate after backend startup * Fix desktop installer assets and setup script skew * Scope setup failure exit to Tauri installer * fix desktop updater production channel * fix desktop auth runtime installer regressions * fix desktop dev cors retry * fix tauri process generation race * feat desktop diagnostics support report * fix tauri apt update best effort * Fix Windows desktop NSIS installer upgrades * Start managed backend after desktop install * Improve NSIS installer branding resolution * Fix assistant-ui internal import * Fix desktop release workflow * Keep desktop auth retry on cached backend --------- Co-authored-by: wasimysaid --- .github/workflows/release-desktop.yml | 45 +- install.ps1 | 217 +++- install.sh | 157 ++- studio/backend/main.py | 2 + studio/frontend/package.json | 1 + studio/frontend/src/app/auth-guards.ts | 5 +- studio/frontend/src/app/provider.tsx | 88 +- studio/frontend/src/app/routes/__root.tsx | 2 +- .../src/components/tauri/startup-screen.tsx | 128 ++- .../src/components/tauri/update-banner.tsx | 90 +- .../src/components/tauri/update-screen.tsx | 48 +- .../src/components/tauri/window-titlebar.tsx | 329 ++++++ studio/frontend/src/features/auth/api.ts | 46 +- .../features/auth/change-password-page.tsx | 2 +- .../frontend/src/features/auth/login-page.tsx | 2 +- .../src/features/auth/tauri-auto-auth.ts | 27 +- .../chat/utils/delete-thread-message.ts | 4 +- .../data-recipes/pages/data-recipes-page.tsx | 2 +- .../data-recipes/pages/edit-recipe-page.tsx | 2 +- .../src/features/export/export-page.tsx | 2 +- .../onboarding/components/wizard-layout.tsx | 2 +- .../recipe-studio/recipe-studio-page.tsx | 2 +- .../src/features/studio/studio-page.tsx | 2 +- .../frontend/src/hooks/use-tauri-backend.ts | 138 ++- studio/frontend/src/hooks/use-tauri-update.ts | 115 +- studio/frontend/src/lib/tauri-diagnostics.ts | 176 ++++ studio/setup.ps1 | 16 +- studio/src-tauri/Cargo.lock | 355 ++++++- studio/src-tauri/Cargo.toml | 3 +- studio/src-tauri/capabilities/default.json | 8 +- studio/src-tauri/src/commands.rs | 227 +++- studio/src-tauri/src/desktop_auth.rs | 169 ++- studio/src-tauri/src/diagnostics/mod.rs | 169 +++ studio/src-tauri/src/diagnostics/phase_log.rs | 823 +++++++++++++++ studio/src-tauri/src/diagnostics/redaction.rs | 198 ++++ studio/src-tauri/src/diagnostics/report.rs | 607 +++++++++++ studio/src-tauri/src/diagnostics/state.rs | 774 ++++++++++++++ studio/src-tauri/src/install.rs | 416 +++++++- studio/src-tauri/src/main.rs | 73 +- studio/src-tauri/src/process.rs | 202 +++- studio/src-tauri/src/update.rs | 139 ++- studio/src-tauri/tauri.conf.json | 40 +- studio/src-tauri/tauri.windows.conf.json | 6 - .../windows/branding/nsis-header.bmp | Bin 25818 -> 102654 bytes .../windows/branding/nsis-sidebar.bmp | Bin 154542 -> 618006 bytes studio/src-tauri/windows/hooks.nsh | 9 +- studio/src-tauri/windows/installer.nsi | 994 ++++++++++++++++++ 47 files changed, 6487 insertions(+), 375 deletions(-) create mode 100644 studio/frontend/src/components/tauri/window-titlebar.tsx create mode 100644 studio/frontend/src/lib/tauri-diagnostics.ts create mode 100644 studio/src-tauri/src/diagnostics/mod.rs create mode 100644 studio/src-tauri/src/diagnostics/phase_log.rs create mode 100644 studio/src-tauri/src/diagnostics/redaction.rs create mode 100644 studio/src-tauri/src/diagnostics/report.rs create mode 100644 studio/src-tauri/src/diagnostics/state.rs create mode 100644 studio/src-tauri/windows/installer.nsi diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml index 2d1c6f51f2..ea82739968 100644 --- a/.github/workflows/release-desktop.yml +++ b/.github/workflows/release-desktop.yml @@ -54,10 +54,43 @@ jobs: with: node-version: 24 + - name: Install pinned Tauri CLI + run: npm install --save-dev --prefix studio @tauri-apps/cli@2.10.1 + + - name: Verify pinned Tauri CLI + shell: bash + run: | + out="$(npx --prefix studio tauri --version)" + echo "$out" + if [ "$out" != "tauri-cli 2.10.1" ]; then + echo "Expected tauri-cli 2.10.1, got $out" >&2 + exit 1 + fi + - name: Install frontend dependencies working-directory: studio/frontend run: npm install + - name: Verify backend package is published + shell: bash + run: | + node <<'JS' + const { readFileSync } = require('node:fs'); + + (async () => { + const cargo = readFileSync('studio/src-tauri/Cargo.toml', 'utf8'); + const match = cargo.match(/^version\s*=\s*"([^"]+)"/m); + if (!match) throw new Error('Could not read desktop app version'); + + const appVersion = match[1]; + const response = await fetch(`https://pypi.org/pypi/unsloth/${appVersion}/json`); + if (!response.ok) { + const message = 'Publish unsloth=={app_version} to PyPI before the desktop release'; + throw new Error(`${message.replace('{app_version}', appVersion)} (HTTP ${response.status})`); + } + })(); + JS + # ── Rust ── - name: Install Rust stable uses: dtolnay/rust-toolchain@stable @@ -105,13 +138,14 @@ jobs: # ── Linux: build + sign + upload ── - name: Build Linux app if: matrix.platform == 'ubuntu-22.04' - uses: tauri-apps/tauri-action@fce9c6108b31ea247710505d3aaaa893ee6768d4 + uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} with: projectPath: studio + tauriScript: npx --prefix . tauri tagName: desktop-v__VERSION__ releaseName: 'Unsloth Studio (Desktop) v__VERSION__' releaseBody: | @@ -123,6 +157,7 @@ jobs: > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} prerelease: false args: -v ${{ matrix.args }} @@ -130,7 +165,7 @@ jobs: # ── macOS: build + sign + notarize + upload ── - name: Build macOS app if: matrix.platform == 'macos-latest' - uses: tauri-apps/tauri-action@fce9c6108b31ea247710505d3aaaa893ee6768d4 + uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} @@ -141,6 +176,7 @@ jobs: APPLE_TEAM_ID: ${{ secrets.APPLE_TEAM_ID }} with: projectPath: studio + tauriScript: npx --prefix . tauri tagName: desktop-v__VERSION__ releaseName: 'Unsloth Studio (Desktop) v__VERSION__' releaseBody: | @@ -152,6 +188,7 @@ jobs: > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} prerelease: false args: -v ${{ matrix.args }} @@ -159,7 +196,7 @@ jobs: # ── Windows: build + sign + upload ── - name: Build Windows app if: matrix.platform == 'windows-latest' - uses: tauri-apps/tauri-action@fce9c6108b31ea247710505d3aaaa893ee6768d4 + uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} @@ -171,6 +208,7 @@ jobs: AZURE_CERTIFICATE_PROFILE_NAME: ${{ secrets.AZURE_CERTIFICATE_PROFILE_NAME }} with: projectPath: studio + tauriScript: npx --prefix . tauri tagName: desktop-v__VERSION__ releaseName: 'Unsloth Studio (Desktop) v__VERSION__' releaseBody: | @@ -182,6 +220,7 @@ jobs: > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. releaseDraft: ${{ inputs.draft }} prerelease: false args: -v ${{ matrix.args }} diff --git a/install.ps1 b/install.ps1 index 2544eaf78a..dcd86eb019 100644 --- a/install.ps1 +++ b/install.ps1 @@ -8,6 +8,79 @@ function Install-UnslothStudio { $ErrorActionPreference = "Stop" $script:UnslothVerbose = ($env:UNSLOTH_VERBOSE -eq "1") + # ── Tauri structured output ── + function Write-TauriLog { + param([string]$Tag, [string]$Message) + if ($TauriMode) { + Write-Host "[TAURI:$Tag] $Message" + } + } + + function Format-TauriDiagBool { + param([bool]$Value) + if ($Value) { return "true" } + return "false" + } + + function Get-TauriDiagArch { + $arch = [string]$env:PROCESSOR_ARCHITECTURE + if ([string]::IsNullOrWhiteSpace($arch)) { + try { $arch = [System.Runtime.InteropServices.RuntimeInformation]::OSArchitecture.ToString() } catch { $arch = "unknown" } + } + $arch = $arch.ToLowerInvariant() + switch ($arch) { + "amd64" { return "x86_64" } + "x64" { return "x86_64" } + "arm64" { return "arm64" } + "x86" { return "x86" } + default { return ($arch -replace '[^a-z0-9_.-]', '_') } + } + } + + function Get-TauriTorchIndexFamily { + param([string]$TorchIndexUrl) + if ($SkipTorch) { return "none" } + if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" } + $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf } + if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf } + return "auto" + } + + function Get-TauriGpuBranch { + param([string]$TorchIndexFamily) + if ($SkipTorch) { return "no_torch" } + if ($TorchIndexFamily -like "cu*") { return "cuda" } + if ($TorchIndexFamily -like "rocm*") { return "rocm" } + if ($TorchIndexFamily -eq "cpu") { return "cpu" } + return "unknown" + } + + function Write-TauriDiag { + param( + [string]$GpuBranch = "unknown", + [string]$TorchIndexFamily = "none", + [string]$PythonVersionForDiag = $PythonVersion + ) + if ([string]::IsNullOrWhiteSpace($PythonVersionForDiag)) { $PythonVersionForDiag = "unknown" } + Write-TauriLog "DIAG" "diag_schema=1 platform=windows arch=$(Get-TauriDiagArch) python_version=$($PythonVersionForDiag.ToLowerInvariant()) skip_torch=$(Format-TauriDiagBool $SkipTorch) mac_intel=false gpu_branch=$GpuBranch torch_index_family=$TorchIndexFamily" + } + + function Exit-InstallFailure { + param( + [Parameter(Mandatory = $true)][string]$Message, + [int]$Code = 1 + ) + if ($Code -eq 0) { $Code = 1 } + Write-TauriLog "ERROR" $Message + if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) { + Restore-StudioVenvRollback + } + if ($TauriMode) { + exit $Code + } + } + # ── Parse flags ── $StudioLocalInstall = $false $PackageName = "unsloth" @@ -26,7 +99,7 @@ function Install-UnslothStudio { $i++ if ($i -ge $argList.Count) { Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red - return + return (Exit-InstallFailure "--package requires an argument.") } $PackageName = $argList[$i] } @@ -42,22 +115,14 @@ function Install-UnslothStudio { $RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) { Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red - return + return (Exit-InstallFailure "--local must be run from the unsloth repo root") } } # Validate --package to prevent injection into shell/Python commands if ($PackageName -notmatch '^[a-zA-Z0-9][a-zA-Z0-9._-]*$') { Write-Host "[ERROR] --package name contains invalid characters (allowed: a-z A-Z 0-9 . _ -)" -ForegroundColor Red - return - } - - # ── Tauri structured output ── - function Write-TauriLog { - param([string]$Tag, [string]$Message) - if ($TauriMode) { - Write-Host "[TAURI:$Tag] $Message" - } + return (Exit-InstallFailure "--package name contains invalid characters") } $PythonVersion = "3.13" @@ -630,7 +695,7 @@ shell.Run cmd, 0, False step "winget" "not available" "Red" substep "Install it from https://aka.ms/getwinget" "Yellow" substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow" - return + return (Exit-InstallFailure "winget is not available") } # ── Helper: detect a working Python 3.11-3.13 on the system ── @@ -749,9 +814,14 @@ shell.Run cmd, 0, False Write-Host " Please install Python $PythonVersion manually from https://www.python.org/downloads/" -ForegroundColor Yellow Write-Host " Make sure to check 'Add Python to PATH' during installation." -ForegroundColor Yellow Write-Host " Then re-run this installer." -ForegroundColor Yellow - return + return (Exit-InstallFailure "Python installation failed") } } + $DiagPythonVersion = $PythonVersion + if ($DetectedPython) { $DiagPythonVersion = $DetectedPython.Version } + $InitialGpuBranch = "unknown" + if ($SkipTorch) { $InitialGpuBranch = "no_torch" } + Write-TauriDiag -GpuBranch $InitialGpuBranch -TorchIndexFamily "none" -PythonVersionForDiag $DiagPythonVersion # ── Install uv if not present ── Write-TauriLog "STEP" "Installing uv package manager" @@ -773,7 +843,7 @@ shell.Run cmd, 0, False if (-not (Get-Command uv -ErrorAction SilentlyContinue)) { step "uv" "could not be installed" "Red" substep "Install it from https://docs.astral.sh/uv/" "Yellow" - return + return (Exit-InstallFailure "uv could not be installed") } # ── Create venv (migrate old layout if possible, otherwise fresh) ── @@ -786,11 +856,68 @@ shell.Run cmd, 0, False $VenvPython = Join-Path $VenvDir "Scripts\python.exe" $_Migrated = $false + $script:StudioVenvRollbackDir = $null + $script:StudioVenvRollbackTarget = $VenvDir + $script:StudioVenvRollbackActive = $false + + function Start-StudioVenvRollback { + param([Parameter(Mandatory = $true)][string]$ExistingDir) + $stamp = Get-Date -Format "yyyyMMddHHmmss" + $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID" + $suffix = 0 + while (Test-Path $candidate) { + $suffix++ + $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix" + } + Move-Item -Path $ExistingDir -Destination $candidate -ErrorAction Stop + $script:StudioVenvRollbackDir = $candidate + $script:StudioVenvRollbackTarget = $ExistingDir + $script:StudioVenvRollbackActive = $true + substep "previous environment preserved for rollback" + } + + function Restore-StudioVenvRollback { + if (-not $script:StudioVenvRollbackActive) { return } + $backup = $script:StudioVenvRollbackDir + $target = $script:StudioVenvRollbackTarget + if (-not $backup -or -not (Test-Path $backup)) { + $script:StudioVenvRollbackActive = $false + return + } + substep "restoring previous environment after failed install..." "Yellow" + try { + if (Test-Path $target) { + Remove-Item -Recurse -Force $target -ErrorAction SilentlyContinue + } + Move-Item -Path $backup -Destination $target -Force -ErrorAction Stop + substep "restored previous environment" + $script:StudioVenvRollbackActive = $false + $script:StudioVenvRollbackDir = $null + } catch { + Write-Host "[WARN] Could not restore previous environment from $backup to $target" -ForegroundColor Yellow + Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow + } + } + + function Complete-StudioVenvRollback { + if (-not $script:StudioVenvRollbackActive) { return } + $backup = $script:StudioVenvRollbackDir + if ($backup -and (Test-Path $backup)) { + Remove-Item -Recurse -Force $backup -ErrorAction SilentlyContinue + } + $script:StudioVenvRollbackActive = $false + $script:StudioVenvRollbackDir = $null + } if (Test-Path $VenvPython) { - # New layout already exists -- nuke for fresh install - substep "removing existing environment for fresh install..." - Remove-Item -Recurse -Force $VenvDir + # New layout already exists -- replace only after preserving rollback copy. + substep "preserving existing environment for rollback..." + try { + Start-StudioVenvRollback -ExistingDir $VenvDir + } catch { + Write-Host "[ERROR] Could not prepare existing environment for reinstall: $($_.Exception.Message)" -ForegroundColor Red + return (Exit-InstallFailure "Could not prepare existing environment for reinstall") + } } elseif (Test-Path (Join-Path $StudioHome ".venv\Scripts\python.exe")) { # Old layout (~/.unsloth/studio/.venv) exists -- validate before migrating $OldVenv = Join-Path $StudioHome ".venv" @@ -799,18 +926,23 @@ shell.Run cmd, 0, False $prevEAP2 = $ErrorActionPreference $ErrorActionPreference = "Continue" try { - & $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null - $torchOk = ($LASTEXITCODE -eq 0) - } catch { $torchOk = $false } + if ($SkipTorch) { + & $OldPy -c "import sys; print(sys.executable)" 2>$null | Out-Null + } else { + & $OldPy -c "import torch; A = torch.ones((2,2)); B = A + A" 2>$null | Out-Null + } + $legacyOk = ($LASTEXITCODE -eq 0) + } catch { $legacyOk = $false } $ErrorActionPreference = $prevEAP2 - if ($torchOk) { + if ($legacyOk) { substep "legacy environment is healthy -- migrating..." Move-Item -Path $OldVenv -Destination $VenvDir -Force substep "moved .venv -> unsloth_studio" $_Migrated = $true } else { substep "legacy environment failed validation -- creating fresh environment" "Yellow" - Remove-Item -Recurse -Force $OldVenv -ErrorAction SilentlyContinue + $invalidVenv = Join-Path $StudioHome (".venv.invalid.{0}.{1}" -f (Get-Date -Format "yyyyMMddHHmmss"), $PID) + Move-Item -Path $OldVenv -Destination $invalidVenv -Force -ErrorAction SilentlyContinue } } elseif (Test-Path (Join-Path $env:USERPROFILE "unsloth_studio\Scripts\python.exe")) { # CWD-relative venv from old install.ps1 -- migrate to absolute path @@ -826,9 +958,8 @@ shell.Run cmd, 0, False substep "$VenvDir" $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" } if ($venvExit -ne 0) { - Write-TauriLog "ERROR" "Failed to create virtual environment (exit code $venvExit)" Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to create virtual environment (exit code $venvExit)" $venvExit) } } else { step "venv" "using migrated environment" @@ -886,6 +1017,9 @@ shell.Run cmd, 0, False return "$baseUrl/cu126" } $TorchIndexUrl = Get-TorchIndexUrl + $TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl + $GpuBranch = Get-TauriGpuBranch $TorchIndexFamily + Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version # ── Print CPU-only hint when no GPU detected ── if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") { @@ -946,14 +1080,14 @@ shell.Run cmd, 0, False } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) } if ($StudioLocalInstall) { substep "overlaying local repo (editable)..." $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } } } elseif ($TorchIndexUrl) { @@ -964,9 +1098,8 @@ shell.Run cmd, 0, False substep "installing PyTorch ($TorchIndexUrl)..." $torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } if ($torchInstallExit -ne 0) { - Write-TauriLog "ERROR" "Failed to install PyTorch (exit code $torchInstallExit)" Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) } } @@ -988,9 +1121,8 @@ shell.Run cmd, 0, False $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } if ($baseInstallExit -ne 0) { - Write-TauriLog "ERROR" "Failed to install unsloth (exit code $baseInstallExit)" Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) } if ($StudioLocalInstall) { @@ -998,7 +1130,7 @@ shell.Run cmd, 0, False $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } } } else { @@ -1009,20 +1141,19 @@ shell.Run cmd, 0, False $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.4.8" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) } substep "overlaying local repo (editable)..." $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --torch-backend=auto -- "$PackageName" } if ($baseInstallExit -ne 0) { - Write-TauriLog "ERROR" "Failed to install unsloth (exit code $baseInstallExit)" Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) } } } @@ -1077,12 +1208,11 @@ shell.Run cmd, 0, False step "setup" "running unsloth studio setup..." $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" if (-not (Test-Path $UnslothExe)) { - Write-TauriLog "ERROR" "unsloth CLI was not installed correctly" Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow - return + return (Exit-InstallFailure "unsloth CLI was not installed correctly") } # Tell setup.ps1 to skip base package installation (install.ps1 already did it) $env:SKIP_STUDIO_BASE = "1" @@ -1104,12 +1234,16 @@ shell.Run cmd, 0, False # and bypass the fast-path version check from PR #4667. $studioArgs = @('studio', 'setup') if ($script:UnslothVerbose) { $studioArgs += '--verbose' } - & $UnslothExe @studioArgs - $setupExit = $LASTEXITCODE + $env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1" + try { + & $UnslothExe @studioArgs + $setupExit = $LASTEXITCODE + } finally { + Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue + } if ($setupExit -ne 0) { - Write-TauriLog "ERROR" "unsloth studio setup failed (exit code $setupExit)" Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red - return + return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit) } # ── Expose `unsloth` via a shim dir containing only unsloth.exe ── @@ -1182,6 +1316,7 @@ shell.Run cmd, 0, False step "path" "added unsloth launcher to PATH" } Refresh-SessionPath # sync current session with registry + Complete-StudioVenvRollback # ── Tauri mode: done, skip shortcuts and auto-launch ── if ($TauriMode) { diff --git a/install.sh b/install.sh index 6c28b6eda4..18c6bba589 100755 --- a/install.sh +++ b/install.sh @@ -162,9 +162,119 @@ tauri_log() { fi } +tauri_diag_marker() { + _diag_gpu_branch="${1:-unknown}" + _diag_torch_index_family="${2:-none}" + tauri_log "DIAG" "diag_schema=1 platform=${OS:-unknown} arch=${_ARCH:-unknown} python_version=${PYTHON_VERSION:-unknown} skip_torch=${SKIP_TORCH:-false} mac_intel=${MAC_INTEL:-false} gpu_branch=${_diag_gpu_branch} torch_index_family=${_diag_torch_index_family}" +} + +_tauri_torch_index_family() { + if [ "${SKIP_TORCH:-false}" = true ]; then + echo "none" + return + fi + _diag_url="${1:-}" + case "$_diag_url" in + */cu118) echo "cu118" ;; + */cu124) echo "cu124" ;; + */cu126) echo "cu126" ;; + */cu128) echo "cu128" ;; + */cu130) echo "cu130" ;; + */cpu) echo "cpu" ;; + */rocm[0-9]*.[0-9]*) + _diag_family=${_diag_url##*/} + case "$_diag_family" in + rocm[0-9]*.[0-9]*) echo "$_diag_family" ;; + *) echo "auto" ;; + esac ;; + "") echo "none" ;; + *) echo "auto" ;; + esac +} + +_tauri_gpu_branch() { + _diag_family="${1:-unknown}" + _diag_radeon="${2:-false}" + if [ "${SKIP_TORCH:-false}" = true ]; then + echo "no_torch" + return + fi + if [ "${OS:-}" = "macos" ]; then + echo "mac" + return + fi + case "$_diag_family" in + cu*) echo "cuda" ;; + rocm*) + if [ "$_diag_radeon" = true ]; then + echo "rocm_radeon" + else + echo "rocm" + fi ;; + radeon) echo "rocm_radeon" ;; + cpu) echo "cpu" ;; + none) echo "no_torch" ;; + *) echo "unknown" ;; + esac +} + PYTHON_VERSION="" # resolved after platform detection STUDIO_HOME="$HOME/.unsloth/studio" VENV_DIR="$STUDIO_HOME/unsloth_studio" +_VENV_ROLLBACK_DIR="" +_VENV_ROLLBACK_TARGET="$VENV_DIR" +_VENV_ROLLBACK_ACTIVE=false + +_start_studio_venv_replacement() { + _existing_dir="$1" + _stamp=$(date +%Y%m%d%H%M%S 2>/dev/null || echo "time") + _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$" + _suffix=0 + while [ -e "$_candidate" ]; do + _suffix=$((_suffix + 1)) + _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$.$_suffix" + done + mv "$_existing_dir" "$_candidate" + _VENV_ROLLBACK_DIR="$_candidate" + _VENV_ROLLBACK_TARGET="$_existing_dir" + _VENV_ROLLBACK_ACTIVE=true + substep "previous environment preserved for rollback" +} + +_restore_studio_venv_replacement() { + [ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0 + [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ] || { + _VENV_ROLLBACK_ACTIVE=false + return 0 + } + substep "restoring previous environment after failed install..." "$C_WARN" + rm -rf "$_VENV_ROLLBACK_TARGET" + if mv "$_VENV_ROLLBACK_DIR" "$_VENV_ROLLBACK_TARGET"; then + substep "restored previous environment" + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" + else + echo "⚠️ Could not restore previous environment from $_VENV_ROLLBACK_DIR to $_VENV_ROLLBACK_TARGET" >&2 + fi +} + +_commit_studio_venv_replacement() { + [ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0 + if [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ]; then + rm -rf "$_VENV_ROLLBACK_DIR" || true + fi + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" +} + +_on_install_exit() { + _status=$? + if [ "$_status" -ne 0 ]; then + _restore_studio_venv_replacement + fi + exit "$_status" +} +trap _on_install_exit EXIT # ── Helper: download a URL to a file (supports curl and wget) ── download() { @@ -828,6 +938,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then SKIP_TORCH=true fi +_TAURI_INITIAL_GPU_BRANCH="unknown" +if [ "$SKIP_TORCH" = true ]; then + _TAURI_INITIAL_GPU_BRANCH="no_torch" +elif [ "$OS" = "macos" ]; then + _TAURI_INITIAL_GPU_BRANCH="mac" +fi +tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none" + # ── Check system dependencies ── # cmake and git are needed by unsloth studio setup to build the GGUF inference # engine (llama.cpp). build-essential and libcurl-dev are also needed on Linux. @@ -883,9 +1001,15 @@ if [ -n "$MISSING" ]; then if command -v apt-get >/dev/null 2>&1; then _smart_apt_install $MISSING else - echo " apt-get is not available. Please install with your package manager:" + echo " Automatic system package installation is supported on apt-based" + echo " Linux distributions (Ubuntu/Debian) only. Please install the" + echo " missing dependencies with your package manager, then re-run setup:" echo " $MISSING" - echo " Then re-run Unsloth Studio setup." + echo "" + echo " Examples:" + echo " Fedora/RHEL: sudo dnf install cmake git gcc gcc-c++ make libcurl-devel" + echo " Arch: sudo pacman -S --needed cmake git base-devel curl" + echo " openSUSE: sudo zypper install cmake git gcc gcc-c++ make libcurl-devel" exit 1 fi ;; @@ -957,12 +1081,19 @@ mkdir -p "$STUDIO_HOME" _MIGRATED=false if [ -x "$VENV_DIR/bin/python" ]; then - # New layout already exists — nuke for fresh install - rm -rf "$VENV_DIR" + # New layout already exists — replace only after preserving rollback copy. + substep "preserving existing environment for rollback..." + _start_studio_venv_replacement "$VENV_DIR" elif [ -x "$STUDIO_HOME/.venv/bin/python" ]; then - # Old layout exists — validate before migrating + # Old layout exists — validate before migrating. + # In no-torch mode, a missing torch package is expected; validate Python only. substep "found legacy Studio environment, validating..." - if "$STUDIO_HOME/.venv/bin/python" -c " + _legacy_ok=false + if [ "$SKIP_TORCH" = true ]; then + if "$STUDIO_HOME/.venv/bin/python" -c "import sys; print(sys.executable)" >/dev/null 2>&1; then + _legacy_ok=true + fi + elif "$STUDIO_HOME/.venv/bin/python" -c " import torch device = 'cuda' if torch.cuda.is_available() else 'cpu' A = torch.ones((10, 10), device=device) @@ -972,13 +1103,17 @@ D = A + B E = D @ C torch.testing.assert_close(torch.unique(E), torch.tensor((20,), device=E.device, dtype=E.dtype)) " >/dev/null 2>&1; then + _legacy_ok=true + fi + if [ "$_legacy_ok" = true ]; then echo "✅ Legacy environment is healthy — migrating..." mv "$STUDIO_HOME/.venv" "$VENV_DIR" echo " Moved ~/.unsloth/studio/.venv → $VENV_DIR" _MIGRATED=true else echo "⚠️ Legacy environment failed validation — creating fresh environment" - rm -rf "$STUDIO_HOME/.venv" + _invalid_venv="$STUDIO_HOME/.venv.invalid.$(date +%Y%m%d%H%M%S 2>/dev/null || echo time).$$" + mv "$STUDIO_HOME/.venv" "$_invalid_venv" 2>/dev/null || true fi fi @@ -1308,6 +1443,12 @@ case "$TORCH_INDEX_URL" in fi ;; esac +_TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL") +if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then + _TAURI_TORCH_INDEX_FAMILY="radeon" +fi +_TAURI_GPU_BRANCH=$(_tauri_gpu_branch "$_TAURI_TORCH_INDEX_FAMILY" "$_amd_gpu_radeon") +tauri_diag_marker "$_TAURI_GPU_BRANCH" "$_TAURI_TORCH_INDEX_FAMILY" # ── Print CPU-only hint when no GPU detected ── case "$TORCH_INDEX_URL" in @@ -1679,6 +1820,8 @@ if [ "$_SETUP_EXIT" -ne 0 ]; then exit "$_SETUP_EXIT" fi +_commit_studio_venv_replacement + # ── Tauri mode: done, skip shortcuts and auto-launch ── if [ "$TAURI_MODE" = true ]; then tauri_log "DONE" "" diff --git a/studio/backend/main.py b/studio/backend/main.py index 9212404b30..aec335fad3 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -187,6 +187,8 @@ if _api_only: "tauri://localhost", # Linux/macOS Tauri webview "http://tauri.localhost", # Windows Tauri webview "http://localhost", # dev fallback + "http://localhost:5173", # Tauri dev/Vite + "http://127.0.0.1:5173", # Tauri dev/Vite fallback ] _cors_origin_regex = None else: diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 7bea824071..3adc7ed0c8 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -43,6 +43,7 @@ "@tanstack/react-router": "^1.159.10", "@tanstack/react-table": "^8.21.3", "@tauri-apps/api": "^2.10.1", + "@tauri-apps/plugin-clipboard-manager": "^2.3.2", "@tauri-apps/plugin-opener": "^2.5.3", "@tauri-apps/plugin-process": "^2.3.1", "@tauri-apps/plugin-updater": "^2.10.1", diff --git a/studio/frontend/src/app/auth-guards.ts b/studio/frontend/src/app/auth-guards.ts index 509b8f61af..52230f0b6f 100644 --- a/studio/frontend/src/app/auth-guards.ts +++ b/studio/frontend/src/app/auth-guards.ts @@ -9,7 +9,6 @@ import { hasRefreshToken, mustChangePassword, refreshSession, - tauriAutoAuth, } from "@/features/auth"; async function hasActiveSession(): Promise { @@ -39,7 +38,7 @@ function authRedirect(to: "/login" | "/change-password"): never { export async function requireAuth(): Promise { if (isTauri) { - await tauriAutoAuth(); + // AppProvider owns backend startup + desktop auth; route guards run before it mounts. return; } @@ -59,7 +58,6 @@ export async function requireAuth(): Promise { export async function requireGuest(): Promise { if (isTauri) { - await tauriAutoAuth(); throw redirect({ to: "/chat" }); } if (!(await hasActiveSession())) return; @@ -68,7 +66,6 @@ export async function requireGuest(): Promise { export async function requirePasswordChangeFlow(): Promise { if (isTauri) { - await tauriAutoAuth(); throw redirect({ to: "/chat" }); } diff --git a/studio/frontend/src/app/provider.tsx b/studio/frontend/src/app/provider.tsx index b75998a169..e82f611552 100644 --- a/studio/frontend/src/app/provider.tsx +++ b/studio/frontend/src/app/provider.tsx @@ -4,12 +4,18 @@ import { StartupScreen } from "@/components/tauri/startup-screen"; import { UpdateBanner } from "@/components/tauri/update-banner"; import { UpdateScreen } from "@/components/tauri/update-screen"; +import { + WindowTitlebar, + shouldUseCustomWindowTitlebar, +} from "@/components/tauri/window-titlebar"; import { Toaster } from "@/components/ui/sonner"; +import { getTauriAuthFailure, tauriAutoAuth } from "@/features/auth"; import { useTauriBackend } from "@/hooks/use-tauri-backend"; import { useTauriUpdate } from "@/hooks/use-tauri-update"; import { isTauri } from "@/lib/api-base"; +import { useRouterState } from "@tanstack/react-router"; import { ThemeProvider } from "next-themes"; -import { useEffect, useRef, type ReactNode } from "react"; +import { useEffect, useRef, useState, type ReactNode } from "react"; interface AppProviderProps { children: ReactNode; @@ -103,6 +109,7 @@ function TauriUpdateLayer({ isExternalServer }: { isExternalServer: boolean }) { error={update.error} onRetry={update.retryUpdate} onSkipRestart={update.skipAndRestart} + onCopyDiagnostics={update.copyDiagnostics} /> ); } @@ -112,22 +119,34 @@ function TauriUpdateLayer({ isExternalServer }: { isExternalServer: boolean }) { status={update.status} info={update.info} dismissed={update.dismissed} + lastFailure={update.lastFailure} isExternalServer={isExternalServer} onInstall={update.installUpdate} onDismiss={update.dismiss} + onCopyDiagnostics={update.copyDiagnostics} /> ); } +const HIDDEN_TITLEBAR_SIDEBAR_ROUTES = new Set([ + "/onboarding", + "/login", + "/change-password", + "/signup", +]); + function TauriWrapper({ children }: { children: ReactNode }) { + const pathname = useRouterState({ select: (s) => s.location.pathname }); const { status, logs, error, isExternalServer, currentStepIndex, progressDetail, elevationPackages, - startInstall, retry, retryInstall, approveElevation, + startInstall, retry, retryInstall, approveElevation, copyDiagnostics, } = useTauriBackend(); const hasResized = useRef(false); const abortRef = useRef(false); + const [desktopAuthReady, setDesktopAuthReady] = useState(!isTauri); + const [desktopAuthRetry, setDesktopAuthRetry] = useState(0); // Show the window once the frontend mounts (for pre-running states) useEffect(() => { @@ -150,24 +169,79 @@ function TauriWrapper({ children }: { children: ReactNode }) { return () => { abortRef.current = true; }; }, [status]); - if (!isTauri) return <>{children}; - if (status === "running") return <>{children}; + useEffect(() => { + if (!isTauri) { + setDesktopAuthReady(true); + return; + } + if (status !== "running") { + setDesktopAuthReady(false); + setDesktopAuthRetry(0); + return; + } - return ( + let disposed = false; + setDesktopAuthReady(false); + tauriAutoAuth({ force: true }).then((authenticated) => { + if (disposed) return; + if (authenticated) { + setDesktopAuthReady(true); + return; + } + if (!getTauriAuthFailure()) { + window.setTimeout(() => { + if (!disposed) setDesktopAuthRetry((value) => value + 1); + }, 500); + } + }); + + return () => { disposed = true; }; + }, [status, desktopAuthRetry]); + + if (!isTauri) return <>{children}; + + const showApp = status === "running" && desktopAuthReady; + const startupStatus = status === "running" ? "starting" : status; + const startupProgressDetail = + status === "running" && !desktopAuthReady + ? "Signing in to desktop session..." + : progressDetail; + + const content = showApp ? ( + <> + + {children} + + ) : ( ); + + if (!shouldUseCustomWindowTitlebar()) return content; + + const showSidebarSurface = + showApp && !HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname); + + return ( +
+ +
+ {content} +
+
+ ); } export function AppProvider({ children }: AppProviderProps) { diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 19b5763557..22d149473c 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -81,7 +81,7 @@ function RootLayout() { pinned={pinned} setPinned={setPinned} togglePinned={togglePinned} - className="!min-h-0 h-dvh overflow-hidden" + className="!min-h-0 h-[calc(100dvh-var(--studio-titlebar-height,0px))] overflow-hidden" > diff --git a/studio/frontend/src/components/tauri/startup-screen.tsx b/studio/frontend/src/components/tauri/startup-screen.tsx index af5f241416..1eb3f81d10 100644 --- a/studio/frontend/src/components/tauri/startup-screen.tsx +++ b/studio/frontend/src/components/tauri/startup-screen.tsx @@ -3,7 +3,9 @@ import { ShimmerButton } from "@/components/ui/shimmer-button"; import type { BackendStatus } from "@/hooks/use-tauri-backend"; +import type { CopySupportDiagnosticsResult } from "@/lib/tauri-diagnostics"; import { AnimatePresence, motion } from "motion/react"; +import { useState } from "react"; interface StartupScreenProps { status: BackendStatus; @@ -17,6 +19,63 @@ interface StartupScreenProps { onRetryInstall: () => void; onApproveElevation: () => void; onStartServer: () => void; + onCopyDiagnostics: () => Promise; +} + +function DiagnosticsCopyActions({ + onCopyDiagnostics, + children, +}: { + onCopyDiagnostics: () => Promise; + children: React.ReactNode; +}) { + const [copying, setCopying] = useState(false); + const [manualReport, setManualReport] = useState(null); + const [manualMessage, setManualMessage] = useState(null); + + async function handleCopyDiagnostics() { + setCopying(true); + try { + const result = await onCopyDiagnostics(); + if (result.ok) { + setManualReport(null); + setManualMessage(null); + } else { + setManualReport(result.report); + setManualMessage(result.error ?? "Clipboard copy failed. Select and copy the diagnostics below."); + } + } catch (error) { + setManualReport(null); + setManualMessage(`Diagnostics copy failed: ${String(error)}`); + } finally { + setCopying(false); + } + } + + return ( +
+
+ void handleCopyDiagnostics()} + > + {copying ? "Copying..." : "Copy Diagnostics"} + + {children} +
+ {manualMessage && ( +

{manualMessage}

+ )} + {manualReport && ( +