Fix the libaray path for probe_server_capabilities() (#5797)
* Fix the libaray path for probe_server_capabilities() even when running something as simple as `./llama-server --help`, the binary still requires correct LD_LIBRARY_PATH to work - or it returns merely an "error while loading shared libraries": "libllama-server-impl.so: cannot open shared object file: No such file or directory" For a local installation with no LD_LIBRARY_PATH specifically set, the probe_server_capabilities() run of `./llana-server --help` should share the same libaray resolution logic as start_llama_server(). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Adjust and readd comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: imagineer99 <samleejackson0@gmail.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
This commit is contained in:
Hua
GitHub
parent
9beefcaf7c
commit
9413e72802
2 changed files with 138 additions and 100 deletions
|
|
@ -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
|
||||
# <binary_dir>/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 <binary_dir>/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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue