Skip to content

Commit 34a91ff

Browse files
committed
feat: add distributed mode (experimental)
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent d3f629f commit 34a91ff

183 files changed

Lines changed: 25404 additions & 2479 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ For older news and full release notes, see [GitHub Releases](https://github.com/
154154
- [Object Detection](https://localai.io/features/object-detection/)
155155
- [Reranker API](https://localai.io/features/reranker/)
156156
- [P2P Inferencing](https://localai.io/features/distribute/)
157+
- [Distributed Mode](https://localai.io/features/distributed-mode/) — Horizontal scaling with PostgreSQL + NATS
157158
- [Model Context Protocol (MCP)](https://localai.io/docs/features/mcp/)
158159
- [Built-in Agents](https://localai.io/features/agents/) — Autonomous AI agents with tool use, RAG, skills, SSE streaming, and [Agent Hub](https://agenthub.localai.io)
159160
- [Backend Gallery](https://localai.io/backends/) — Install/remove backends on the fly via OCI images

backend/backend.proto

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ service Backend {
5151
rpc StartQuantization(QuantizationRequest) returns (QuantizationJobResult) {}
5252
rpc QuantizationProgress(QuantizationProgressRequest) returns (stream QuantizationProgressUpdate) {}
5353
rpc StopQuantization(QuantizationStopRequest) returns (Result) {}
54+
5455
}
5556

5657
// Define the empty request
@@ -676,3 +677,4 @@ message QuantizationProgressUpdate {
676677
message QuantizationStopRequest {
677678
string job_id = 1;
678679
}
680+

backend/cpp/llama-cpp/grpc-server.cpp

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,10 @@
2222
#include <grpcpp/ext/proto_server_reflection_plugin.h>
2323
#include <grpcpp/grpcpp.h>
2424
#include <grpcpp/health_check_service_interface.h>
25+
#include <grpcpp/security/server_credentials.h>
2526
#include <regex>
2627
#include <atomic>
28+
#include <cstdlib>
2729
#include <mutex>
2830
#include <signal.h>
2931
#include <thread>
@@ -37,6 +39,47 @@ using grpc::Server;
3739
using grpc::ServerBuilder;
3840
using grpc::ServerContext;
3941
using grpc::Status;
42+
43+
// gRPC bearer token auth via AuthMetadataProcessor for distributed mode.
44+
// Reads LOCALAI_GRPC_AUTH_TOKEN from the environment. When set, rejects
45+
// requests without a matching "authorization: Bearer <token>" metadata header.
46+
class TokenAuthMetadataProcessor : public grpc::AuthMetadataProcessor {
47+
public:
48+
explicit TokenAuthMetadataProcessor(const std::string& token) : token_(token) {}
49+
50+
bool IsBlocking() const override { return false; }
51+
52+
grpc::Status Process(const InputMetadata& auth_metadata,
53+
grpc::AuthContext* /*context*/,
54+
OutputMetadata* /*consumed_auth_metadata*/,
55+
OutputMetadata* /*response_metadata*/) override {
56+
auto it = auth_metadata.find("authorization");
57+
if (it != auth_metadata.end()) {
58+
std::string expected = "Bearer " + token_;
59+
std::string got(it->second.data(), it->second.size());
60+
// Constant-time comparison
61+
if (expected.size() == got.size() && ct_memcmp(expected.data(), got.data(), expected.size()) == 0) {
62+
return grpc::Status::OK;
63+
}
64+
}
65+
return grpc::Status(grpc::StatusCode::UNAUTHENTICATED, "invalid token");
66+
}
67+
68+
private:
69+
std::string token_;
70+
71+
// Minimal constant-time comparison (avoids OpenSSL dependency)
72+
static int ct_memcmp(const void* a, const void* b, size_t n) {
73+
const unsigned char* pa = static_cast<const unsigned char*>(a);
74+
const unsigned char* pb = static_cast<const unsigned char*>(b);
75+
unsigned char result = 0;
76+
for (size_t i = 0; i < n; i++) {
77+
result |= pa[i] ^ pb[i];
78+
}
79+
return result;
80+
}
81+
};
82+
4083
// END LocalAI
4184

4285

@@ -2760,11 +2803,24 @@ int main(int argc, char** argv) {
27602803
BackendServiceImpl service(ctx_server);
27612804

27622805
ServerBuilder builder;
2763-
builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
2806+
// Add bearer token auth via AuthMetadataProcessor if LOCALAI_GRPC_AUTH_TOKEN is set
2807+
const char* auth_token = std::getenv("LOCALAI_GRPC_AUTH_TOKEN");
2808+
std::shared_ptr<grpc::ServerCredentials> creds;
2809+
if (auth_token != nullptr && auth_token[0] != '\0') {
2810+
creds = grpc::InsecureServerCredentials();
2811+
creds->SetAuthMetadataProcessor(
2812+
std::make_shared<TokenAuthMetadataProcessor>(auth_token));
2813+
std::cout << "gRPC auth enabled via LOCALAI_GRPC_AUTH_TOKEN" << std::endl;
2814+
} else {
2815+
creds = grpc::InsecureServerCredentials();
2816+
}
2817+
2818+
builder.AddListeningPort(server_address, creds);
27642819
builder.RegisterService(&service);
27652820
builder.SetMaxMessageSize(50 * 1024 * 1024); // 50MB
27662821
builder.SetMaxSendMessageSize(50 * 1024 * 1024); // 50MB
27672822
builder.SetMaxReceiveMessageSize(50 * 1024 * 1024); // 50MB
2823+
27682824
std::unique_ptr<Server> server(builder.BuildAndStart());
27692825
// run the HTTP server in a thread - see comment below
27702826
std::thread t([&]()

backend/python/ace-step/backend.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,10 @@
1919
import backend_pb2
2020
import backend_pb2_grpc
2121
import grpc
22+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
23+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
24+
from grpc_auth import get_auth_interceptors
25+
2226
from acestep.inference import (
2327
GenerationParams,
2428
GenerationConfig,
@@ -444,6 +448,8 @@ def serve(address):
444448
("grpc.max_send_message_length", 50 * 1024 * 1024),
445449
("grpc.max_receive_message_length", 50 * 1024 * 1024),
446450
],
451+
452+
interceptors=get_auth_interceptors(),
447453
)
448454
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
449455
server.add_insecure_port(address)

backend/python/chatterbox/backend.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,10 @@
1616
from chatterbox.tts import ChatterboxTTS
1717
from chatterbox.mtl_tts import ChatterboxMultilingualTTS
1818
import grpc
19+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
20+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
21+
from grpc_auth import get_auth_interceptors
22+
1923
import tempfile
2024

2125
def is_float(s):
@@ -225,7 +229,9 @@ def serve(address):
225229
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
226230
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
227231
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
228-
])
232+
],
233+
interceptors=get_auth_interceptors(),
234+
)
229235
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
230236
server.add_insecure_port(address)
231237
server.start()

backend/python/common/grpc_auth.py

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
"""Shared gRPC bearer token authentication interceptor for LocalAI Python backends.
2+
3+
When the environment variable LOCALAI_GRPC_AUTH_TOKEN is set, requests without
4+
a valid Bearer token in the 'authorization' metadata header are rejected with
5+
UNAUTHENTICATED. When the variable is empty or unset, no authentication is
6+
performed (backward compatible).
7+
"""
8+
9+
import hmac
10+
import os
11+
12+
import grpc
13+
14+
15+
class _AbortHandler(grpc.RpcMethodHandler):
16+
"""A method handler that immediately aborts with UNAUTHENTICATED."""
17+
18+
def __init__(self):
19+
self.request_streaming = False
20+
self.response_streaming = False
21+
self.request_deserializer = None
22+
self.response_serializer = None
23+
self.unary_unary = self._abort
24+
self.unary_stream = None
25+
self.stream_unary = None
26+
self.stream_stream = None
27+
28+
@staticmethod
29+
def _abort(request, context):
30+
context.abort(grpc.StatusCode.UNAUTHENTICATED, "invalid token")
31+
32+
33+
class TokenAuthInterceptor(grpc.ServerInterceptor):
34+
"""Sync gRPC server interceptor that validates a bearer token."""
35+
36+
def __init__(self, token: str):
37+
self._token = token
38+
self._abort_handler = _AbortHandler()
39+
40+
def intercept_service(self, continuation, handler_call_details):
41+
metadata = dict(handler_call_details.invocation_metadata)
42+
auth = metadata.get("authorization", "")
43+
expected = "Bearer " + self._token
44+
if not hmac.compare_digest(auth, expected):
45+
return self._abort_handler
46+
return continuation(handler_call_details)
47+
48+
49+
class AsyncTokenAuthInterceptor(grpc.aio.ServerInterceptor):
50+
"""Async gRPC server interceptor that validates a bearer token."""
51+
52+
def __init__(self, token: str):
53+
self._token = token
54+
55+
async def intercept_service(self, continuation, handler_call_details):
56+
metadata = dict(handler_call_details.invocation_metadata)
57+
auth = metadata.get("authorization", "")
58+
expected = "Bearer " + self._token
59+
if not hmac.compare_digest(auth, expected):
60+
return _AbortHandler()
61+
return await continuation(handler_call_details)
62+
63+
64+
def get_auth_interceptors(*, aio: bool = False):
65+
"""Return a list of gRPC interceptors for bearer token auth.
66+
67+
Args:
68+
aio: If True, return async-compatible interceptors for grpc.aio.server().
69+
If False (default), return sync interceptors for grpc.server().
70+
71+
Returns an empty list when LOCALAI_GRPC_AUTH_TOKEN is not set.
72+
"""
73+
token = os.environ.get("LOCALAI_GRPC_AUTH_TOKEN", "")
74+
if not token:
75+
return []
76+
if aio:
77+
return [AsyncTokenAuthInterceptor(token)]
78+
return [TokenAuthInterceptor(token)]

backend/python/coqui/backend.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@
1515
from TTS.api import TTS
1616

1717
import grpc
18+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
19+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
20+
from grpc_auth import get_auth_interceptors
21+
1822

1923

2024
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
@@ -93,7 +97,9 @@ def serve(address):
9397
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
9498
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
9599
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
96-
])
100+
],
101+
interceptors=get_auth_interceptors(),
102+
)
97103
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
98104
server.add_insecure_port(address)
99105
server.start()

backend/python/diffusers/backend.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@
2222
import backend_pb2_grpc
2323

2424
import grpc
25+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
26+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
27+
from grpc_auth import get_auth_interceptors
28+
2529

2630
# Import dynamic loader for pipeline discovery
2731
from diffusers_dynamic_loader import (
@@ -1042,7 +1046,9 @@ def serve(address):
10421046
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
10431047
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
10441048
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
1045-
])
1049+
],
1050+
interceptors=get_auth_interceptors(),
1051+
)
10461052
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
10471053
server.add_insecure_port(address)
10481054
server.start()

