diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 50144893e1..4035188e88 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -127,6 +127,15 @@ LLAMA_SERVER_NOT_FOUND_DETAIL = ( "then try again. (Advanced: set LLAMA_SERVER_PATH to an existing binary.)" ) +# Shared by the route, pre-teardown and post-metadata rejections (#7205). +_VULKAN_DIFFUSION_GPU_IDS_ERROR = ( + "GPU selection (gpu_ids) is not supported for a DiffusionGemma " + "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " + "its device by CUDA physical index, which has no defined mapping " + "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " + "device." +) + # llama-server can serve HTTP 200 while running a model entirely on CPU when a # GPU backend fails to init (#5807 / #5106 / #5830). Classify the startup log so @@ -4710,6 +4719,13 @@ class LlamaCppBackend: probe._read_gguf_metadata(gguf_path) return probe._is_diffusion + def _reject_vulkan_diffusion_gpu_ids_before_teardown( + self, gguf_path: str, model_identifier: str + ) -> None: + """Reject Vulkan + gpu_ids for diffusion GGUFs before Phase 1 teardown.""" + if self._gguf_path_is_diffusion(gguf_path, model_identifier): + raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR) + def _read_gguf_metadata(self, gguf_path: str) -> None: """Read context_length, architecture params, and chat_template from a GGUF header. @@ -6515,12 +6531,7 @@ class LlamaCppBackend: f"present. Available Vulkan devices: {sorted(_pf_probed)}." ) - # A remote uncached GGUF may only reveal that it needs the - # single-device diffusion runner after download. On Vulkan, an - # explicit gpu_ids request cannot be mapped from ggml ordinals to - # that runner's CUDA physical index. Download and classify the main - # file before killing the healthy server so this late rejection is - # non-destructive. The Phase 2 call below reuses this cached path. + # Classify before killing the healthy server (#7205); Phase 2 reuses this path. _preflight_model_path = None if is_vulkan_backend and gpu_ids and hf_repo: _resolved_repo = _resolve_repo_id_casing(hf_repo) @@ -6537,14 +6548,17 @@ class LlamaCppBackend: hf_variant = hf_variant, hf_token = hf_token, ) - if self._gguf_path_is_diffusion(_preflight_model_path, model_identifier): - raise ValueError( - "GPU selection (gpu_ids) is not supported for a DiffusionGemma " - "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " - "its device by CUDA physical index, which has no defined mapping " - "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " - "device." - ) + self._reject_vulkan_diffusion_gpu_ids_before_teardown( + _preflight_model_path, + model_identifier, + ) + elif is_vulkan_backend and gpu_ids and gguf_path and not hf_repo: + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + self._reject_vulkan_diffusion_gpu_ids_before_teardown( + gguf_path, + model_identifier, + ) # ── Phase 1: kill old process (under lock, fast) ────────── with self._lock: @@ -6621,18 +6635,9 @@ class LlamaCppBackend: # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; # serve them with the diffusion runner (same OpenAI-compat interface). if self._is_diffusion: - # The diffusion runner pins its child by CUDA visibility mask, so a - # ggml Vulkan ordinal cannot be honored (wrong GPU / CPU fallback). - # Route and remote-download preflights reject before teardown; keep - # this as a final defense if classification ever disagrees. + # Final defense: route and pre-teardown preflights reject before Phase 1. if is_vulkan_backend and gpu_ids: - raise ValueError( - "GPU selection (gpu_ids) is not supported for a DiffusionGemma " - "GGUF on a Vulkan llama.cpp build: the diffusion runner selects " - "its device by CUDA physical index, which has no defined mapping " - "to ggml Vulkan device ordinals. Omit gpu_ids to use the default " - "device." - ) + raise ValueError(_VULKAN_DIFFUSION_GPU_IDS_ERROR) # Not a tensor/layer GGUF: clear any preserved-fallback flag from a # prior load (this path skips the command builder that clears it). self._layer_preserves_tensor_intent = False diff --git a/studio/backend/tests/test_gpu_memory_mode.py b/studio/backend/tests/test_gpu_memory_mode.py index 271a882b11..fdfcbc1610 100644 --- a/studio/backend/tests/test_gpu_memory_mode.py +++ b/studio/backend/tests/test_gpu_memory_mode.py @@ -22,6 +22,7 @@ and MoE offload itself (``--fit off``). These tests pin: from __future__ import annotations import inspect +import struct import sys import types as _types from pathlib import Path @@ -702,6 +703,15 @@ def test_remote_vulkan_diffusion_preflight_runs_before_teardown(monkeypatch): assert "model_path = _preflight_model_path or self._download_gguf(" in src +def test_local_vulkan_diffusion_preflight_runs_before_teardown(): + src = inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model) + local_preflight = src.index( + "self._reject_vulkan_diffusion_gpu_ids_before_teardown(\n gguf_path," + ) + teardown = src.index("# ── Phase 1: kill old process") + assert local_preflight < teardown + + def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch): backend = LlamaCppBackend() killed = [] @@ -737,6 +747,165 @@ def test_remote_vulkan_diffusion_rejection_keeps_active_server(monkeypatch): assert killed == [] +def test_remote_vulkan_preflight_download_failure_keeps_active_server(monkeypatch, tmp_path): + # A resolvable shard-1 file does not prove the variant is complete, so download + # failures must surface from the pre-teardown _download_gguf, not after the kill. + import hub.utils.gguf as hub_gguf + + cached_shard = tmp_path / "model-00001-of-00003.gguf" + cached_shard.write_bytes(b"GGUF") + monkeypatch.setattr( + hub_gguf, + "resolve_local_gguf_path", + lambda _repo, _variant: str(cached_shard), + ) + + for failure in ( + FileNotFoundError("shard 2 of 3 missing"), + OSError("[Errno 28] No space left on device"), + ConnectionError("hub unreachable"), + ): + backend = LlamaCppBackend() + order = [] + + def _download(_failure = failure, **_kwargs): + order.append("download") + raise _failure + + monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama") + monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True) + monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)]) + monkeypatch.setattr(backend, "_download_gguf", _download) + monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: False) + monkeypatch.setattr(backend, "_kill_process", lambda: order.append("kill")) + monkeypatch.setattr(llama_cpp_module, "_resolve_repo_id_casing", lambda repo: repo) + monkeypatch.setattr( + llama_cpp_module, + "_hf_offline_if_dns_dead", + lambda: __import__("contextlib").nullcontext(), + ) + + with pytest.raises(type(failure)): + backend.load_model( + hf_repo = "owner/model", + hf_variant = "Q4_K_M", + model_identifier = "owner/model", + gpu_ids = [0], + ) + + assert order == ["download"], failure + + +def test_local_vulkan_diffusion_rejection_keeps_active_server(monkeypatch, tmp_path): + gguf_path = tmp_path / "diffusion.gguf" + gguf_path.write_bytes(b"GGUF") + + backend = LlamaCppBackend() + killed = [] + monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama") + monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True) + monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)]) + monkeypatch.setattr(backend, "_gguf_path_is_diffusion", lambda *_args: True) + monkeypatch.setattr(backend, "_kill_process", lambda: killed.append(True)) + + with pytest.raises(ValueError, match = "DiffusionGemma"): + backend.load_model( + gguf_path = str(gguf_path), + model_identifier = "local/diffusion", + gpu_ids = [0], + ) + + assert killed == [] + + +class _ReachedServerStart(Exception): + """Marks a load getting past the pre-teardown preflight.""" + + +def _write_gguf_header( + path: Path, + architecture: str, + *, + diffusion: bool = False, +) -> str: + """Smallest GGUF the header probe can classify: arch, plus the canvas marker.""" + + def _kv_str(key: str, value: str) -> bytes: + kb, vb = key.encode(), value.encode() + return ( + struct.pack(" bytes: + kb = key.encode() + return struct.pack(" LlamaCppBackend: + backend = LlamaCppBackend() + monkeypatch.setattr(backend, "_find_llama_server_binary", lambda **_kwargs: "/bin/llama") + monkeypatch.setattr(backend, "_is_vulkan_backend", lambda _binary = None: True) + monkeypatch.setattr(backend, "_get_gpu_memory", lambda _binary = None: [(0, 1024, 2048)]) + monkeypatch.setattr(backend, "_kill_process", lambda: killed.append(True)) + return backend + + +def test_local_vulkan_pre_teardown_reads_the_real_gguf_header(monkeypatch, tmp_path): + # Classify from the header, not from Vulkan + gpu_ids alone: normal GGUFs load. + killed = [] + backend = _vulkan_pinned_backend(monkeypatch, killed) + monkeypatch.setattr( + backend, + "_wait_for_vram_settle", + lambda **_kwargs: (_ for _ in ()).throw(_ReachedServerStart()), + ) + + with pytest.raises(_ReachedServerStart): + backend.load_model( + gguf_path = _write_gguf_header(tmp_path / "chat.gguf", "llama"), + model_identifier = "local/chat", + gpu_ids = [0], + ) + + assert killed == [True] + + +def test_local_vulkan_diffusion_header_rejects_before_teardown(monkeypatch, tmp_path): + # Same path, real DiffusionGemma canvas marker: rejected with the server intact. + killed = [] + backend = _vulkan_pinned_backend(monkeypatch, killed) + + with pytest.raises(ValueError, match = "DiffusionGemma"): + backend.load_model( + gguf_path = _write_gguf_header(tmp_path / "d.gguf", "gemma3", diffusion = True), + model_identifier = "local/diffusion", + gpu_ids = [0], + ) + + assert killed == [] + + +def test_local_vulkan_missing_gguf_is_reported_before_teardown(monkeypatch, tmp_path): + # The preflight existence check must not cost the live model either. + killed = [] + backend = _vulkan_pinned_backend(monkeypatch, killed) + + with pytest.raises(FileNotFoundError): + backend.load_model( + gguf_path = str(tmp_path / "absent.gguf"), + model_identifier = "local/missing", + gpu_ids = [0], + ) + + assert killed == [] + + def test_start_diffusion_server_resets_tensor_parallel(): # A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model # phase 1 only kills the process, it skips the unload reset). Diffusion is never