diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2ee7c64e4c..66f79bcb70 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4765,7 +4765,11 @@ def _guard_chat_load_against_training( hf_token = hf_token, max_seq_length = max_seq_length, llama_extra_args = llama_extra_args, - n_parallel = n_parallel, + # The diffusion runner never receives --parallel (load_model hands off + # to _start_diffusion_server before the slot plumbing), so its cache is + # always single-slot; sizing it for more would 409 a load that fits. An + # unclassified GGUF keeps the requested count, which is the safe side. + n_parallel = 1 if diffusion_kind is True else n_parallel, cache_type_kv = cache_type_kv, tensor_parallel = ( _effective_tensor_parallel(llama_extra_args, tensor_parallel) diff --git a/studio/backend/tests/test_parallel_slots_per_load.py b/studio/backend/tests/test_parallel_slots_per_load.py index 2d881fd6cc..ca69ae06fa 100644 --- a/studio/backend/tests/test_parallel_slots_per_load.py +++ b/studio/backend/tests/test_parallel_slots_per_load.py @@ -23,6 +23,7 @@ from __future__ import annotations import inspect import re +import struct import sys import types as _types from pathlib import Path @@ -397,3 +398,112 @@ def test_clamp_sits_between_the_echo_and_the_fit(): assert pending < clamp, "the requested count is captured before the clamp" assert clamp < estimate, "the fit must be estimated from the effective slot count" assert clamp < commit, "the committed effective count is the clamped one" + + +# ── Training-guard sizing ──────────────────────────────────────────── + + +def _write_swa_gguf(path: Path) -> str: + """Smallest DiffusionGemma-shaped header the KV estimator can size: the + canvas marker that routes it to the diffusion runner, plus the sliding-window + dims that make llama.cpp's SWA cache slot-scaled (it allocates + ``n_swa * n_seq_max + n_ubatch`` cells, or one such cache per stream).""" + + def _kv_str(key: str, value: str) -> bytes: + kb, vb = key.encode(), value.encode() + return ( + struct.pack(" bytes: + kb = key.encode() + return struct.pack(" float: + """Run the training guard over a local GGUF and return the size it budgeted.""" + import routes.inference as inf + + seen = {} + + core_training = _types.ModuleType("core.training") + core_training.get_training_backend = lambda: _types.SimpleNamespace( + is_training_active = lambda: True + ) + + def _can_load(**kwargs): + seen.update(kwargs) + return True, {"mode": "single_device"} + + training_vram = _types.ModuleType("routes.training_vram") + training_vram.can_load_chat_during_training = _can_load + monkeypatch.setitem(sys.modules, "core.training", core_training) + monkeypatch.setitem(sys.modules, "routes.training_vram", training_vram) + + monkeypatch.setattr(inf, "_classify_diffusion_gguf", lambda _config: diffusion) + monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda *a, **k: False)) + monkeypatch.setattr(LlamaCppBackend, "_effective_gpu_count", staticmethod(lambda *a, **k: 1)) + monkeypatch.setattr(LlamaCppBackend, "_diffusion_gpu_arg", staticmethod(lambda *a, **k: "0")) + # Pin the --kv-unified probe so the estimate cannot depend on whether this + # machine happens to have a llama-server binary installed. + monkeypatch.setattr( + LlamaCppBackend, + "probe_server_capabilities", + classmethod(lambda cls, binary = None: {}), + ) + + inf._guard_chat_load_against_training( + _types.SimpleNamespace(is_gguf = True, gguf_file = gguf_path, identifier = "local/model"), + model_identifier = "local/model", + hf_token = None, + load_in_4bit = False, + max_seq_length = 8192, + requested_gpu_ids = None, + n_parallel = n_parallel, + gpu_memory_mode = "auto", + ) + return seen["required_override_gb"] + + +def test_training_guard_sizes_a_diffusion_gguf_at_one_slot(monkeypatch, tmp_path): + # The diffusion runner ignores --parallel, so slots must not inflate the + # coexistence estimate and 409 a load that would have fitted beside training. + gguf = _write_swa_gguf(tmp_path / "diffusion.gguf") + one = _guard_required_gb(monkeypatch, gguf, n_parallel = 1, diffusion = True) + many = _guard_required_gb(monkeypatch, gguf, n_parallel = 8, diffusion = True) + assert one == many + + +def test_training_guard_still_sizes_slots_for_an_ordinary_gguf(monkeypatch, tmp_path): + # llama-server really does allocate per-slot SWA cells, so the reduction + # above must be scoped to diffusion and not flatten every GGUF to one slot. + gguf = _write_swa_gguf(tmp_path / "chat.gguf") + one = _guard_required_gb(monkeypatch, gguf, n_parallel = 1, diffusion = False) + many = _guard_required_gb(monkeypatch, gguf, n_parallel = 8, diffusion = False) + assert many > one + + +def test_training_guard_keeps_slots_for_an_unclassified_gguf(monkeypatch, tmp_path): + # None means the header was inconclusive: it may still be a llama-server + # GGUF, so keep the larger estimate rather than under-size against training. + gguf = _write_swa_gguf(tmp_path / "unknown.gguf") + one = _guard_required_gb(monkeypatch, gguf, n_parallel = 1, diffusion = None) + many = _guard_required_gb(monkeypatch, gguf, n_parallel = 8, diffusion = None) + assert many > one diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 810b390862..82b34cfb8d 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1735,6 +1735,11 @@ async function autoLoadSmallestModel(): Promise<{ ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), kvCacheDtype: loadResp.cache_type_kv ?? null, loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + // Slots are GGUF-only and this branch never sends them, so clear the + // control and its baseline like the interactive/compare load paths: a + // staged override would otherwise be saved for a model it cannot reach. + nParallel: null, + loadedNParallel: null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, // Non-GGUF response: clears any stale GPU baseline a prior manual-GPU @@ -2008,6 +2013,12 @@ async function autoLoadSmallestModel(): Promise<{ ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), kvCacheDtype: loadResp.cache_type_kv ?? null, loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + // The request above deliberately omits n_parallel, so clear both: a + // staged override left over from a preset would otherwise read as + // applied and make the next Apply reload at a count this load never + // sent, against a baseline the status seed fills with the default. + nParallel: null, + loadedNParallel: null, tensorParallel: loadResp.tensor_parallel ?? false, loadedTensorParallel: loadResp.tensor_parallel ?? false, ...loadedGpuMemoryFields(loadResp), diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 7fb8ecdcee..86acefb3a4 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -332,6 +332,12 @@ export function applyActiveModelStatusToStore( (prevState.loadedNParallel === null || hydratingExistingModel) && { loadedNParallel: status.requested_parallel_slots, }), + // Slots are per-model, so a model/variant change underneath this tab blanks + // the control the way performLoad's cross-model reset does. Without it the + // previous model's explicit count follows onto the new model, and saving or + // reloading there pins it. The baseline above still carries the new model's + // resolved count for the rollback. + ...(seedLoadParams && hydratingExistingModel && { nParallel: null }), // Re-seed on first hydration, model/variant changes, or a same-model backend // placement change. gpuStatusFields preserves dirty local edits in the last // case while advancing their loaded baselines. diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index e7d1785ca1..dded1d18ab 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -680,6 +680,55 @@ def test_parallel_slots_setting_wired_end_to_end(): assert 'config.nParallel ?? "",' in sidebar +def test_parallel_slots_control_cleared_when_the_load_never_sent_them(): + """`nParallel` is the editable control ("blank = follow the server default") + and `loadedNParallel` the rollback baseline. Every success path that does + NOT send a slot count must therefore blank the control, or a value staged + for some other model shows as applied, is persisted into that model's + per-model config (`isDefaultConfig` keys on nParallel, so a phantom count + turns a Save-nothing click into a stored override) and is re-sent by the + next Apply. Three paths were missing the clear; each assertion below is the + only thing pinning one of them.""" + status = _read("features/chat/lib/apply-inference-status-to-store.ts") + # A model/variant swap underneath this tab (another client, the CLI) must + # reset the control like performLoad's cross-model reset does, or model A's + # explicit count follows onto model B. + assert ( + "...(seedLoadParams && hydratingExistingModel && { nParallel: null })," + in status + ) + # ... while still never adopting the RESOLVED echo into the control. + assert "nParallel: status.requested_parallel_slots," not in status + + adapter = _read("features/chat/api/chat-adapter.ts") + # Slice the two success branches of loadAutoLoadCandidate apart, and bound + # the second one at the shared tail, or it would swallow the fresh-default + # path below and stay green when this branch loses its clear. + candidate = adapter.split("async function loadAutoLoadCandidate", 1)[1] + gguf_branch, non_gguf_rest = candidate.split( + 'if (candidate.kind === "gguf") {', 1 + )[1].split("\n } else {\n", 1) + non_gguf_branch = non_gguf_rest.split("if (!(loadResp.is_lora ?? false)) {", 1)[0] + # The cached-GGUF branch keeps the remembered override (it sends it)... + assert "nParallel: config.nParallel ?? null," in gguf_branch + assert "nParallel: null," not in gguf_branch + # ... the safetensors fallback sends no slots, so it clears both. Without + # this the count survives on a model whose run-settings form does not even + # render the field, leaving it unreachable and unclearable from the UI. + assert "nParallel: null," in non_gguf_branch + assert "loadedNParallel: null," in non_gguf_branch + + fresh_default = adapter.split("No downloaded models found. Fetching", 1)[1].split( + 'showAutoLoadSuccess("Loaded Qwen', 1 + )[0] + # The fresh-default download deliberately omits n_parallel from its request, + # so its success state must clear both too; otherwise the control reads as + # an unapplied edit forever against the baseline the status seed fills in. + assert "n_parallel" not in fresh_default.split("saveSpeculativeType", 1)[0] + assert "nParallel: null," in fresh_default + assert "loadedNParallel: null," in fresh_default + + def test_vulkan_inference_devices_are_the_pickable_set(): """GGUF loads run through llama-server, so on a Vulkan build the picker must offer the inference inventory (ggml ordinals, the space `--device Vulkan`