backend/python/faster-qwen3-tts/backend.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@
1515
import soundfile as sf
1616

1717
import grpc
18+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
19+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
20+
from grpc_auth import get_auth_interceptors
21+
1822

1923

2024
def is_float(s):
@@ -165,6 +169,8 @@ def serve(address):
165169
('grpc.max_send_message_length', 50 * 1024 * 1024),
166170
('grpc.max_receive_message_length', 50 * 1024 * 1024),
167171
]
172+
,
173+
interceptors=get_auth_interceptors(),
168174
)
169175
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
170176
server.add_insecure_port(address)

backend/python/faster-whisper/backend.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@
1414
from faster_whisper import WhisperModel
1515

1616
import grpc
17+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'common'))
18+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'common'))
19+
from grpc_auth import get_auth_interceptors
20+
1721

1822

1923
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
@@ -70,7 +74,9 @@ def serve(address):
7074
('grpc.max_message_length', 50 * 1024 * 1024), # 50MB
7175
('grpc.max_send_message_length', 50 * 1024 * 1024), # 50MB
7276
('grpc.max_receive_message_length', 50 * 1024 * 1024), # 50MB
73-
])
77+
],
78+
interceptors=get_auth_interceptors(),
79+
)
7480
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
7581
server.add_insecure_port(address)
7682
server.start()

0 commit comments

Comments
 (0)