fix(grpc): let the transport answer through the wrappers, not only past gRPC

Re-review round 2. One blocking defect, and it was the concern I filed myself
last round and mis-scoped as a future trap. It was live, and it sat on the most
destructive reaping path of the five.

RouteResult.Client is an InFlightTrackingClient, over a FileStagingClient when a
stager is configured. model_router puts that on the cached remote model and
pkg/model's checkIsLoaded asks IT whether the transport failed. Both wrappers
embed grpc.Backend, which does not declare LastDialError, so the type assertion
read nil and the guard added last round fell straight through to the old
eviction. That eviction sends backend.stop over NATS to every node holding the
model and deletes every replica row, where the other sites delete one. The spec
covering it built a bare client by hand, which is why it passed while production
did not.

This is the third time in this task a correct fix was disarmed one layer out, so
the fix is a mechanism rather than two methods. BackendUnwrapper is one line per
decorator, LastDialErrorOf walks the chain, and both consumers now call it
instead of each keeping its own assertion. One implementation, no per-caller
policy to get wrong.

Sweeping every type that embeds or holds a grpc.Backend found a third decorator
the review had not named, and it is itself a reaping consumer of the same
collapsed signal. ConnectionEvictingClient is built for remote models in
initializers.go and its evict callback runs ShutdownModel; it fires during
INFERENCE rather than on a health check, so a tunnel blip mid-request was enough
to stop a model that was loaded and serving. It consults the transport first
now. A locally spawned backend has no custom transport, so that path is
unchanged byte for byte. Everything else touching a Backend is a consumer rather
than a decorator; there is no fourth.

The probe cache joiner shape is pinned. It was the right design last round with
nothing holding it: the mutation back to a closed-over variable passed all 602
specs in the package. Eight goroutines coalesced on a probe that blocks on a
channel now assert every joiner gets the leader's REASON and not just its
answer, which is the difference between a leader declining to reap and its seven
joiners reaping on the leader's own observation.

The LastDialError scope note claimed an exactness it does not have at
checkIsLoaded, which reads a shared long-lived client after releasing opMutex.
It now says which caller is not exact, why the imprecision is accepted there,
and what making it exact would cost.

The four-outcome table in the docs still said a worker with no live owner is
treated as absent and rescheduled, contradicting the code and the paragraph nine
lines below it. None of those outcomes is absence any more, and the table says
so, names the fifth, and points at the heartbeat as the thing that does decide
presence.

Five mutations, each reddening named specs, including the two the reviewer found
surviving.

Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-09-01 16:33:59 +00:00
1 parent a8ac2af167
commit b4d8e23abb
11 files changed
+362 -33

No files matched your search

