Revert RAG auto Vulkan guard, guard multi-backend Vulkan detection, and preserve forced Vulkan across updates

Author: oobabooga

studio/backend/core/inference/llama_cpp.py

@@ -2316,17 +2316,27 @@ def _get_gguf_size_bytes(model_path: str) -> int:
 
     @staticmethod
     def _is_vulkan_backend(binary: Optional[str] = None) -> bool:
-        """True if the installed llama.cpp build is the Vulkan one.
-
-        Builds are single-backend, so the presence of the Vulkan ggml
-        backend library next to llama-server is sufficient. Used to keep
-        the free-memory probe and the GPU pin in the same device-index
-        space (ggml's Vulkan ordinals, not nvidia-smi order).
+        """True if the installed llama.cpp build is Vulkan-only.
+
+        The official prebuilts are single-backend, so the Vulkan ggml backend
+        library next to llama-server identifies a Vulkan build. Used to keep the
+        free-memory probe and the GPU pin in ggml's Vulkan device-index space.
+        Guard the assumption for a custom multi-backend build: if a CUDA or HIP
+        ggml library sits alongside Vulkan, defer to that (the nvidia-smi/torch
+        probes and the CUDA/HIP pin), since those backends are torch-usable and
+        the memory probe/pin are better understood there than on Vulkan.
         """
         binary = binary or LlamaCppBackend._find_llama_server_binary()
         if not binary:
             return False
-        return (_llama_lib_dir(binary) / _vulkan_lib_filename()).is_file()
+        lib_dir = _llama_lib_dir(binary)
+        if not (lib_dir / _vulkan_lib_filename()).is_file():
+            return False
+        for _backend in ("cuda", "hip"):
+            sibling = f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so"
+            if (lib_dir / sibling).is_file():
+                return False
+        return True
 
     @staticmethod
     def _resolve_visible_physical_ids() -> Optional[list[int]]:

studio/backend/core/rag/embeddings.py

@@ -278,15 +278,9 @@ def _resolve_auto() -> str:
     -- or ST if its binary is missing. GPU check is torch-free (nvidia-smi)."""
     from core.inference.llama_cpp import LlamaCppBackend
 
-    binary = LlamaCppBackend._find_llama_server_binary()
-    # A Vulkan build's GPU is not torch-usable (torch has no Vulkan backend), so
-    # its free-memory report must not steer auto to sentence-transformers, which
-    # would then run on CPU/XPU; the llama-server Vulkan build is the fast path.
-    if not LlamaCppBackend._is_vulkan_backend(binary) and LlamaCppBackend._get_gpu_free_memory(
-        binary
-    ):
+    if LlamaCppBackend._get_gpu_free_memory():
         return "sentence-transformers"
-    if binary:
+    if LlamaCppBackend._find_llama_server_binary():
         return "llama-server"
     return "sentence-transformers"
 

studio/backend/tests/test_llama_cpp_update.py

@@ -448,6 +448,48 @@ def _on_start(cmd):
     assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
 
 
+def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
+    # A Vulkan install (marker asset carries 'vulkan') must re-assert
+    # UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to
+    # CUDA/ROCm and silently replaces the Vulkan build.
+    install_dir = tmp_path / "llama.cpp"
+    binary = _write_install(
+        install_dir,
+        "b9493",
+        repo = "ggml-org/llama.cpp",
+        asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz",
+    )
+    monkeypatch.setattr(upd, "_find_binary", lambda: binary)
+    monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
+    monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
+
+    def _on_start(cmd):
+        _write_install(
+            install_dir,
+            "b9518",
+            repo = "ggml-org/llama.cpp",
+            asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
+        )
+
+    popen_kwargs: dict = {}
+    _patch_installer_popen(
+        monkeypatch,
+        lines = ["installed\n"],
+        on_start = _on_start,
+        captured_kwargs = popen_kwargs,
+    )
+
+    assert upd.start_update()["started"] is True
+    deadline = time.time() + 10
+    while time.time() < deadline:
+        job = upd.get_update_status()["job"]
+        if job["state"] in ("success", "error"):
+            break
+        time.sleep(0.05)
+    assert job["state"] == "success", job
+    assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
+
+
 def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
     install_dir = tmp_path / "llama.cpp"
     binary = _write_install(install_dir, "b9595")

studio/backend/tests/test_llama_cpp_vulkan_probe.py

@@ -53,7 +53,11 @@ def _build_loggers_stub():
 _maybe_stub("structlog", lambda: _types.ModuleType("structlog"))
 
 from core.inference import llama_cpp as _llama_mod  # noqa: E402
-from core.inference.llama_cpp import LlamaCppBackend, _vulkan_lib_filename  # noqa: E402
+from core.inference.llama_cpp import (  # noqa: E402
+    LlamaCppBackend,
+    _llama_lib_dir,
+    _vulkan_lib_filename,
+)
 
 MIB = 1024 * 1024
 GIB = 1024 * MIB
