diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 00dfe14e63..74e7a8f220 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1281,12 +1281,14 @@ class LlamaCppBackend: supports_ctx_checkpoints = False supports_no_cache_prompt = False try: + probe_env = cls._llama_server_env_for_binary(bin_path) result = subprocess.run( [bin_path, "--help"], capture_output = True, text = True, timeout = 10, check = False, + env = probe_env, ) help_text = (result.stdout or "") + "\n" + (result.stderr or "") # Split into per-flag blocks (each --flag line + its indented @@ -1795,6 +1797,100 @@ class LlamaCppBackend: path_dirs.append(cuda_bin_x64) return path_dirs + @staticmethod + def _llama_server_env_for_binary(binary: str) -> dict[str, str]: + """Build a subprocess env that lets llama-server resolve native libs.""" + env = child_env_without_native_path_secret() + binary_dir = str(Path(binary).parent) + + if sys.platform == "win32": + # Ordering: see _build_windows_path_dirs. #5106. + path_dirs = LlamaCppBackend._build_windows_path_dirs( + binary_dir, + sys.prefix, + os.environ.get("CUDA_PATH", ""), + ) + existing_path = env.get("PATH", "") + env["PATH"] = ";".join(path_dirs) + ";" + existing_path + + # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile + # kernel files (rocblas/library/*.dat + *.hsaco); the DLL searches + # /rocblas/library/ which doesn't exist. + _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", "")) + if _hip_path: + _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library") + if os.path.isdir(_rocblas_lib): + env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib) + else: + # Linux: LD_LIBRARY_PATH for shared libs next to the binary plus + # CUDA runtime libs (libcudart, libcublas, etc.) + import platform + + lib_dirs = [] + # WSL: system HIP before the bundle's (which segfaults on /dev/dxg). + for _wsl_rocm in _wsl_system_rocm_lib_dirs(): + lib_dirs.append(_wsl_rocm) + if lib_dirs: + env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") + lib_dirs.append(binary_dir) + _arch = platform.machine() # x86_64, aarch64, etc. + + # Pip-installed nvidia CUDA runtime libs. The prebuilt binary links + # libcudart.so.13 / libcublas.so.13 which live here, not in + # /usr/local/cuda. + import glob as _glob + + for _nv_pattern in [ + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cu*", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "cudnn", + "lib", + ), + os.path.join( + sys.prefix, + "lib", + "python*", + "site-packages", + "nvidia", + "nvjitlink", + "lib", + ), + ]: + for _nv_dir in _glob.glob(_nv_pattern): + if os.path.isdir(_nv_dir): + lib_dirs.append(_nv_dir) + + for cuda_lib in [ + "/usr/local/cuda/lib64", + f"/usr/local/cuda/targets/{_arch}-linux/lib", + # Fallback CUDA compat paths (e.g. binary built with CUDA 12 + # where default /usr/local/cuda is CUDA 13+). + "/usr/local/cuda-12/lib64", + "/usr/local/cuda-12.8/lib64", + f"/usr/local/cuda-12/targets/{_arch}-linux/lib", + f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", + ]: + if os.path.isdir(cuda_lib): + lib_dirs.append(cuda_lib) + existing_ld = env.get("LD_LIBRARY_PATH", "") + new_ld = ":".join(lib_dirs) + env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld + + return env + @staticmethod def _select_gpus( model_size_bytes: int, @@ -4168,8 +4264,15 @@ class LlamaCppBackend: logger.info(f"Starting llama-server: {' '.join(_log_cmd)}") # Library paths so llama-server finds its shared libs and CUDA DLLs. - env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) + env = self._llama_server_env_for_binary(binary) + + # Windows + full offload: PASSIVE OMP + 2 threads stop + # spin-wait burning CPU. CPU/partial offload keeps default + # OMP parallelism. #5692. + if sys.platform == "win32" and full_offload_tuning_active: + env.setdefault("OMP_WAIT_POLICY", "PASSIVE") + if not threads_overridden: + env.setdefault("OMP_NUM_THREADS", "2") # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use # shared system RAM. setdefault so a user value wins. @@ -4185,104 +4288,6 @@ class LlamaCppBackend: f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) - if sys.platform == "win32": - # Ordering: see _build_windows_path_dirs. #5106. - path_dirs = self._build_windows_path_dirs( - binary_dir, - sys.prefix, - os.environ.get("CUDA_PATH", ""), - ) - existing_path = env.get("PATH", "") - env["PATH"] = ";".join(path_dirs) + ";" + existing_path - - # Windows + full offload: PASSIVE OMP + 2 threads stop - # spin-wait burning CPU. CPU/partial offload keeps - # default OMP parallelism. #5692. - if full_offload_tuning_active: - env.setdefault("OMP_WAIT_POLICY", "PASSIVE") - if not threads_overridden: - env.setdefault("OMP_NUM_THREADS", "2") - - # ROCm: the prebuilt bundles rocblas.dll but NOT the Tensile - # kernel files (rocblas/library/*.dat + *.hsaco); the DLL - # searches /rocblas/library/ which doesn't exist - # -> silent crash on the first GEMM. ROCBLAS_TENSILE_LIBPATH - # repoints that search at the ROCm install. - _hip_path = os.environ.get("HIP_PATH", os.environ.get("ROCM_PATH", "")) - if _hip_path: - _rocblas_lib = os.path.join(_hip_path, "bin", "rocblas", "library") - if os.path.isdir(_rocblas_lib): - env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib) - else: - # Linux: LD_LIBRARY_PATH for shared libs next to the binary - # plus CUDA runtime libs (libcudart, libcublas, etc.) - import platform - - lib_dirs = [] - # WSL: system HIP before the bundle's (which segfaults on - # /dev/dxg). Mirror install_llama_prebuilt.binary_env, which - # validates the prebuilt with this same ordering. - for _wsl_rocm in _wsl_system_rocm_lib_dirs(): - lib_dirs.append(_wsl_rocm) - if lib_dirs: - env.setdefault("HSA_ENABLE_DXG_DETECTION", "1") - lib_dirs.append(binary_dir) - _arch = platform.machine() # x86_64, aarch64, etc. - - # Pip-installed nvidia CUDA runtime libs. The prebuilt - # binary links libcudart.so.13 / libcublas.so.13 which live - # here, not in /usr/local/cuda. - import glob as _glob - - for _nv_pattern in [ - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cu*", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "cudnn", - "lib", - ), - os.path.join( - sys.prefix, - "lib", - "python*", - "site-packages", - "nvidia", - "nvjitlink", - "lib", - ), - ]: - for _nv_dir in _glob.glob(_nv_pattern): - if os.path.isdir(_nv_dir): - lib_dirs.append(_nv_dir) - - for cuda_lib in [ - "/usr/local/cuda/lib64", - f"/usr/local/cuda/targets/{_arch}-linux/lib", - # Fallback CUDA compat paths (e.g. binary built with - # CUDA 12 where default /usr/local/cuda is CUDA 13+). - "/usr/local/cuda-12/lib64", - "/usr/local/cuda-12.8/lib64", - f"/usr/local/cuda-12/targets/{_arch}-linux/lib", - f"/usr/local/cuda-12.8/targets/{_arch}-linux/lib", - ]: - if os.path.isdir(cuda_lib): - lib_dirs.append(cuda_lib) - existing_ld = env.get("LD_LIBRARY_PATH", "") - new_ld = ":".join(lib_dirs) - env["LD_LIBRARY_PATH"] = f"{new_ld}:{existing_ld}" if existing_ld else new_ld - # Pin to selected GPU(s). On ROCm, narrowing only # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full # set, so set HIP_VISIBLE_DEVICES too. diff --git a/studio/backend/tests/test_llama_cpp_mtp_detection.py b/studio/backend/tests/test_llama_cpp_mtp_detection.py index 2cc0fb4184..14c3848d33 100644 --- a/studio/backend/tests/test_llama_cpp_mtp_detection.py +++ b/studio/backend/tests/test_llama_cpp_mtp_detection.py @@ -10,6 +10,7 @@ the _already_in_target_state mirror that prevents needless reloads. from __future__ import annotations import inspect +import os import struct import sys import types as _types @@ -566,6 +567,38 @@ def test_probe_server_capabilities_detects_draft_mtp(tmp_path): assert caps["supports_mtp"] is True +@_NEEDS_BASH +def test_probe_server_capabilities_uses_binary_library_env(tmp_path, monkeypatch): + fake = _make_fake_llama_server( + tmp_path / "llama-server", + "--spec-type none,mtp,ngram-simple\n", + ) + captured = {} + + monkeypatch.setattr( + "core.inference.llama_cpp.child_env_without_native_path_secret", + lambda: {"LD_LIBRARY_PATH": "/already-there"}, + ) + + def fake_run(cmd, **kwargs): + captured["cmd"] = cmd + captured["env"] = kwargs.get("env") + return _types.SimpleNamespace(stdout = "--spec-type none,mtp,ngram-simple\n", stderr = "") + + monkeypatch.setattr("core.inference.llama_cpp.subprocess.run", fake_run) + + _clear_caps_cache() + caps = LlamaCppBackend.probe_server_capabilities(str(fake)) + + assert caps["found"] is True + assert caps["supports_mtp"] is True + assert captured["cmd"] == [str(fake), "--help"] + assert captured["env"] is not None + ld_dirs = captured["env"]["LD_LIBRARY_PATH"].split(os.pathsep) + assert str(fake.parent) in ld_dirs + assert "/already-there" in ld_dirs + + @_NEEDS_BASH def test_probe_server_capabilities_detects_renamed_mtp(tmp_path): # Renamed upstream: draft-mtp -> mtp.