From b4d8e23abb9d4d5ba60e910237d44aacc18ae678 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 1 Sep 2026 16:33:59 +0000 Subject: [PATCH] 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 --- core/services/nodes/file_staging_client.go | 8 ++ core/services/nodes/inflight.go | 23 +++- core/services/nodes/interfaces.go | 11 +- core/services/nodes/probe_cache_test.go | 55 +++++++++ core/services/nodes/wrapper_transport_test.go | 115 ++++++++++++++++++ docs/content/features/distributed-mode.md | 9 +- pkg/grpc/backend.go | 48 ++++++++ pkg/grpc/client.go | 29 +++-- pkg/model/connection_evicting_client.go | 33 ++++- pkg/model/loader.go | 17 ++- pkg/model/remote_unroutable_internal_test.go | 47 +++++++ 11 files changed, 362 insertions(+), 33 deletions(-) create mode 100644 core/services/nodes/wrapper_transport_test.go diff --git a/core/services/nodes/file_staging_client.go b/core/services/nodes/file_staging_client.go index bfc202c82..d73a59267 100644 --- a/core/services/nodes/file_staging_client.go +++ b/core/services/nodes/file_staging_client.go @@ -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{ diff --git a/core/services/nodes/inflight.go b/core/services/nodes/inflight.go index 3102a3254..f46bb3e5b 100644 --- a/core/services/nodes/inflight.go +++ b/core/services/nodes/inflight.go @@ -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, diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index d10289838..c8aa23323 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -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 } diff --git a/core/services/nodes/probe_cache_test.go b/core/services/nodes/probe_cache_test.go index 58e6fa111..d28c4c7b7 100644 --- a/core/services/nodes/probe_cache_test.go +++ b/core/services/nodes/probe_cache_test.go @@ -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 diff --git a/core/services/nodes/wrapper_transport_test.go b/core/services/nodes/wrapper_transport_test.go new file mode 100644 index 000000000..ff504929a --- /dev/null +++ b/core/services/nodes/wrapper_transport_test.go @@ -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)) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 4f10775e1..a4700ec4f 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -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 diff --git a/pkg/grpc/backend.go b/pkg/grpc/backend.go index 2bb013de0..1126e39b7 100644 --- a/pkg/grpc/backend.go +++ b/pkg/grpc/backend.go @@ -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 diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index e132086a0..fe4a5d182 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -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() diff --git a/pkg/model/connection_evicting_client.go b/pkg/model/connection_evicting_client.go index 00d42d200..bde2333bf 100644 --- a/pkg/model/connection_evicting_client.go +++ b/pkg/model/connection_evicting_client.go @@ -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 --- diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 1193a18d8..91fdbff67 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -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) } diff --git a/pkg/model/remote_unroutable_internal_test.go b/pkg/model/remote_unroutable_internal_test.go index 25a440d80..06883176f 100644 --- a/pkg/model/remote_unroutable_internal_test.go +++ b/pkg/model/remote_unroutable_internal_test.go @@ -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()) + }) +})