This commit is contained in:
@@ -0,0 +1,87 @@
|
||||
"""Shared gRPC bearer token authentication interceptor for LocalAI Python backends.
|
||||
|
||||
When the environment variable LOCALAI_GRPC_AUTH_TOKEN is set, requests without
|
||||
a valid Bearer token in the 'authorization' metadata header are rejected with
|
||||
UNAUTHENTICATED. When the variable is empty or unset, no authentication is
|
||||
performed (backward compatible).
|
||||
"""
|
||||
|
||||
import hmac
|
||||
import os
|
||||
|
||||
import grpc
|
||||
|
||||
from parent_watch import start_parent_death_watcher
|
||||
|
||||
|
||||
class _AbortHandler(grpc.RpcMethodHandler):
|
||||
"""A method handler that immediately aborts with UNAUTHENTICATED."""
|
||||
|
||||
def __init__(self):
|
||||
self.request_streaming = False
|
||||
self.response_streaming = False
|
||||
self.request_deserializer = None
|
||||
self.response_serializer = None
|
||||
self.unary_unary = self._abort
|
||||
self.unary_stream = None
|
||||
self.stream_unary = None
|
||||
self.stream_stream = None
|
||||
|
||||
@staticmethod
|
||||
def _abort(request, context):
|
||||
context.abort(grpc.StatusCode.UNAUTHENTICATED, "invalid token")
|
||||
|
||||
|
||||
class TokenAuthInterceptor(grpc.ServerInterceptor):
|
||||
"""Sync gRPC server interceptor that validates a bearer token."""
|
||||
|
||||
def __init__(self, token: str):
|
||||
self._token = token
|
||||
self._abort_handler = _AbortHandler()
|
||||
|
||||
def intercept_service(self, continuation, handler_call_details):
|
||||
metadata = dict(handler_call_details.invocation_metadata)
|
||||
auth = metadata.get("authorization", "")
|
||||
expected = "Bearer " + self._token
|
||||
if not hmac.compare_digest(auth, expected):
|
||||
return self._abort_handler
|
||||
return continuation(handler_call_details)
|
||||
|
||||
|
||||
class AsyncTokenAuthInterceptor(grpc.aio.ServerInterceptor):
|
||||
"""Async gRPC server interceptor that validates a bearer token."""
|
||||
|
||||
def __init__(self, token: str):
|
||||
self._token = token
|
||||
|
||||
async def intercept_service(self, continuation, handler_call_details):
|
||||
metadata = dict(handler_call_details.invocation_metadata)
|
||||
auth = metadata.get("authorization", "")
|
||||
expected = "Bearer " + self._token
|
||||
if not hmac.compare_digest(auth, expected):
|
||||
return _AbortHandler()
|
||||
return await continuation(handler_call_details)
|
||||
|
||||
|
||||
def get_auth_interceptors(*, aio: bool = False):
|
||||
"""Return a list of gRPC interceptors for bearer token auth.
|
||||
|
||||
Args:
|
||||
aio: If True, return async-compatible interceptors for grpc.aio.server().
|
||||
If False (default), return sync interceptors for grpc.server().
|
||||
|
||||
Returns an empty list when LOCALAI_GRPC_AUTH_TOKEN is not set.
|
||||
"""
|
||||
# Arm the best-effort parent-death backstop here: this is the single helper
|
||||
# every LocalAI Python backend invokes exactly once while building its gRPC
|
||||
# server (mirroring how the Go watcher arms in pkg/grpc's shared serve path).
|
||||
# start_parent_death_watcher() is idempotent and a no-op when disabled or on
|
||||
# unsupported platforms — see parent_watch.py.
|
||||
start_parent_death_watcher()
|
||||
|
||||
token = os.environ.get("LOCALAI_GRPC_AUTH_TOKEN", "")
|
||||
if not token:
|
||||
return []
|
||||
if aio:
|
||||
return [AsyncTokenAuthInterceptor(token)]
|
||||
return [TokenAuthInterceptor(token)]
|
||||
@@ -0,0 +1,578 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
#
|
||||
# use the library by adding the following line to a script:
|
||||
# source $(dirname $0)/../common/libbackend.sh
|
||||
#
|
||||
# If you want to limit what targets a backend can be used on, set the variable LIMIT_TARGETS to a
|
||||
# space separated list of valid targets BEFORE sourcing the library, for example to only allow a backend
|
||||
# to be used on CUDA and CPU backends:
|
||||
#
|
||||
# LIMIT_TARGETS="cublas cpu"
|
||||
# source $(dirname $0)/../common/libbackend.sh
|
||||
#
|
||||
# You can use any valid BUILD_TYPE or BUILD_PROFILE, if you need to limit a backend to CUDA 12 only:
|
||||
#
|
||||
# LIMIT_TARGETS="cublas12"
|
||||
# source $(dirname $0)/../common/libbackend.sh
|
||||
#
|
||||
# You can switch between uv (conda-like) and pip installation methods by setting USE_PIP:
|
||||
# USE_PIP=true source $(dirname $0)/../common/libbackend.sh
|
||||
#
|
||||
# ===================== user-configurable defaults =====================
|
||||
PYTHON_VERSION="${PYTHON_VERSION:-3.10}" # e.g. 3.10 / 3.11 / 3.12 / 3.13
|
||||
PYTHON_PATCH="${PYTHON_PATCH:-18}" # e.g. 18 -> 3.10.18 ; 13 -> 3.11.13
|
||||
PY_STANDALONE_TAG="${PY_STANDALONE_TAG:-20250818}" # release tag date
|
||||
# Enable/disable bundling of a portable Python build
|
||||
PORTABLE_PYTHON="${PORTABLE_PYTHON:-false}"
|
||||
|
||||
# If you want to fully pin the filename (including tuned CPU targets), set:
|
||||
# PORTABLE_PY_FILENAME="cpython-3.10.18+20250818-x86_64_v3-unknown-linux-gnu-install_only.tar.gz"
|
||||
: "${PORTABLE_PY_FILENAME:=}"
|
||||
: "${PORTABLE_PY_SHA256:=}" # optional; if set we verify the download
|
||||
# =====================================================================
|
||||
|
||||
# Default to uv if USE_PIP is not set
|
||||
if [ "x${USE_PIP:-}" == "x" ]; then
|
||||
USE_PIP=false
|
||||
fi
|
||||
|
||||
# ----------------------- helpers -----------------------
|
||||
function _is_musl() {
|
||||
# detect musl (Alpine, etc)
|
||||
if command -v ldd >/dev/null 2>&1; then
|
||||
ldd --version 2>&1 | grep -qi musl && return 0
|
||||
fi
|
||||
# busybox-ish fallback
|
||||
if command -v getconf >/dev/null 2>&1; then
|
||||
getconf GNU_LIBC_VERSION >/dev/null 2>&1 || return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
function _triple() {
|
||||
local os="" arch="" libc="gnu"
|
||||
case "$(uname -s)" in
|
||||
Linux*) os="unknown-linux" ;;
|
||||
Darwin*) os="apple-darwin" ;;
|
||||
MINGW*|MSYS*|CYGWIN*) os="pc-windows-msvc" ;; # best-effort for Git Bash
|
||||
*) echo "Unsupported OS $(uname -s)"; exit 1;;
|
||||
esac
|
||||
|
||||
case "$(uname -m)" in
|
||||
x86_64) arch="x86_64" ;;
|
||||
aarch64|arm64) arch="aarch64" ;;
|
||||
armv7l) arch="armv7" ;;
|
||||
i686|i386) arch="i686" ;;
|
||||
ppc64le) arch="ppc64le" ;;
|
||||
s390x) arch="s390x" ;;
|
||||
riscv64) arch="riscv64" ;;
|
||||
*) echo "Unsupported arch $(uname -m)"; exit 1;;
|
||||
esac
|
||||
|
||||
if [[ "$os" == "unknown-linux" ]]; then
|
||||
if _is_musl; then
|
||||
libc="musl"
|
||||
else
|
||||
libc="gnu"
|
||||
fi
|
||||
echo "${arch}-${os}-${libc}"
|
||||
else
|
||||
echo "${arch}-${os}"
|
||||
fi
|
||||
}
|
||||
|
||||
function _portable_dir() {
|
||||
echo "${EDIR}/python"
|
||||
}
|
||||
|
||||
function _portable_bin() {
|
||||
# python-build-standalone puts python in ./bin
|
||||
echo "$(_portable_dir)/bin"
|
||||
}
|
||||
|
||||
function _portable_python() {
|
||||
if [ -x "$(_portable_bin)/python3" ]; then
|
||||
echo "$(_portable_bin)/python3"
|
||||
else
|
||||
echo "$(_portable_bin)/python"
|
||||
fi
|
||||
}
|
||||
|
||||
|
||||
# macOS loader env for the portable CPython
|
||||
_macosPortableEnv() {
|
||||
if [ "$(uname -s)" = "Darwin" ]; then
|
||||
export DYLD_LIBRARY_PATH="$(_portable_dir)/lib${DYLD_LIBRARY_PATH:+:${DYLD_LIBRARY_PATH}}"
|
||||
export DYLD_FALLBACK_LIBRARY_PATH="$(_portable_dir)/lib${DYLD_FALLBACK_LIBRARY_PATH:+:${DYLD_FALLBACK_LIBRARY_PATH}}"
|
||||
fi
|
||||
}
|
||||
|
||||
# Good hygiene on macOS for downloaded/extracted trees
|
||||
_unquarantinePortablePython() {
|
||||
if [ "$(uname -s)" = "Darwin" ]; then
|
||||
command -v xattr >/dev/null 2>&1 && xattr -dr com.apple.quarantine "$(_portable_dir)" || true
|
||||
fi
|
||||
}
|
||||
|
||||
# ------------------ ### PORTABLE PYTHON ------------------
|
||||
function ensurePortablePython() {
|
||||
local pdir="$(_portable_dir)"
|
||||
local pbin="$(_portable_bin)"
|
||||
local pyexe
|
||||
|
||||
if [ -x "${pbin}/python3" ] || [ -x "${pbin}/python" ]; then
|
||||
_macosPortableEnv
|
||||
return 0
|
||||
fi
|
||||
|
||||
mkdir -p "${pdir}"
|
||||
local triple="$(_triple)"
|
||||
|
||||
local full_ver="${PYTHON_VERSION}.${PYTHON_PATCH}"
|
||||
local fn=""
|
||||
if [ -n "${PORTABLE_PY_FILENAME}" ]; then
|
||||
fn="${PORTABLE_PY_FILENAME}"
|
||||
else
|
||||
# generic asset name: cpython-<full_ver>+<tag>-<triple>-install_only.tar.gz
|
||||
fn="cpython-${full_ver}+${PY_STANDALONE_TAG}-${triple}-install_only.tar.gz"
|
||||
fi
|
||||
|
||||
local url="https://github.com/astral-sh/python-build-standalone/releases/download/${PY_STANDALONE_TAG}/${fn}"
|
||||
local tmp="${pdir}/${fn}"
|
||||
echo "Downloading portable Python: ${fn}"
|
||||
# curl with retries; fall back to wget if needed
|
||||
if command -v curl >/dev/null 2>&1; then
|
||||
curl -L --fail --retry 3 --retry-delay 1 -o "${tmp}" "${url}"
|
||||
else
|
||||
wget -O "${tmp}" "${url}"
|
||||
fi
|
||||
|
||||
if [ -n "${PORTABLE_PY_SHA256}" ]; then
|
||||
echo "${PORTABLE_PY_SHA256} ${tmp}" | sha256sum -c -
|
||||
fi
|
||||
|
||||
echo "Extracting ${fn} -> ${pdir}"
|
||||
# always a .tar.gz (we purposely choose install_only)
|
||||
tar -xzf "${tmp}" -C "${pdir}"
|
||||
rm -f "${tmp}"
|
||||
|
||||
# Some archives nest a directory; if so, flatten to ${pdir}
|
||||
# Find the first dir with a 'bin/python*'
|
||||
local inner
|
||||
inner="$(find "${pdir}" -type f -path "*/bin/python*" -maxdepth 3 2>/dev/null | head -n1 || true)"
|
||||
if [ -n "${inner}" ]; then
|
||||
local inner_root
|
||||
inner_root="$(dirname "$(dirname "${inner}")")" # .../bin -> root
|
||||
if [ "${inner_root}" != "${pdir}" ]; then
|
||||
# move contents up one level
|
||||
shopt -s dotglob
|
||||
mv "${inner_root}/"* "${pdir}/"
|
||||
rm -rf "${inner_root}"
|
||||
shopt -u dotglob
|
||||
fi
|
||||
fi
|
||||
|
||||
_unquarantinePortablePython
|
||||
_macosPortableEnv
|
||||
# Make sure it's runnable
|
||||
pyexe="$(_portable_python)"
|
||||
"${pyexe}" -V
|
||||
}
|
||||
|
||||
# init handles the setup of the library
|
||||
function init() {
|
||||
BACKEND_NAME=${PWD##*/}
|
||||
MY_DIR=$(realpath "$(dirname "$0")")
|
||||
BUILD_PROFILE=$(getBuildProfile)
|
||||
|
||||
EDIR=${MY_DIR}
|
||||
if [ "x${ENV_DIR:-}" != "x" ]; then
|
||||
EDIR=${ENV_DIR}
|
||||
fi
|
||||
|
||||
if [ ! -z "${LIMIT_TARGETS:-}" ]; then
|
||||
isValidTarget=$(checkTargets ${LIMIT_TARGETS})
|
||||
if [ ${isValidTarget} != true ]; then
|
||||
echo "${BACKEND_NAME} can only be used on the following targets: ${LIMIT_TARGETS}"
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
echo "Initializing libbackend for ${BACKEND_NAME}"
|
||||
}
|
||||
|
||||
|
||||
# getBuildProfile will inspect the system to determine which build profile is appropriate:
|
||||
# returns one of the following:
|
||||
# - cublas12
|
||||
# - cublas13
|
||||
# - hipblas
|
||||
# - intel
|
||||
function getBuildProfile() {
|
||||
if [ x"${BUILD_TYPE:-}" == "xcublas" ] || [ x"${BUILD_TYPE:-}" == "xl4t" ]; then
|
||||
if [ ! -z "${CUDA_MAJOR_VERSION:-}" ]; then
|
||||
echo ${BUILD_TYPE}${CUDA_MAJOR_VERSION}
|
||||
else
|
||||
echo ${BUILD_TYPE}
|
||||
fi
|
||||
return 0
|
||||
fi
|
||||
|
||||
if [ -d "/opt/intel" ]; then
|
||||
echo "intel"
|
||||
return 0
|
||||
fi
|
||||
|
||||
if [ -n "${BUILD_TYPE:-}" ]; then
|
||||
echo ${BUILD_TYPE}
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "cpu"
|
||||
}
|
||||
|
||||
|
||||
# Make the venv relocatable:
|
||||
# - rewrite venv/bin/python{,3} to relative symlinks into $(_portable_dir)
|
||||
# - normalize entrypoint shebangs to /usr/bin/env python3
|
||||
# - optionally update pyvenv.cfg to point to the portable Python directory (only at runtime)
|
||||
# Usage: _makeVenvPortable [--update-pyvenv-cfg]
|
||||
_makeVenvPortable() {
|
||||
local update_pyvenv_cfg=false
|
||||
if [ "${1:-}" = "--update-pyvenv-cfg" ]; then
|
||||
update_pyvenv_cfg=true
|
||||
fi
|
||||
|
||||
local venv_dir="${EDIR}/venv"
|
||||
local vbin="${venv_dir}/bin"
|
||||
|
||||
[ -d "${vbin}" ] || return 0
|
||||
|
||||
# 1) Replace python symlinks with relative ones to ../../python/bin/python3
|
||||
# (venv/bin -> venv -> EDIR -> python/bin)
|
||||
local rel_py='../../python/bin/python3'
|
||||
|
||||
for name in python3 python; do
|
||||
if [ -e "${vbin}/${name}" ] || [ -L "${vbin}/${name}" ]; then
|
||||
rm -f "${vbin}/${name}"
|
||||
fi
|
||||
done
|
||||
ln -s "${rel_py}" "${vbin}/python3"
|
||||
ln -s "python3" "${vbin}/python"
|
||||
|
||||
# 2) Update pyvenv.cfg to point to the portable Python directory (only at runtime)
|
||||
# Use absolute path resolved at runtime so it works when the venv is copied
|
||||
if [ "$update_pyvenv_cfg" = "true" ]; then
|
||||
local pyvenv_cfg="${venv_dir}/pyvenv.cfg"
|
||||
if [ -f "${pyvenv_cfg}" ]; then
|
||||
local portable_dir="$(_portable_dir)"
|
||||
# Resolve to absolute path - this ensures it works when the backend is copied
|
||||
# Only resolve if the directory exists (it should if ensurePortablePython was called)
|
||||
if [ -d "${portable_dir}" ]; then
|
||||
portable_dir="$(cd "${portable_dir}" && pwd)"
|
||||
else
|
||||
# Fallback to relative path if directory doesn't exist yet
|
||||
portable_dir="../python"
|
||||
fi
|
||||
local sed_i=(sed -i)
|
||||
# macOS/BSD sed needs a backup suffix; GNU sed doesn't. Make it portable:
|
||||
if sed --version >/dev/null 2>&1; then
|
||||
sed_i=(sed -i)
|
||||
else
|
||||
sed_i=(sed -i '')
|
||||
fi
|
||||
# Update the home field in pyvenv.cfg
|
||||
# Handle both absolute paths (starting with /) and relative paths
|
||||
if grep -q "^home = " "${pyvenv_cfg}"; then
|
||||
"${sed_i[@]}" "s|^home = .*|home = ${portable_dir}|" "${pyvenv_cfg}"
|
||||
else
|
||||
# If home field doesn't exist, add it
|
||||
echo "home = ${portable_dir}" >> "${pyvenv_cfg}"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
# 3) Rewrite shebangs of entry points to use env, so the venv is relocatable
|
||||
# Only touch text files that start with #! and reference the current venv.
|
||||
local ve_abs="${vbin}/python"
|
||||
local sed_i=(sed -i)
|
||||
# macOS/BSD sed needs a backup suffix; GNU sed doesn't. Make it portable:
|
||||
if sed --version >/dev/null 2>&1; then
|
||||
sed_i=(sed -i)
|
||||
else
|
||||
sed_i=(sed -i '')
|
||||
fi
|
||||
|
||||
for f in "${vbin}"/*; do
|
||||
[ -f "$f" ] || continue
|
||||
# Fast path: check first two bytes (#!)
|
||||
head -c2 "$f" 2>/dev/null | grep -q '^#!' || continue
|
||||
# Only rewrite if the shebang mentions the (absolute) venv python
|
||||
if head -n1 "$f" | grep -Fq "${ve_abs}"; then
|
||||
"${sed_i[@]}" '1s|^#!.*$|#!/usr/bin/env python3|' "$f"
|
||||
chmod +x "$f" 2>/dev/null || true
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
|
||||
# Apply the venv to the current process: VIRTUAL_ENV, PATH, PYTHONHOME hygiene.
|
||||
# Equivalent to the runtime portion of `source bin/activate`, but computed from
|
||||
# $EDIR (resolved at runtime via realpath) instead of the path baked into
|
||||
# bin/activate at venv-create time. `uv venv` (and `python -m venv`) both bake
|
||||
# the create-time absolute path in, so sourcing activate on a relocated venv —
|
||||
# e.g. one built at /vllm/venv inside a Docker stage and unpacked under
|
||||
# /backends/cuda13-vllm-development/venv at runtime — silently prepends a
|
||||
# stale, non-existent path to $PATH. Doing the setup ourselves sidesteps that;
|
||||
# this is the same approach `uv run` takes internally.
|
||||
_activateVenv() {
|
||||
export VIRTUAL_ENV="${EDIR}/venv"
|
||||
export PATH="${EDIR}/venv/bin:${PATH}"
|
||||
unset PYTHONHOME
|
||||
}
|
||||
|
||||
# ensureVenv makes sure that the venv for the backend both exists, and is activated.
|
||||
#
|
||||
# This function is idempotent, so you can call it as many times as you want and it will
|
||||
# always result in an activated virtual environment
|
||||
function ensureVenv() {
|
||||
local interpreter=""
|
||||
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ] || [ -e "$(_portable_python)" ]; then
|
||||
echo "Using portable Python"
|
||||
ensurePortablePython
|
||||
interpreter="$(_portable_python)"
|
||||
else
|
||||
# Prefer system python${PYTHON_VERSION}, else python3, else fall back to bundled
|
||||
if command -v python${PYTHON_VERSION} >/dev/null 2>&1; then
|
||||
interpreter="python${PYTHON_VERSION}"
|
||||
elif command -v python3 >/dev/null 2>&1; then
|
||||
interpreter="python3"
|
||||
else
|
||||
echo "No suitable system Python found, bootstrapping portable build..."
|
||||
ensurePortablePython
|
||||
interpreter="$(_portable_python)"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ ! -d "${EDIR}/venv" ]; then
|
||||
if [ "x${USE_PIP}" == "xtrue" ]; then
|
||||
# --copies is only needed when we will later relocate the venv via
|
||||
# _makeVenvPortable (PORTABLE_PYTHON=true). Some Python builds —
|
||||
# notably macOS system Python — refuse to create a venv with
|
||||
# --copies because the build doesn't support it. Fall back to
|
||||
# symlinks in that case.
|
||||
local venv_args=""
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
|
||||
venv_args="--copies"
|
||||
fi
|
||||
"${interpreter}" -m venv ${venv_args} "${EDIR}/venv"
|
||||
_activateVenv
|
||||
"${interpreter}" -m pip install --upgrade pip
|
||||
else
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
|
||||
uv venv --python "${interpreter}" "${EDIR}/venv"
|
||||
else
|
||||
uv venv --python "${PYTHON_VERSION}" "${EDIR}/venv"
|
||||
fi
|
||||
fi
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
|
||||
# During install, only update symlinks and shebangs, not pyvenv.cfg
|
||||
_makeVenvPortable
|
||||
fi
|
||||
fi
|
||||
|
||||
# We call it here to make sure that when we source a venv we can still use python as expected
|
||||
if [ -x "$(_portable_python)" ]; then
|
||||
_macosPortableEnv
|
||||
fi
|
||||
|
||||
if [ "x${VIRTUAL_ENV:-}" != "x${EDIR}/venv" ]; then
|
||||
_activateVenv
|
||||
fi
|
||||
}
|
||||
|
||||
|
||||
function runProtogen() {
|
||||
ensureVenv
|
||||
|
||||
# Match grpcio-tools to the grpcio already installed by the backend's
|
||||
# requirements. grpcio and grpcio-tools are released in lockstep, and the
|
||||
# protoc that grpcio-tools bundles stamps a Protobuf "gencode" version into
|
||||
# backend_pb2.py. Left unpinned, `uv pip install grpcio-tools` pulls the
|
||||
# newest release, whose newer gencode (e.g. 7.35.0) trips Protobuf's
|
||||
# runtime >= gencode guarantee at import time when a backend caps the
|
||||
# protobuf runtime lower (vLLM pins it to 6.33.6), crashing the backend with
|
||||
# "grpc service not ready" before it ever loads a model. Pinning
|
||||
# grpcio-tools to the installed grpcio version keeps the gencode in step with
|
||||
# the runtime. Falls back to unpinned when grpcio isn't installed yet.
|
||||
# See mudler/LocalAI#10718.
|
||||
local grpcio_tools_spec="grpcio-tools"
|
||||
local grpcio_version
|
||||
grpcio_version="$(python -c 'import importlib.metadata as m; print(m.version("grpcio"))' 2>/dev/null || true)"
|
||||
if [ -n "${grpcio_version}" ]; then
|
||||
grpcio_tools_spec="grpcio-tools==${grpcio_version}"
|
||||
fi
|
||||
|
||||
if [ "x${USE_PIP}" == "xtrue" ]; then
|
||||
pip install "${grpcio_tools_spec}"
|
||||
else
|
||||
uv pip install "${grpcio_tools_spec}"
|
||||
fi
|
||||
pushd "${EDIR}" >/dev/null
|
||||
# use the venv python (ensures correct interpreter & sys.path)
|
||||
python -m grpc_tools.protoc -I../../ -I./ --python_out=. --grpc_python_out=. backend.proto
|
||||
popd >/dev/null
|
||||
}
|
||||
|
||||
|
||||
# installRequirements looks for several requirements files and if they exist runs the install for them in order
|
||||
#
|
||||
# - requirements-install.txt
|
||||
# - requirements.txt
|
||||
# - requirements-${BUILD_TYPE}.txt
|
||||
# - requirements-${BUILD_PROFILE}.txt
|
||||
#
|
||||
# BUILD_PROFILE is a more specific version of BUILD_TYPE, ex: cuda-12 or cuda-13
|
||||
# it can also include some options that we do not have BUILD_TYPES for, ex: intel
|
||||
#
|
||||
# NOTE: for BUILD_PROFILE==intel, this function does NOT automatically use the Intel python package index.
|
||||
# you may want to add the following line to a requirements-intel.txt if you use one:
|
||||
#
|
||||
# --index-url https://download.pytorch.org/whl/xpu
|
||||
#
|
||||
# If you need to add extra flags into the pip install command you can do so by setting the variable EXTRA_PIP_INSTALL_FLAGS
|
||||
# before calling installRequirements. For example:
|
||||
#
|
||||
# source $(dirname $0)/../common/libbackend.sh
|
||||
# EXTRA_PIP_INSTALL_FLAGS="--no-build-isolation"
|
||||
# installRequirements
|
||||
function installRequirements() {
|
||||
ensureVenv
|
||||
declare -a requirementFiles=(
|
||||
"${EDIR}/requirements-install.txt"
|
||||
"${EDIR}/requirements.txt"
|
||||
"${EDIR}/requirements-${BUILD_TYPE:-}.txt"
|
||||
)
|
||||
|
||||
if [ "x${BUILD_TYPE:-}" != "x${BUILD_PROFILE}" ]; then
|
||||
requirementFiles+=("${EDIR}/requirements-${BUILD_PROFILE}.txt")
|
||||
fi
|
||||
if [ "x${BUILD_TYPE:-}" == "x" ]; then
|
||||
requirementFiles+=("${EDIR}/requirements-cpu.txt")
|
||||
fi
|
||||
requirementFiles+=("${EDIR}/requirements-after.txt")
|
||||
if [ "x${BUILD_TYPE:-}" != "x${BUILD_PROFILE}" ]; then
|
||||
requirementFiles+=("${EDIR}/requirements-${BUILD_PROFILE}-after.txt")
|
||||
fi
|
||||
|
||||
# This is needed to build wheels that e.g. depends on Python.h
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ]; then
|
||||
export C_INCLUDE_PATH="${C_INCLUDE_PATH:-}:$(_portable_dir)/include/python${PYTHON_VERSION}"
|
||||
fi
|
||||
|
||||
for reqFile in ${requirementFiles[@]}; do
|
||||
if [ -f "${reqFile}" ]; then
|
||||
echo "starting requirements install for ${reqFile}"
|
||||
if [ "x${USE_PIP}" == "xtrue" ]; then
|
||||
pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --requirement "${reqFile}"
|
||||
else
|
||||
uv pip install ${EXTRA_PIP_INSTALL_FLAGS:-} --requirement "${reqFile}"
|
||||
fi
|
||||
echo "finished requirements install for ${reqFile}"
|
||||
fi
|
||||
done
|
||||
|
||||
runProtogen
|
||||
}
|
||||
|
||||
# startBackend discovers and runs the backend GRPC server
|
||||
#
|
||||
# You can specify a specific backend file to execute by setting BACKEND_FILE before calling startBackend.
|
||||
# example:
|
||||
#
|
||||
# source ../common/libbackend.sh
|
||||
# BACKEND_FILE="${MY_DIR}/source/backend.py"
|
||||
# startBackend $@
|
||||
#
|
||||
# valid filenames for autodiscovered backend servers are:
|
||||
# - server.py
|
||||
# - backend.py
|
||||
# - ${BACKEND_NAME}.py
|
||||
function startBackend() {
|
||||
ensureVenv
|
||||
# Update pyvenv.cfg before running to ensure paths are correct for current location
|
||||
# This is critical when the backend position is dynamic (e.g., copied from container)
|
||||
if [ "x${PORTABLE_PYTHON}" == "xtrue" ] || [ -x "$(_portable_python)" ]; then
|
||||
_makeVenvPortable --update-pyvenv-cfg
|
||||
fi
|
||||
|
||||
# Set up GPU library paths if a lib directory exists
|
||||
# This allows backends to include their own GPU libraries (CUDA, ROCm, etc.)
|
||||
if [ -d "${EDIR}/lib" ]; then
|
||||
export LD_LIBRARY_PATH="${EDIR}/lib:${LD_LIBRARY_PATH:-}"
|
||||
echo "Added ${EDIR}/lib to LD_LIBRARY_PATH for GPU libraries"
|
||||
fi
|
||||
|
||||
if [ ! -z "${BACKEND_FILE:-}" ]; then
|
||||
exec "${EDIR}/venv/bin/python" "${BACKEND_FILE}" "$@"
|
||||
elif [ -e "${MY_DIR}/server.py" ]; then
|
||||
exec "${EDIR}/venv/bin/python" "${MY_DIR}/server.py" "$@"
|
||||
elif [ -e "${MY_DIR}/backend.py" ]; then
|
||||
exec "${EDIR}/venv/bin/python" "${MY_DIR}/backend.py" "$@"
|
||||
elif [ -e "${MY_DIR}/${BACKEND_NAME}.py" ]; then
|
||||
exec "${EDIR}/venv/bin/python" "${MY_DIR}/${BACKEND_NAME}.py" "$@"
|
||||
fi
|
||||
}
|
||||
|
||||
|
||||
# runUnittests discovers and runs python unittests
|
||||
#
|
||||
# You can specify a specific test file to use by setting TEST_FILE before calling runUnittests.
|
||||
# example:
|
||||
#
|
||||
# source ../common/libbackend.sh
|
||||
# TEST_FILE="${MY_DIR}/source/test.py"
|
||||
# runUnittests $@
|
||||
#
|
||||
# be default a file named test.py in the backends directory will be used
|
||||
function runUnittests() {
|
||||
ensureVenv
|
||||
if [ ! -z "${TEST_FILE:-}" ]; then
|
||||
testDir=$(dirname "$(realpath "${TEST_FILE}")")
|
||||
testFile=$(basename "${TEST_FILE}")
|
||||
pushd "${testDir}" >/dev/null
|
||||
python -m unittest "${testFile}"
|
||||
popd >/dev/null
|
||||
elif [ -f "${MY_DIR}/test.py" ]; then
|
||||
pushd "${MY_DIR}" >/dev/null
|
||||
python -m unittest test.py
|
||||
popd >/dev/null
|
||||
else
|
||||
echo "no tests defined for ${BACKEND_NAME}"
|
||||
fi
|
||||
}
|
||||
|
||||
|
||||
##################################################################################
|
||||
# Below here are helper functions not intended to be used outside of the library #
|
||||
##################################################################################
|
||||
|
||||
# checkTargets determines if the current BUILD_TYPE or BUILD_PROFILE is in a list of valid targets
|
||||
function checkTargets() {
|
||||
targets=$@
|
||||
declare -a targets=($targets)
|
||||
for target in ${targets[@]}; do
|
||||
if [ "x${BUILD_TYPE:-}" == "x${target}" ]; then
|
||||
echo true; return 0
|
||||
fi
|
||||
if [ "x${BUILD_PROFILE}" == "x${target}" ]; then
|
||||
echo true; return 0
|
||||
fi
|
||||
done
|
||||
echo false
|
||||
}
|
||||
|
||||
init
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Shared utilities for the mlx and mlx-vlm gRPC backends.
|
||||
|
||||
These helpers wrap mlx-lm's and mlx-vlm's native tool-parser modules, which
|
||||
auto-detect the right parser from the model's chat template. Each tool
|
||||
module exposes ``tool_call_start``, ``tool_call_end`` and
|
||||
``parse_tool_call(text, tools) -> dict | list[dict]``.
|
||||
|
||||
The split-reasoning helper is generic enough to work with any think-start /
|
||||
think-end delimiter pair.
|
||||
"""
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import uuid
|
||||
|
||||
|
||||
def split_reasoning(text, think_start, think_end):
|
||||
"""Split ``<think>...</think>`` blocks out of ``text``.
|
||||
|
||||
Returns ``(reasoning_content, remaining_text)``. When ``think_start`` is
|
||||
empty or not found, returns ``("", text)`` unchanged.
|
||||
"""
|
||||
if not think_start or not text:
|
||||
return "", text
|
||||
if think_start not in text:
|
||||
# Models like Qwen3.5 open assistant turns already INSIDE thinking, so
|
||||
# the generated text carries only the closing tag. Everything before it
|
||||
# is reasoning that would otherwise leak into the content.
|
||||
if think_end and think_end in text:
|
||||
head, _, tail = text.partition(think_end)
|
||||
return head.strip(), tail.strip()
|
||||
return "", text
|
||||
pattern = re.compile(
|
||||
re.escape(think_start) + r"(.*?)" + re.escape(think_end or ""),
|
||||
re.DOTALL,
|
||||
)
|
||||
reasoning_parts = pattern.findall(text)
|
||||
if not reasoning_parts:
|
||||
return "", text
|
||||
remaining = pattern.sub("", text).strip()
|
||||
return "\n".join(p.strip() for p in reasoning_parts), remaining
|
||||
|
||||
|
||||
def parse_tool_calls(text, tool_module, tools):
|
||||
"""Extract tool calls from ``text`` using a mlx-lm tool module.
|
||||
|
||||
Ports the ``process_tool_calls`` logic from
|
||||
``mlx_vlm/server.py`` (v0.10 onwards). ``tool_module`` must expose
|
||||
``tool_call_start``, ``tool_call_end`` and ``parse_tool_call``.
|
||||
|
||||
Returns ``(calls, remaining_text)`` where ``calls`` is a list of dicts:
|
||||
|
||||
[{"index": int, "id": str, "name": str, "arguments": str (JSON)}]
|
||||
|
||||
and ``remaining_text`` is the free-form text with the tool call blocks
|
||||
removed. ``(calls, text)`` is returned unchanged if ``tool_module`` is
|
||||
``None`` or the start delimiter isn't present.
|
||||
"""
|
||||
if tool_module is None or not text:
|
||||
return [], text
|
||||
start = getattr(tool_module, "tool_call_start", None)
|
||||
end = getattr(tool_module, "tool_call_end", None)
|
||||
parse_fn = getattr(tool_module, "parse_tool_call", None)
|
||||
if not start or parse_fn is None or start not in text:
|
||||
return [], text
|
||||
|
||||
if end == "" or end is None:
|
||||
pattern = re.compile(
|
||||
re.escape(start) + r".*?(?:\n|$)",
|
||||
re.DOTALL,
|
||||
)
|
||||
else:
|
||||
pattern = re.compile(
|
||||
re.escape(start) + r".*?" + re.escape(end),
|
||||
re.DOTALL,
|
||||
)
|
||||
|
||||
matches = pattern.findall(text)
|
||||
if not matches:
|
||||
return [], text
|
||||
|
||||
remaining = pattern.sub(" ", text).strip()
|
||||
calls = []
|
||||
for match in matches:
|
||||
call_body = match.strip().removeprefix(start)
|
||||
if end:
|
||||
call_body = call_body.removesuffix(end)
|
||||
call_body = call_body.strip()
|
||||
try:
|
||||
parsed = parse_fn(call_body, tools)
|
||||
except Exception as e:
|
||||
print(
|
||||
f"[mlx_utils] Invalid tool call: {call_body!r} ({e})",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
if not isinstance(parsed, list):
|
||||
parsed = [parsed]
|
||||
for tc in parsed:
|
||||
calls.append(
|
||||
{
|
||||
"index": len(calls),
|
||||
"id": str(uuid.uuid4()),
|
||||
"name": (tc.get("name") or "").strip(),
|
||||
"arguments": json.dumps(tc.get("arguments", {}), ensure_ascii=False),
|
||||
}
|
||||
)
|
||||
return calls, remaining
|
||||
@@ -0,0 +1,75 @@
|
||||
"""Unit tests for the mlx/mlx-vlm shared helpers (mlx_utils.py).
|
||||
|
||||
Run standalone (Python standard library only, no backend venv needed):
|
||||
python3 -m unittest mlx_utils_test
|
||||
|
||||
These mirror the server-less helper tests in backend/python/mlx/test.py
|
||||
(TestSharedHelpers), but live here so they run on any platform: the mlx
|
||||
test module imports grpc/backend_pb2 at import time and needs the MLX venv,
|
||||
whereas mlx_utils only needs the standard library.
|
||||
"""
|
||||
|
||||
import types
|
||||
import unittest
|
||||
|
||||
from mlx_utils import parse_tool_calls, split_reasoning
|
||||
|
||||
|
||||
class TestSplitReasoning(unittest.TestCase):
|
||||
def test_both_tags(self):
|
||||
r, c = split_reasoning(
|
||||
"<think>step 1\nstep 2</think>The answer is 42.", "<think>", "</think>"
|
||||
)
|
||||
self.assertEqual(r, "step 1\nstep 2")
|
||||
self.assertEqual(c, "The answer is 42.")
|
||||
|
||||
def test_implicit_opener_only_closing_tag(self):
|
||||
# Qwen3.5 opens the assistant turn already inside thinking, so the
|
||||
# output carries only the closing tag; everything before it is reasoning.
|
||||
r, c = split_reasoning(
|
||||
"The user is asking about the weather.\n</think>\n\nThe weather in Rome is sunny.",
|
||||
"<think>",
|
||||
"</think>",
|
||||
)
|
||||
self.assertEqual(r, "The user is asking about the weather.")
|
||||
self.assertEqual(c, "The weather in Rome is sunny.")
|
||||
|
||||
def test_no_tags_at_all(self):
|
||||
r, c = split_reasoning("just text", "<think>", "</think>")
|
||||
self.assertEqual(r, "")
|
||||
self.assertEqual(c, "just text")
|
||||
|
||||
def test_empty_think_end_and_no_opener_match(self):
|
||||
# No think_end to anchor on, and the opener is absent → return unchanged.
|
||||
r, c = split_reasoning("no opener here", "<think>", "")
|
||||
self.assertEqual(r, "")
|
||||
self.assertEqual(c, "no opener here")
|
||||
|
||||
def test_empty_text(self):
|
||||
r, c = split_reasoning("", "<think>", "</think>")
|
||||
self.assertEqual(r, "")
|
||||
self.assertEqual(c, "")
|
||||
|
||||
|
||||
class TestParseToolCalls(unittest.TestCase):
|
||||
def test_with_shim(self):
|
||||
tm = types.SimpleNamespace(
|
||||
tool_call_start="<tool_call>",
|
||||
tool_call_end="</tool_call>",
|
||||
parse_tool_call=lambda body, tools: {
|
||||
"name": "get_weather",
|
||||
"arguments": {"location": body.strip()},
|
||||
},
|
||||
)
|
||||
calls, remaining = parse_tool_calls(
|
||||
"Sure: <tool_call>Paris</tool_call>", tm, tools=None
|
||||
)
|
||||
self.assertEqual(len(calls), 1)
|
||||
self.assertEqual(calls[0]["name"], "get_weather")
|
||||
self.assertEqual(calls[0]["arguments"], '{"location": "Paris"}')
|
||||
self.assertEqual(calls[0]["index"], 0)
|
||||
self.assertNotIn("<tool_call>", remaining)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Parent-death watcher (best-effort backstop) for LocalAI Python backends.
|
||||
|
||||
LocalAI spawns each backend as a child process and, on a clean shutdown, tears
|
||||
it down itself (SIGTERM -> grace -> SIGKILL). That graceful path only runs when
|
||||
LocalAI receives a catchable signal and lives long enough to run its handlers.
|
||||
If LocalAI is SIGKILLed (e.g. a supervising process's grace period elapses
|
||||
first), that teardown never runs and this backend would be reparented to init
|
||||
and linger, holding GPU/VRAM and its listen port.
|
||||
|
||||
The watcher here is a best-effort backstop for exactly that case: it does NOT
|
||||
replace the graceful teardown, it only covers the "parent vanished without
|
||||
cleaning up" path. It detects reparenting: when the process that spawned this
|
||||
backend dies, the kernel reparents us to the nearest sub-reaper or to init
|
||||
(PID 1), so os.getppid() stops matching the value captured at startup. This
|
||||
getppid() approach is portable across Linux/macOS (unlike the Linux-only
|
||||
PR_SET_PDEATHSIG), which is why it is used here, mirroring the Go backends'
|
||||
pkg/grpc/parentwatch.go and the C++ backends' parent_watch.h. It is disabled on
|
||||
Windows, which has no equivalent orphan-reparenting semantics.
|
||||
|
||||
Env vars (shared verbatim across the Go, C++ and Python backends):
|
||||
LOCALAI_BACKEND_PARENT_WATCH enabled by default; a falsey value
|
||||
("false"/"0"/"no"/"off", case-insensitive)
|
||||
disables it.
|
||||
LOCALAI_BACKEND_PARENT_WATCH_INTERVAL poll interval as a Go-style duration
|
||||
string ("500ms", "2s", "1m") or a bare
|
||||
number of seconds. Defaults to 2s.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
|
||||
ENV_PARENT_WATCH = "LOCALAI_BACKEND_PARENT_WATCH"
|
||||
ENV_PARENT_WATCH_INTERVAL = "LOCALAI_BACKEND_PARENT_WATCH_INTERVAL"
|
||||
|
||||
_DEFAULT_INTERVAL_SECONDS = 2.0
|
||||
|
||||
# Guard so repeated calls (e.g. get_auth_interceptors invoked more than once)
|
||||
# only ever arm a single watcher thread per process.
|
||||
_started = False
|
||||
_started_lock = threading.Lock()
|
||||
|
||||
|
||||
def _enabled():
|
||||
"""Report whether the watcher should run in this process."""
|
||||
# Windows does not reparent orphans to a well-known init PID, so the
|
||||
# getppid() heuristic used here doesn't apply there.
|
||||
if os.name == "nt" or sys.platform.startswith("win"):
|
||||
return False
|
||||
val = os.environ.get(ENV_PARENT_WATCH, "").strip().lower()
|
||||
if val in ("false", "0", "no", "off"):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _interval_seconds():
|
||||
"""Return the configured poll interval in seconds, or the default.
|
||||
|
||||
Accepts Go-style duration strings ("500ms", "2s", "1m") for cross-language
|
||||
parity, or a bare number interpreted as seconds.
|
||||
"""
|
||||
raw = os.environ.get(ENV_PARENT_WATCH_INTERVAL, "").strip()
|
||||
if not raw:
|
||||
return _DEFAULT_INTERVAL_SECONDS
|
||||
# Split numeric prefix from unit suffix.
|
||||
i = 0
|
||||
while i < len(raw) and (raw[i].isdigit() or raw[i] == "." or (i == 0 and raw[i] in "+-")):
|
||||
i += 1
|
||||
if i == 0:
|
||||
return _DEFAULT_INTERVAL_SECONDS
|
||||
try:
|
||||
num = float(raw[:i])
|
||||
except ValueError:
|
||||
return _DEFAULT_INTERVAL_SECONDS
|
||||
unit = raw[i:].lower()
|
||||
if unit == "ms":
|
||||
seconds = num / 1000.0
|
||||
elif unit in ("s", ""):
|
||||
seconds = num
|
||||
elif unit == "m":
|
||||
seconds = num * 60.0
|
||||
else:
|
||||
return _DEFAULT_INTERVAL_SECONDS
|
||||
return seconds if seconds > 0 else _DEFAULT_INTERVAL_SECONDS
|
||||
|
||||
|
||||
def _parent_died(orig_ppid):
|
||||
"""Report whether this process has been reparented away from orig_ppid.
|
||||
|
||||
Reparenting is the standard POSIX signal that the original parent (here, the
|
||||
LocalAI process that spawned this backend) has exited: the orphan is handed
|
||||
to the nearest sub-reaper or to init (PID 1), so os.getppid() no longer
|
||||
matches the value captured at startup.
|
||||
"""
|
||||
ppid = os.getppid()
|
||||
return ppid != orig_ppid or ppid == 1
|
||||
|
||||
|
||||
def _watch(orig_ppid, interval, on_death):
|
||||
"""Poll until _parent_died reports the original parent is gone, then call
|
||||
on_death. Blocks, so run it on its own (daemon) thread."""
|
||||
import time
|
||||
|
||||
while True:
|
||||
time.sleep(interval)
|
||||
if _parent_died(orig_ppid):
|
||||
on_death()
|
||||
return
|
||||
|
||||
|
||||
def start_parent_death_watcher():
|
||||
"""Install the best-effort safety net described in this module's docstring.
|
||||
|
||||
No-op when disabled, on Windows, when already orphaned at startup
|
||||
(os.getppid() <= 1), or if already started. This is a backstop alongside —
|
||||
never a replacement for — LocalAI's graceful teardown.
|
||||
"""
|
||||
global _started
|
||||
if not _enabled():
|
||||
return
|
||||
with _started_lock:
|
||||
if _started:
|
||||
return
|
||||
orig_ppid = os.getppid()
|
||||
# A parent of 1 (or less) at startup means we were already orphaned (or
|
||||
# launched directly under init) — there is no original parent to watch.
|
||||
if orig_ppid <= 1:
|
||||
return
|
||||
interval = _interval_seconds()
|
||||
|
||||
def on_death():
|
||||
print(
|
||||
"backend parent process (pid {}) exited without stopping this "
|
||||
"backend; self-terminating to avoid orphaning".format(orig_ppid),
|
||||
file=sys.stderr,
|
||||
flush=True,
|
||||
)
|
||||
# Immediate, non-cleanup exit: this is a shutdown safety net and the
|
||||
# normal graceful path is already gone.
|
||||
os._exit(1)
|
||||
|
||||
thread = threading.Thread(
|
||||
target=_watch,
|
||||
args=(orig_ppid, interval, on_death),
|
||||
name="parent-death-watcher",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
_started = True
|
||||
@@ -0,0 +1,150 @@
|
||||
"""Unit tests for the parent-death watcher (parent_watch.py).
|
||||
|
||||
Run standalone (Python standard library only, no backend venv needed):
|
||||
python3 -m unittest parent_watch_test
|
||||
|
||||
The core test (test_detects_reparent) builds a genuine two-level process tree
|
||||
(test -> middle -> grandchild) with os.fork, lets the middle process die, and
|
||||
asserts the grandchild's parent_watch._watch detects the reparenting and
|
||||
self-terminates — mirroring the Go test in pkg/grpc/parentwatch_test.go and the
|
||||
C++ test in backend/cpp/llama-cpp/parent_watch_test.cpp.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
|
||||
import parent_watch
|
||||
|
||||
|
||||
class TestParentWatchEnvParsing(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._saved = {
|
||||
k: os.environ.get(k)
|
||||
for k in (parent_watch.ENV_PARENT_WATCH, parent_watch.ENV_PARENT_WATCH_INTERVAL)
|
||||
}
|
||||
for k in self._saved:
|
||||
os.environ.pop(k, None)
|
||||
|
||||
def tearDown(self):
|
||||
for k, v in self._saved.items():
|
||||
if v is None:
|
||||
os.environ.pop(k, None)
|
||||
else:
|
||||
os.environ[k] = v
|
||||
|
||||
def test_interval_default(self):
|
||||
self.assertEqual(parent_watch._interval_seconds(), 2.0)
|
||||
|
||||
def test_interval_units(self):
|
||||
cases = {"500ms": 0.5, "2s": 2.0, "1m": 60.0, "3": 3.0, "0.5s": 0.5}
|
||||
for raw, expected in cases.items():
|
||||
os.environ[parent_watch.ENV_PARENT_WATCH_INTERVAL] = raw
|
||||
self.assertAlmostEqual(parent_watch._interval_seconds(), expected, msg=raw)
|
||||
|
||||
def test_interval_garbage_falls_back(self):
|
||||
os.environ[parent_watch.ENV_PARENT_WATCH_INTERVAL] = "garbage"
|
||||
self.assertEqual(parent_watch._interval_seconds(), 2.0)
|
||||
|
||||
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "POSIX only")
|
||||
def test_enabled_default(self):
|
||||
self.assertTrue(parent_watch._enabled())
|
||||
|
||||
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "POSIX only")
|
||||
def test_disabled_by_falsey(self):
|
||||
for val in ("false", "0", "no", "off", "OFF", " False "):
|
||||
os.environ[parent_watch.ENV_PARENT_WATCH] = val
|
||||
self.assertFalse(parent_watch._enabled(), msg=val)
|
||||
|
||||
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "POSIX only")
|
||||
def test_enabled_by_truthy(self):
|
||||
for val in ("true", "1", "yes", "on"):
|
||||
os.environ[parent_watch.ENV_PARENT_WATCH] = val
|
||||
self.assertTrue(parent_watch._enabled(), msg=val)
|
||||
|
||||
|
||||
@unittest.skipIf(os.name == "nt" or sys.platform.startswith("win"), "fork/reparent is POSIX only")
|
||||
class TestParentWatchReparent(unittest.TestCase):
|
||||
def _wait_for_file(self, path, timeout=10.0):
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if os.path.exists(path):
|
||||
return True
|
||||
time.sleep(0.02)
|
||||
return False
|
||||
|
||||
def test_detects_reparent(self):
|
||||
tmpdir = tempfile.mkdtemp(prefix="parentwatch_test_")
|
||||
ready_file = os.path.join(tmpdir, "ready")
|
||||
exited_file = os.path.join(tmpdir, "exited")
|
||||
|
||||
middle = os.fork()
|
||||
if middle == 0:
|
||||
# ---- middle process ----
|
||||
grandchild = os.fork()
|
||||
if grandchild == 0:
|
||||
# ---- grandchild process: arm the REAL watcher against middle ----
|
||||
orig_ppid = os.getppid()
|
||||
|
||||
def on_death():
|
||||
with open(exited_file, "w") as f:
|
||||
f.write("1")
|
||||
os._exit(7)
|
||||
|
||||
threading.Thread(
|
||||
target=parent_watch._watch,
|
||||
args=(orig_ppid, 0.05, on_death),
|
||||
daemon=True,
|
||||
).start()
|
||||
|
||||
# Safety valve: never linger if something goes wrong.
|
||||
def bail():
|
||||
time.sleep(30)
|
||||
os._exit(2)
|
||||
|
||||
threading.Thread(target=bail, daemon=True).start()
|
||||
|
||||
# Signal readiness only after the watcher captured orig_ppid.
|
||||
with open(ready_file, "w") as f:
|
||||
f.write(str(os.getpid()))
|
||||
while True:
|
||||
time.sleep(1)
|
||||
else:
|
||||
# middle: wait until grandchild is ready, then exit to orphan it.
|
||||
if not self._wait_for_file(ready_file):
|
||||
os._exit(5)
|
||||
os._exit(0)
|
||||
|
||||
# ---- test (top) process ----
|
||||
os.waitpid(middle, 0) # reap middle only; grandchild is orphaned
|
||||
|
||||
self.assertTrue(os.path.exists(ready_file), "grandchild never signaled readiness")
|
||||
self.assertTrue(
|
||||
self._wait_for_file(exited_file),
|
||||
"watcher did not detect parent death within timeout",
|
||||
)
|
||||
|
||||
# Best-effort cleanup: kill the grandchild if it somehow survived.
|
||||
try:
|
||||
with open(ready_file) as f:
|
||||
pid = int(f.read().strip())
|
||||
if pid > 1:
|
||||
os.kill(pid, 9)
|
||||
except (OSError, ValueError):
|
||||
pass
|
||||
for p in (ready_file, exited_file):
|
||||
try:
|
||||
os.remove(p)
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
os.rmdir(tmpdir)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Generic utilities shared across Python gRPC backends.
|
||||
|
||||
These helpers don't depend on any specific inference framework and can be
|
||||
imported by any backend that needs to parse LocalAI gRPC options or build a
|
||||
chat-template-compatible message list from proto Message objects.
|
||||
"""
|
||||
import json
|
||||
|
||||
|
||||
def parse_options(options_list):
|
||||
"""Parse Options[] list of ``key:value`` strings into a dict.
|
||||
|
||||
Supports type inference for common cases (bool, int, float). Unknown or
|
||||
mixed-case values are returned as strings.
|
||||
|
||||
Used by LoadModel to extract backend-specific options passed via
|
||||
``ModelOptions.Options`` in ``backend.proto``.
|
||||
"""
|
||||
opts = {}
|
||||
for opt in options_list:
|
||||
if ":" not in opt:
|
||||
continue
|
||||
key, value = opt.split(":", 1)
|
||||
key = key.strip()
|
||||
value = value.strip()
|
||||
# Try type conversion
|
||||
if value.lower() in ("true", "false"):
|
||||
opts[key] = value.lower() == "true"
|
||||
else:
|
||||
try:
|
||||
opts[key] = int(value)
|
||||
except ValueError:
|
||||
try:
|
||||
opts[key] = float(value)
|
||||
except ValueError:
|
||||
opts[key] = value
|
||||
return opts
|
||||
|
||||
|
||||
def messages_to_dicts(proto_messages):
|
||||
"""Convert proto ``Message`` objects to dicts suitable for ``apply_chat_template``.
|
||||
|
||||
Handles: ``role``, ``content``, ``name``, ``tool_call_id``,
|
||||
``reasoning_content``, ``tool_calls`` (JSON string → Python list).
|
||||
|
||||
HuggingFace chat templates (and their MLX/vLLM wrappers) expect a list of
|
||||
plain dicts — proto Message objects don't work directly with Jinja, so
|
||||
this conversion is needed before every ``apply_chat_template`` call.
|
||||
"""
|
||||
result = []
|
||||
for msg in proto_messages:
|
||||
d = {"role": msg.role, "content": msg.content or ""}
|
||||
if msg.name:
|
||||
d["name"] = msg.name
|
||||
if msg.tool_call_id:
|
||||
d["tool_call_id"] = msg.tool_call_id
|
||||
if msg.reasoning_content:
|
||||
d["reasoning_content"] = msg.reasoning_content
|
||||
if msg.tool_calls:
|
||||
try:
|
||||
tool_calls = json.loads(msg.tool_calls)
|
||||
# Chat templates (e.g. Qwen) iterate function.arguments as a
|
||||
# mapping, but the OpenAI wire format carries it as a JSON
|
||||
# string — decode it back so the template's .items() works.
|
||||
for tc in tool_calls:
|
||||
fn = tc.get("function") if isinstance(tc, dict) else None
|
||||
if isinstance(fn, dict) and isinstance(fn.get("arguments"), str):
|
||||
try:
|
||||
fn["arguments"] = json.loads(fn["arguments"])
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
d["tool_calls"] = tool_calls
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
result.append(d)
|
||||
return result
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Unit tests for the shared python backend helpers (python_utils.py).
|
||||
|
||||
Run standalone (Python standard library only, no backend venv needed):
|
||||
python3 -m unittest python_utils_test
|
||||
|
||||
These mirror the server-less helper tests in backend/python/mlx/test.py
|
||||
(TestSharedHelpers), but live here so they run on any platform: the mlx
|
||||
test module imports grpc/backend_pb2 at import time and needs the MLX venv,
|
||||
whereas python_utils has no third-party dependency. Proto Message objects
|
||||
are faked with types.SimpleNamespace (real proto fields default to "").
|
||||
"""
|
||||
|
||||
import json
|
||||
import types
|
||||
import unittest
|
||||
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
|
||||
|
||||
def _msg(**fields):
|
||||
"""Fake a proto Message: every unset field is the empty string, as protobuf."""
|
||||
defaults = {
|
||||
"role": "",
|
||||
"content": "",
|
||||
"name": "",
|
||||
"tool_call_id": "",
|
||||
"reasoning_content": "",
|
||||
"tool_calls": "",
|
||||
}
|
||||
defaults.update(fields)
|
||||
return types.SimpleNamespace(**defaults)
|
||||
|
||||
|
||||
class TestParseOptions(unittest.TestCase):
|
||||
def test_type_inference(self):
|
||||
opts = parse_options(
|
||||
["temperature:0.7", "max_tokens:128", "trust:true", "name:hello", "no_colon_skipped"]
|
||||
)
|
||||
self.assertEqual(opts["temperature"], 0.7)
|
||||
self.assertEqual(opts["max_tokens"], 128)
|
||||
self.assertIs(opts["trust"], True)
|
||||
self.assertEqual(opts["name"], "hello")
|
||||
self.assertNotIn("no_colon_skipped", opts)
|
||||
|
||||
|
||||
class TestMessagesToDicts(unittest.TestCase):
|
||||
def test_basic_fields(self):
|
||||
out = messages_to_dicts(
|
||||
[
|
||||
_msg(role="user", content="hi"),
|
||||
_msg(role="tool", content="42", tool_call_id="call_1", name="f"),
|
||||
]
|
||||
)
|
||||
self.assertEqual(out[0], {"role": "user", "content": "hi"})
|
||||
self.assertEqual(out[1]["tool_call_id"], "call_1")
|
||||
self.assertEqual(out[1]["name"], "f")
|
||||
|
||||
def test_tool_call_arguments_string_decoded_to_mapping(self):
|
||||
# OpenAI wire format ships function.arguments as a JSON *string*; chat
|
||||
# templates iterate it as a mapping, so it must come back as a dict.
|
||||
out = messages_to_dicts(
|
||||
[
|
||||
_msg(
|
||||
role="assistant",
|
||||
tool_calls=json.dumps(
|
||||
[
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"location": "Rome"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
args = out[0]["tool_calls"][0]["function"]["arguments"]
|
||||
self.assertEqual(args, {"location": "Rome"})
|
||||
self.assertEqual(dict(args.items()), {"location": "Rome"})
|
||||
|
||||
def test_tool_call_arguments_already_mapping_is_idempotent(self):
|
||||
out = messages_to_dicts(
|
||||
[
|
||||
_msg(
|
||||
role="assistant",
|
||||
tool_calls=json.dumps(
|
||||
[{"function": {"name": "f", "arguments": {"a": 1}}}]
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
self.assertEqual(out[0]["tool_calls"][0]["function"]["arguments"], {"a": 1})
|
||||
|
||||
def test_tool_call_arguments_invalid_json_left_as_string(self):
|
||||
out = messages_to_dicts(
|
||||
[
|
||||
_msg(
|
||||
role="assistant",
|
||||
tool_calls=json.dumps(
|
||||
[{"function": {"name": "f", "arguments": "not-json"}}]
|
||||
),
|
||||
)
|
||||
]
|
||||
)
|
||||
self.assertEqual(out[0]["tool_calls"][0]["function"]["arguments"], "not-json")
|
||||
|
||||
def test_tool_call_without_function_key(self):
|
||||
out = messages_to_dicts(
|
||||
[_msg(role="assistant", tool_calls=json.dumps([{"id": "call_1"}]))]
|
||||
)
|
||||
self.assertEqual(out[0]["tool_calls"], [{"id": "call_1"}])
|
||||
|
||||
def test_tool_calls_invalid_json_dropped(self):
|
||||
out = messages_to_dicts([_msg(role="assistant", tool_calls="{not json")])
|
||||
self.assertNotIn("tool_calls", out[0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,13 @@
|
||||
.DEFAULT_GOAL := install
|
||||
|
||||
.PHONY: install
|
||||
install:
|
||||
bash install.sh
|
||||
|
||||
.PHONY: protogen-clean
|
||||
protogen-clean:
|
||||
$(RM) backend_pb2_grpc.py backend_pb2.py
|
||||
|
||||
.PHONY: clean
|
||||
clean: protogen-clean
|
||||
rm -rf venv __pycache__
|
||||
Executable
+4
@@ -0,0 +1,4 @@
|
||||
#!/usr/bin/env python3
|
||||
import grpc
|
||||
import backend_pb2
|
||||
import backend_pb2_grpc
|
||||
Executable
+19
@@ -0,0 +1,19 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
# This is here because the Intel pip index is broken and returns 200 status codes for every package name, it just doesn't return any package links.
|
||||
# This makes uv think that the package exists in the Intel pip index, and by default it stops looking at other pip indexes once it finds a match.
|
||||
# We need uv to continue falling through to the pypi default index to find optimum[openvino] in the pypi index
|
||||
# the --upgrade actually allows us to *downgrade* torch to the version provided in the Intel pip index
|
||||
if [ "x${BUILD_PROFILE}" == "xintel" ]; then
|
||||
EXTRA_PIP_INSTALL_FLAGS+=" --upgrade --index-strategy=unsafe-first-match"
|
||||
fi
|
||||
|
||||
installRequirements
|
||||
@@ -0,0 +1,11 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
runProtogen
|
||||
@@ -0,0 +1,2 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/rocm7.0
|
||||
torch
|
||||
@@ -0,0 +1,4 @@
|
||||
--extra-index-url https://download.pytorch.org/whl/xpu
|
||||
torch==2.8.0
|
||||
oneccl_bind_pt==2.8.0+xpu
|
||||
optimum[openvino]
|
||||
@@ -0,0 +1,3 @@
|
||||
grpcio==1.80.0
|
||||
protobuf
|
||||
grpcio-tools
|
||||
Executable
+9
@@ -0,0 +1,9 @@
|
||||
#!/bin/bash
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
startBackend $@
|
||||
Executable
+11
@@ -0,0 +1,11 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
|
||||
backend_dir=$(dirname $0)
|
||||
if [ -d $backend_dir/common ]; then
|
||||
source $backend_dir/common/libbackend.sh
|
||||
else
|
||||
source $backend_dir/../common/libbackend.sh
|
||||
fi
|
||||
|
||||
runUnittests
|
||||
@@ -0,0 +1,43 @@
|
||||
"""vLLM-specific helpers for the vllm and vllm-omni gRPC backends.
|
||||
|
||||
Generic helpers (``parse_options``, ``messages_to_dicts``) live in
|
||||
``python_utils`` and are re-exported here for backwards compatibility with
|
||||
existing imports in both backends.
|
||||
"""
|
||||
import sys
|
||||
|
||||
from python_utils import messages_to_dicts, parse_options
|
||||
|
||||
__all__ = ["parse_options", "messages_to_dicts", "setup_parsers"]
|
||||
|
||||
|
||||
def setup_parsers(opts):
|
||||
"""Return ``(tool_parser_cls, reasoning_parser_cls)`` from an opts dict.
|
||||
|
||||
Uses vLLM's native ``ToolParserManager`` / ``ReasoningParserManager``.
|
||||
Returns ``(None, None)`` if vLLM isn't installed or the requested
|
||||
parser name can't be resolved.
|
||||
"""
|
||||
tool_parser_cls = None
|
||||
reasoning_parser_cls = None
|
||||
|
||||
tool_parser_name = opts.get("tool_parser")
|
||||
reasoning_parser_name = opts.get("reasoning_parser")
|
||||
|
||||
if tool_parser_name:
|
||||
try:
|
||||
from vllm.tool_parsers import ToolParserManager
|
||||
tool_parser_cls = ToolParserManager.get_tool_parser(tool_parser_name)
|
||||
print(f"[vllm_utils] Loaded tool_parser: {tool_parser_name}", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"[vllm_utils] Failed to load tool_parser {tool_parser_name}: {e}", file=sys.stderr)
|
||||
|
||||
if reasoning_parser_name:
|
||||
try:
|
||||
from vllm.reasoning import ReasoningParserManager
|
||||
reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(reasoning_parser_name)
|
||||
print(f"[vllm_utils] Loaded reasoning_parser: {reasoning_parser_name}", file=sys.stderr)
|
||||
except Exception as e:
|
||||
print(f"[vllm_utils] Failed to load reasoning_parser {reasoning_parser_name}: {e}", file=sys.stderr)
|
||||
|
||||
return tool_parser_cls, reasoning_parser_cls
|
||||
Reference in New Issue
Block a user