@@ -37,6 +37,14 @@ type FileStagingClient struct {
remoteModelPath string // set during LoadModel from staged ModelPath
}
// Unwrap exposes the client this one decorates, so grpc.LastDialErrorOf can see
// past it. Without it a staged client answers "the transport was fine" for
// every dial, because embedding grpc.Backend inherits only what Backend
// declares and DialErrorReporter is deliberately not on Backend.
func (f *FileStagingClient) Unwrap() grpc.Backend { return f.Backend }
var _ grpc.BackendUnwrapper = (*FileStagingClient)(nil)
// NewFileStagingClient creates a new file staging wrapper.
func NewFileStagingClient(inner grpc.Backend, stager FileStager, nodeID string) *FileStagingClient {
return &FileStagingClient{
+19 -4
View File
@@ -30,10 +30,15 @@ import (
type InFlightTrackingClient struct {
grpc.ControlBackend // passthrough for control-plane / streaming-constructor methods
inner grpc.InferenceBackend // tracked inference methods delegate here
registry InFlightTracker
nodeID string
modelName string
replicaIndex int
// wrapped is the SAME object as ControlBackend and inner, kept at its full
// type so Unwrap can hand it back. The two fields above are deliberately
// narrowed to the sub-interfaces, which is what gives the compile-time
// guarantee below, and neither of them can be returned as a grpc.Backend.
wrapped grpc.Backend
registry InFlightTracker
nodeID string
modelName string
replicaIndex int
firstOnce sync.Once // guards onFirstComplete
onFirstComplete func() // called once after the first tracked inference call completes
@@ -44,11 +49,21 @@ type InFlightTrackingClient struct {
// InferenceBackend method is left unwrapped.
var _ grpc.Backend = (*InFlightTrackingClient)(nil)
// And it must stay transparent to grpc.LastDialErrorOf. This is the wrapper
// SmartRouter puts on every routed client, so a remote model's cached client is
// one of these; without Unwrap, the transport guard in pkg/model reads nil for
// every model the router produced and evicts on a tunnel blip.
var _ grpc.BackendUnwrapper = (*InFlightTrackingClient)(nil)
// Unwrap exposes the client this one decorates.
func (c *InFlightTrackingClient) Unwrap() grpc.Backend { return c.wrapped }
// NewInFlightTrackingClient wraps a gRPC backend client with in-flight tracking.
func NewInFlightTrackingClient(inner grpc.Backend, registry InFlightTracker, nodeID, modelName string, replicaIndex int) *InFlightTrackingClient {
return &InFlightTrackingClient{
ControlBackend: inner,
inner: inner,
wrapped: inner,
registry: registry,
nodeID: nodeID,
modelName: modelName,
+5 -6
View File
@@ -208,13 +208,12 @@ var ErrNoWorkerDialer = fmt.Errorf("%w: no worker tunnel dialer is configured",
//
// A client that reports nothing (no custom dialer, or a test double) yields
// nil, which means "the call reached a backend" and preserves the behaviour
// every non-distributed caller has always had.
// every non-distributed caller has always had. Decorators are looked through;
// see grpc.BackendUnwrapper for why that is not optional.
func unroutable(client grpc.Backend) error {
reporter, ok := client.(grpc.DialErrorReporter)
if !ok {
return nil
}
dialErr := reporter.LastDialError()
// LastDialErrorOf and not a type assertion: the assertion could not see
// past a decorator, and SmartRouter hands every routed client out wrapped.
dialErr := grpc.LastDialErrorOf(client)
if dialErr == nil {
return nil
}
+55
View File
@@ -1,6 +1,7 @@
package nodes
import (
"errors"
"sync"
"sync/atomic"
"time"
@@ -104,6 +105,60 @@ var _ = Describe("probeCache", func() {
}
})
It("hands every coalesced joiner the leader's REASON, not just its answer", func() {
// The hole this shape exists to close, and the one a closed-over
// variable reintroduces. The reason is written only in the goroutine
// that runs the probe; every caller coalesced into that flight would
// read its own unset variable and see nil. In production that means the
// leader correctly declines to reap a replica on an unreachable worker
// while all seven joiners reap it, on the leader's own observation.
c := newProbeCache(time.Minute)
unreached := errors.New("no route to the worker")
// The probe blocks until every goroutine is inside flight.Do, so the
// joiners are genuinely coalesced rather than serialised. Released by a
// channel, so nothing here waits on a clock.
entered := make(chan struct{})
release := make(chan struct{})
var calls int32
probe := func() (bool, error) {
atomic.AddInt32(&calls, 1)
close(entered)
<-release
return false, unreached
}
const N = 8
start := make(chan struct{})
var wg sync.WaitGroup
reasons := make([]error, N)
alive := make([]bool, N)
for i := 0; i < N; i++ {
wg.Add(1)
go func(i int) {
defer GinkgoRecover()
defer wg.Done()
<-start
alive[i], reasons[i] = c.DoOrCachedResult("k", probe)
}(i)
}
close(start)
// Only unblock the leader once at least one goroutine is inside the
// probe; the rest are then either waiting on the flight or about to be.
<-entered
close(release)
wg.Wait()
Expect(atomic.LoadInt32(&calls)).To(Equal(int32(1)),
"singleflight must collapse %d concurrent probes into one", N)
for i := range reasons {
Expect(alive[i]).To(BeFalse(), "goroutine %d saw a different answer", i)
Expect(reasons[i]).To(MatchError(unreached),
"goroutine %d joined the flight and got the answer without the reason, which is how a joiner reaps what the leader would not", i)
}
})
It("treats different keys independently", func() {
c := newProbeCache(time.Minute)
var aCalls, bCalls int32
@@ -0,0 +1,115 @@
// SPDX-License-Identifier: MIT
package nodes
import (
"context"
"errors"
"net"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/services/cluster"
grpc "github.com/mudler/LocalAI/pkg/grpc"
)
// The reviewer's spec, plus the production shape it was pointing at.
//
// The guard added at the fourth reap site asks the client whether the TRANSPORT
// failed. In production that client is not the one the factory built: SmartRouter
// hands out result.Client, which is an *InFlightTrackingClient, over a
// *FileStagingClient whenever a stager is configured. Both embed grpc.Backend,
// which does not declare LastDialError, so a type assertion on the outermost
// type read nil and the guard was inert for exactly the models the router
// produces. Every spec that constructed a raw client by hand passed anyway.
//
// This is the third time in this task that a correct fix was disarmed by a
// layer further out, which is why the mechanism is now one walker rather than a
// per-caller assertion.
var _ = Describe("the transport answer through the wrappers the router builds", func() {
var (
cause error
raw grpc.Backend
)
BeforeEach(func() {
cause = errors.New("cluster: no route from this replica to that worker")
f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) {
return func(context.Context, string) (net.Conn, error) { return nil, cause }
})
Expect(err).ToNot(HaveOccurred())
raw, err = f.NewClientForNode("X", "10.0.0.1:9001", false)
Expect(err).ToNot(HaveOccurred())
// Provoke one dial so there is something to report.
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
_, _ = raw.HealthCheck(ctx)
})
It("the raw factory client reports, as designed", func() {
Expect(unroutable(raw)).To(MatchError(ErrWorkerUnroutable))
})
It("reports through the in-flight tracker, which is what RouteResult.Client is", func() {
tracked := NewInFlightTrackingClient(raw, &fakeModelRouter{}, "X", "m", 0)
Expect(unroutable(tracked)).To(MatchError(ErrWorkerUnroutable))
})
It("reports through the file staging client, which buildClientForAddr adds", func() {
staged := NewFileStagingClient(raw, nil, "X")
Expect(unroutable(staged)).To(MatchError(ErrWorkerUnroutable))
})
It("reports through BOTH, nested the way production nests them", func() {
// buildClientForAddr wraps in staging, newRouteResult wraps that in
// tracking, model_router puts the result on the cached model, and
// pkg/model's checkIsLoaded asks it. Two layers, and a walker that
// stopped at one would still be wrong here.
nested := NewInFlightTrackingClient(NewFileStagingClient(raw, nil, "X"), &fakeModelRouter{}, "X", "m", 0)
Expect(unroutable(nested)).To(MatchError(ErrWorkerUnroutable))
})
It("still reports nothing through the wrappers when the dial succeeded", func() {
// The other direction, so forwarding cannot pass by always answering
// "unroutable": a backend that genuinely died must still be reapable
// through the same wrappers.
listener, err := net.Listen("tcp", "127.0.0.1:0")
Expect(err).ToNot(HaveOccurred())
DeferCleanup(func() { _ = listener.Close() })
f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) {
var d net.Dialer
return func(ctx context.Context, _ string) (net.Conn, error) {
return d.DialContext(ctx, "tcp", listener.Addr().String())
}
})
Expect(err).ToNot(HaveOccurred())
live, err := f.NewClientForNode("X", "10.0.0.1:9001", false)
Expect(err).ToNot(HaveOccurred())
_, _ = live.HealthCheck(context.Background())
nested := NewInFlightTrackingClient(NewFileStagingClient(live, nil, "X"), &fakeModelRouter{}, "X", "m", 0)
Expect(unroutable(nested)).To(BeNil())
})
It("keeps the cluster condition matchable through the wrappers", func() {
// Not merely "something failed". The five conditions have to survive
// the decorators as well as gRPC, or the consumers are guessing again.
routed := errors.New("x")
f, err := NewTunnelClientFactory("", func(string) func(ctx context.Context, addr string) (net.Conn, error) {
return func(context.Context, string) (net.Conn, error) { return nil, routed }
})
Expect(err).ToNot(HaveOccurred())
c, err := f.NewClientForNode("X", "10.0.0.1:9001", false)
Expect(err).ToNot(HaveOccurred())
routed = cluster.ErrNoRoute
_, _ = c.HealthCheck(context.Background())
nested := NewInFlightTrackingClient(NewFileStagingClient(c, nil, "X"), &fakeModelRouter{}, "X", "m", 0)
got := unroutable(nested)
Expect(got).To(MatchError(cluster.ErrNoRoute))
Expect(got).ToNot(MatchError(cluster.ErrNoConnection))
})
})
+6 -3
View File
@@ -163,16 +163,19 @@ A worker's tunnel lands on exactly one replica, so with N replicas behind a load
The dialling replica states how much time its own client has left in the frame that opens the relayed stream, and the owner bounds its work by the smaller of that and its own 15s ceiling. Neither number can lengthen the other: a patient client cannot park the owning replica, and an impatient one cannot be kept waiting on a budget it did not ask for.
Four outcomes are kept apart on purpose, because they call for different actions:
These outcomes are kept apart on purpose, because they call for different actions:
| Outcome | What it means | What acts on it |
|---|---|---|
| No live owner | No replica holds this worker's tunnel | The worker is treated as absent; its models can be rescheduled |
| No live owner | No replica holds this worker's tunnel | No route right now; the worker's models are **left alone** |
| Not the owner | The routing was stale | Resolve the owner again |
| Peer unreachable | A replica exists and will not answer | Retry |
| No relay path | This replica cannot reach the owner at all | Report; requests here fail until it can |
| The worker refused | The worker answered and said no | Report; the worker is connected |
Only the first is absence. The others are never reported as it, and that is not a stylistic preference: a scheduler told that a connected worker has gone away reclaims every model it is running.
**None of them is absence.** A worker's presence is its **heartbeat**, and a route to it is a separate fact that can be false while the worker is registered, heartbeating and serving every request another replica sends it. So the frontend answers "no route", never "this worker is gone", and nothing on this list causes a model to be rescheduled or a `node_models` row to be deleted.
That distinction is the whole point rather than a nicety. A scheduler told that a connected worker has gone away stops its backend and reclaims every model it is running, and the events that produce "no route" are ordinary ones: a frontend replica restarting, an ownership row a moment stale, a worker that has not dialled its tunnel yet. A worker is treated as absent only when its **heartbeat** goes stale, which is a separate mechanism with its own threshold (see `--stale-node-threshold`).
#### There is no frontend-side fallback, and upgrade order matters
+48
View File
@@ -72,6 +72,54 @@ type DialErrorReporter interface {
LastDialError() error
}
// BackendUnwrapper is implemented by a Backend that DECORATES another one.
//
// Every wrapper in this codebase must implement it, and the reason is a defect
// that shipped: a wrapper embeds the Backend interface, so it inherits every
// declared method and NOTHING else. DialErrorReporter is deliberately not
// declared on Backend, so a wrapped client silently stopped answering "did the
// transport fail" and the guard built on that answer read nil in production
// while passing every spec that constructed a raw client by hand.
//
// Implementing this is what makes a decorator transparent to LastDialErrorOf,
// and it is one line rather than a re-implementation per wrapper, so there is
// no per-wrapper policy to get wrong.
type BackendUnwrapper interface {
Unwrap() Backend
}
// maxBackendUnwrapDepth bounds the walk below. Three wrappers exist today and
// they nest at most two deep; the bound is a guard against a cycle a future
// wrapper could introduce, not a limit anything real approaches.
const maxBackendUnwrapDepth = 16
// LastDialErrorOf reports why the most recent dial under b failed, looking
// THROUGH any decorators, or nil when the dial succeeded or nothing under b has
// a custom transport.
//
// It is the single implementation of that question. Its callers
// (core/services/nodes and pkg/model) each had their own type assertion, and an
// assertion cannot see past a wrapper: in production the client handed to
// pkg/model is an *InFlightTrackingClient over a *FileStagingClient over the
// real one, so both callers were asking a wrapper that had no answer and
// reading nil as "the transport was fine".
func LastDialErrorOf(b Backend) error {
for range maxBackendUnwrapDepth {
if b == nil {
return nil
}
if reporter, ok := b.(DialErrorReporter); ok {
return reporter.LastDialError()
}
wrapper, ok := b.(BackendUnwrapper)
if !ok {
return nil
}
b = wrapper.Unwrap()
}
return nil
}
func buildClient(address string, parallel bool, wd WatchDog, enableWatchDog bool, token string) *Client {
if !enableWatchDog {
wd = nil
+20 -9
View File
@@ -1444,15 +1444,26 @@ func (c *Client) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts ..
// tell them apart, with the original error VALUE intact, so
// core/services/cluster's sentinels survive the trip.
//
// Scope, stated exactly. This is the last dial on this CLIENT, not the last
// dial for a particular RPC. A client used for one probe and closed gives exact
// attribution, which is how every reaping path in core/services/nodes uses it.
// A client shared across concurrent RPCs can attribute a dial failure to the
// wrong one; both directions of that error are safe, because a caller consults
// this only when its RPC already failed, and the outcomes are "treat a dead
// backend as unreachable-for-now" (the row survives one extra round) or "treat
// a transport failure as a backend failure" (the behaviour before this
// existed).
// Scope, stated exactly, including where it is NOT exact.
//
// This is the last dial on this CLIENT, not the last dial for a particular RPC.
// Three of the four callers build a client for one probe and close it, so
// attribution there is exact. The fourth, pkg/model's checkIsLoaded, reads the
// model's long-lived SHARED client and consults this after HealthCheck has
// released opMutex, so a concurrent RPC on the same client can record or clear
// the value inside that window. An earlier version of this comment claimed
// exactness for all four; it was wrong.
//
// The imprecision is accepted there rather than designed away, and the reason
// is which way it can go. A caller consults this only when its own RPC already
// failed, so the two outcomes are: a concurrent dial FAILURE makes a genuinely
// dead backend look unreachable-for-now, and its row survives one extra round
// until the transport recovers; or a concurrent dial SUCCESS clears the value
// and a transport failure reads as a backend failure, which is exactly the
// behaviour that existed before any of this. Neither is a new hazard, and the
// second requires a transport that recovered inside the window. Making it exact
// would mean threading a per-call handle through every Backend method, which is
// a far larger change than the failure it would prevent.
func (c *Client) LastDialError() error {
c.dialErrMu.Lock()
defer c.dialErrMu.Unlock()
+27 -6
View File
@@ -22,6 +22,12 @@ type ConnectionEvictingClient struct {
once sync.Once
}
var _ grpc.BackendUnwrapper = (*ConnectionEvictingClient)(nil)
// Unwrap exposes the client this one decorates, so grpc.LastDialErrorOf can see
// past it.
func (c *ConnectionEvictingClient) Unwrap() grpc.Backend { return c.Backend }
func newConnectionEvictingClient(inner grpc.Backend, modelID string, evict func()) grpc.Backend {
return &ConnectionEvictingClient{
Backend: inner,
@@ -31,13 +37,28 @@ func newConnectionEvictingClient(inner grpc.Backend, modelID string, evict func(
}
func (c *ConnectionEvictingClient) checkErr(err error) {
if err != nil && isConnectionError(err) {
c.once.Do(func() {
xlog.Warn("Connection error during inference, evicting model from cache",
"model", c.modelID, "error", err)
c.evict()
})
if err == nil || !isConnectionError(err) {
return
}
// The fifth site of the same shape, and the one reached during INFERENCE
// rather than a health check. evict() runs ShutdownModel, which for a remote
// model sends backend.stop over NATS to every node holding it and deletes
// every replica row. In distributed mode the client underneath reaches the
// backend over the worker's tunnel, and a failure of THAT transport arrives
// as the same codes.Unavailable a dead backend produces; evicting on it
// stops a model that is loaded and serving, on a worker that is
// heartbeating. A locally spawned backend has no custom transport, so this
// reports nil and the behaviour there is exactly what it always was.
if dialErr := grpc.LastDialErrorOf(c.Backend); dialErr != nil {
xlog.Warn("Inference failed because the worker could not be reached; keeping the model",
"model", c.modelID, "error", dialErr)
return
}
c.once.Do(func() {
xlog.Warn("Connection error during inference, evicting model from cache",
"model", c.modelID, "error", err)
c.evict()
})
}
// --- Intercepted inference methods ---
+12 -5
View File
@@ -695,6 +695,12 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model {
// codes.Unavailable a dead worker produces. Evicting on it would
// unload a model that is loaded and serving. The client records
// which of the two happened; see grpc.DialErrorReporter.
// The client here is long-lived and shared, so this reads the last
// dial on it rather than the one this HealthCheck made; see
// (*grpc.Client).LastDialError for why that imprecision is
// accepted. Both directions of it land on behaviour that already
// existed, and the common case (a worker with no route at all) has
// no concurrent success to clear the value.
if dialErr := transportFailure(client); dialErr != nil {
xlog.Warn("Remote model health check could not reach the worker, keeping cached",
"model", s, "error", dialErr)
@@ -735,9 +741,10 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model {
// backend, where the address IS a socket on this machine and a failed
// connection really does mean the process died.
func transportFailure(client grpc.Backend) error {
reporter, ok := client.(grpc.DialErrorReporter)
if !ok {
return nil
}
return reporter.LastDialError()
// LastDialErrorOf and not a type assertion. The client reaching this
// function for a routed remote model is an *InFlightTrackingClient, often
// over a *FileStagingClient, and an assertion on the outermost type reads
// nil for both: they embed grpc.Backend, which does not declare
// LastDialError. That is exactly how this guard shipped inert.
return grpc.LastDialErrorOf(client)
}
@@ -11,6 +11,7 @@ import (
. "github.com/onsi/gomega"
grpc "github.com/mudler/LocalAI/pkg/grpc"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
"github.com/mudler/LocalAI/pkg/system"
)
@@ -59,3 +60,49 @@ var _ = Describe("the health check on a remote model whose transport failed", fu
Expect(stillThere).To(BeFalse())
})
})
var _ = Describe("the eviction wrapper on a remote model whose transport failed", func() {
// The FIFTH site of the same shape, found by sweeping the decorators rather
// than being named. initializers.go builds this wrapper for exactly the
// remote models the router produces, and its evict callback runs
// ShutdownModel, which sends backend.stop over NATS to every node holding
// the model and deletes every replica row. It fires during INFERENCE, not
// on a health check, so a tunnel blip mid-request was enough.
failingDial := func(cause error) grpc.Backend {
return grpc.NewClientWithDialer("10.0.0.1:9001", false, nil, false, "",
func(context.Context, string) (net.Conn, error) { return nil, cause })
}
It("does not evict when the worker could not be reached", func() {
evicted := 0
client := newConnectionEvictingClient(
failingDial(errors.New("cluster: no route from this replica to that worker")),
"remote-model", func() { evicted++ })
_, err := client.Predict(context.Background(), &pb.PredictOptions{})
Expect(err).To(HaveOccurred())
Expect(evicted).To(BeZero(),
"a worker this frontend cannot route to must not have its backend stopped and its rows deleted")
})
It("still evicts when the worker WAS reached and the connection failed", func() {
// The other direction. No custom dialer, so nothing reports a transport
// failure and a connection error means what it always meant.
evicted := 0
client := newConnectionEvictingClient(
grpc.NewClientWithToken("127.0.0.1:1", false, nil, false, ""),
"dead-model", func() { evicted++ })
_, err := client.Predict(context.Background(), &pb.PredictOptions{})
Expect(err).To(HaveOccurred())
Expect(evicted).To(Equal(1))
})
It("is transparent to the transport question, so a wrapper of it still sees through", func() {
client := newConnectionEvictingClient(
failingDial(errors.New("cluster: no route from this replica to that worker")),
"remote-model", func() {})
_, _ = client.Predict(context.Background(), &pb.PredictOptions{})
Expect(grpc.LastDialErrorOf(client)).ToNot(BeNil())
})
})