mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-03 12:57:02 -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)]
|