@@ -152,5 +156,20 @@ def test_vulkan_pin_args_uses_device_names_not_env_mask():
     assert LlamaCppBackend._vulkan_pin_args([]) == []
 
 
+def test_vulkan_only_build_is_detected(tmp_path):
+    binary = _make_vulkan_install(tmp_path)
+    assert LlamaCppBackend._is_vulkan_backend(binary) is True
+
+
+def test_multi_backend_build_is_not_vulkan_only(tmp_path):
+    # A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be
+    # treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan
+    # device; defer to the CUDA/HIP path instead.
+    binary = _make_vulkan_install(tmp_path)
+    cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so"
+    (_llama_lib_dir(binary) / cuda).write_bytes(b"stub")
+    assert LlamaCppBackend._is_vulkan_backend(binary) is False
+
+
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__, "-v"]))

studio/backend/tests/test_rag_embed_llama_server.py

@@ -52,12 +52,8 @@ def _mock_auto(monkeypatch, *, gpus, binary):
     from core.inference.llama_cpp import LlamaCppBackend
 
     monkeypatch.setattr(config, "EMBED_BACKEND", "auto")
-    # _resolve_auto passes the resolved binary to _get_gpu_free_memory and probes
-    # _is_vulkan_backend (a Vulkan GPU is not torch-usable); accept the arg and
-    # keep these non-Vulkan so the CUDA/no-CUDA selection is what's exercised.
-    monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda *a: gpus))
+    monkeypatch.setattr(LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda: gpus))
     monkeypatch.setattr(LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary))
-    monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda *a: False))
 
 
 def _stub_st_load(monkeypatch):

studio/backend/tests/test_rag_embeddings.py

@@ -220,43 +220,3 @@ def test_st_encode_failure_without_llama_binary_reraises(monkeypatch):
     embeddings._reset_backend()
     with pytest.raises(RuntimeError, match = "CUDA error during encode"):
         embeddings.encode(["alpha", "beta"])
-
-
-def _patch_auto_backend(monkeypatch, *, binary, is_vulkan, gpu_free):
-    """Stub the hardware probes _resolve_auto reads."""
-    from core.inference.llama_cpp import LlamaCppBackend
-
-    monkeypatch.setattr(LlamaCppBackend, "_find_llama_server_binary", staticmethod(lambda: binary))
-    monkeypatch.setattr(
-        LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda b = None: is_vulkan)
-    )
-    monkeypatch.setattr(
-        LlamaCppBackend, "_get_gpu_free_memory", staticmethod(lambda b = None: gpu_free)
-    )
-
-
-def test_resolve_auto_cuda_gpu_picks_sentence_transformers(monkeypatch):
-    # A torch-usable (CUDA/ROCm) GPU -> ST wins bulk indexing.
-    _patch_auto_backend(
-        monkeypatch, binary = "/fake/llama-server", is_vulkan = False, gpu_free = [(0, 24000)]
-    )
-    assert embeddings._resolve_auto() == "sentence-transformers"
-
-
-def test_resolve_auto_vulkan_gpu_picks_llama_server(monkeypatch):
-    # A Vulkan build's GPU is not torch-usable, so its free-memory report must
-    # NOT steer auto to ST (which would run on CPU/XPU); llama-server can use it.
-    _patch_auto_backend(
-        monkeypatch, binary = "/fake/llama-server", is_vulkan = True, gpu_free = [(0, 8000)]
-    )
-    assert embeddings._resolve_auto() == "llama-server"
-
-
-def test_resolve_auto_cpu_host_with_binary_picks_llama_server(monkeypatch):
-    _patch_auto_backend(monkeypatch, binary = "/fake/llama-server", is_vulkan = False, gpu_free = [])
-    assert embeddings._resolve_auto() == "llama-server"
-
-
-def test_resolve_auto_cpu_host_without_binary_picks_sentence_transformers(monkeypatch):
-    _patch_auto_backend(monkeypatch, binary = None, is_vulkan = False, gpu_free = [])
-    assert embeddings._resolve_auto() == "sentence-transformers"

studio/backend/utils/llama_cpp_update.py

@@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
         logger.info("llama update: installing", cmd = " ".join(cmd))
         # Stream progress lines into job["progress"].
         env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
+        # Preserve a Vulkan install across updates. The marker asset records the
+        # backend, but detect_host on a CUDA/ROCm box would re-route to that
+        # backend and silently replace the Vulkan build; re-assert the Vulkan
+        # choice via the same env flag setup uses (mirrors _rocm_install_args).
+        if asset and "vulkan" in asset.lower():
+            env["UNSLOTH_FORCE_VULKAN"] = "1"
         proc = subprocess.Popen(
             cmd,
             stdout = subprocess.PIPE,