mirror of
https://github.com/mudler/LocalAI.git
synced 2026-08-01 11:00:24 -04:00
a4e6e01e4d
* fix(grpc): self-terminate backend workers when LocalAI dies non-gracefully Symptom: a backend model-worker subprocess (the per-model gRPC server LocalAI spawns) can be orphaned and linger — holding VRAM and its listen port — if the LocalAI process is killed non-gracefully (e.g. a supervisor's graceful-shutdown grace period elapses and LocalAI is SIGKILLed) before its own teardown runs. Root cause: LocalAI's graceful teardown (pkg/signals/handler.go installs the SIGINT/SIGTERM handler; core/cli/run.go registers app.Shutdown -> ModelLoader.StopAllGRPC -> process.Stop in pkg/model/process.go) only runs when LocalAI receives a catchable signal and survives long enough to run its handlers. Backends are spawned via github.com/mudler/go-processmanager v0.1.1, whose getSysProcAttr() sets Setpgid:true (own process group, so the group can be signalled) but never PR_SET_PDEATHSIG/Pdeathsig, and exposes no Config field or option for a caller to inject/extend SysProcAttr. LocalAI fully delegates spawning to that library (it never builds the exec.Cmd itself), so it cannot set a kernel parent-death signal at the spawn site. If LocalAI is SIGKILLed, nothing tells the backend to exit and it is reparented to init. Fix: add a best-effort, backend-side safety net at the one shared choke point every out-of-process Go backend routes through — grpc.StartServer / RunServer in pkg/grpc. On startup it captures getppid() and polls; when the process is reparented (getppid changes / becomes 1 — the standard POSIX signal the original parent died) it logs and self-terminates. getppid() reparent detection is portable (Linux + macOS), unlike Linux-only PR_SET_PDEATHSIG. Toggle via LOCALAI_BACKEND_PARENT_WATCH (default on; off on Windows) and LOCALAI_BACKEND_PARENT_WATCH_INTERVAL. This is strictly a backstop alongside the existing graceful SIGTERM->grace->SIGKILL teardown, which is unchanged. Scope/limitations: covers Go-based backends (everything using pkg/grpc). The C++ backends (e.g. llama-cpp) and Python backends do not route through pkg/grpc and are not covered by this mechanism — they would each need an equivalent parent-death check (follow-up). The fully general fix is for go-processmanager to expose SysProcAttr injection so LocalAI can set Pdeathsig at spawn for every backend regardless of language (suggested upstream follow-up; out of scope for this LocalAI-only PR). Test: pkg/grpc/parentwatch_test.go builds a real test -> middle -> grandchild process tree, lets the middle process exit to orphan the grandchild running the real watchParentDeath, and asserts it detects the reparent and self-terminates. Unix-only (build-tagged), runs in CI (Linux). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * fix(process): extend parent-death backstop to C++ and Python backends The Go parent-death watcher (pkg/grpc/parentwatch.go, commit 772b435d5) only protects backends that route through pkg/grpc. C++ and Python backends don't, so the originally-reported case — the llama.cpp gRPC worker surviving a non-graceful LocalAI death — was still uncovered. Extend the same best-effort backstop to both languages, reusing the exact mechanism and semantics: - capture getppid() at startup, skip if already orphaned (<=1) - a background thread polls getppid() and self-exits on reparenting (getppid() != orig || == 1), portable across Linux/macOS, no-op on Windows - same env vars: LOCALAI_BACKEND_PARENT_WATCH (default on; falsy false/0/no/off disable) and LOCALAI_BACKEND_PARENT_WATCH_INTERVAL (default 2s; accepts Go-style durations like 500ms/2s/1m) C++: implemented in backend/cpp/llama-cpp (the reported, most-used C++ backend) as a dependency-free header parent_watch.h, wired into grpc-server.cpp's main() and copied at build time via prepare.sh. C++ backends have no shared server scaffolding, so other C++ backends (ds4, ik-llama-cpp, privacy-filter, ...) are not yet covered and would each need the same one-line include+call as follow-ups. Python: implemented once in the shared common/parent_watch.py and armed from common/grpc_auth.py's get_auth_interceptors() — the single helper every one of the 35 Python backends invokes while building its gRPC server — so all Python backends (and future ones) are covered with no per-backend edits and no duplicated implementation. Tests (real process-tree reparent detection, mirroring the Go test): - backend/cpp/llama-cpp/parent_watch_test.cpp (via run-unit-tests.sh) - backend/python/common/parent_watch_test.py (python -m unittest) Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
"""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)]
|