feat(distributed): workers no longer need inbound ports (worker tunnel, phase 2)

.golangci.yml

@@ -19,7 +19,22 @@ linters:
     - staticcheck
   enable:
     - forbidigo
+    # gocritic is enabled for ONE checker: ruleguard, which runs the rules in
+    # hack/lint/. Every other gocritic check is off (disable-all below), so
+    # this adds no style noise; it is here purely as the gate that catches a
+    # gRPC backend wrapper written without Unwrap. See
+    # hack/lint/backend_wrappers.go for why that cannot be a compile-time
+    # assertion.
+    - gocritic
   settings:
+    gocritic:
+      disable-all: true
+      enabled-checks:
+        - ruleguard
+      settings:
+        ruleguard:
+          failOn: all
+          rules: '${base-path}/hack/lint/backend_wrappers.go'
     forbidigo:
       forbid:
         - pattern: '^t\.Errorf$'
@@ -126,3 +141,10 @@ linters:
       - path: ^backend/go/whisper/sources/
         text: 'http\.(DefaultClient|Get|Post|PostForm|Head)'
         linters: [forbidigo]
+      # Test doubles embed grpc.Backend to inherit the interface's method set
+      # over a NIL value; they decorate nothing, hold no inner client, and have
+      # no transport answer to forward. The rule targets production wrappers,
+      # which is where swallowing that answer deletes replica rows.
+      # gocritic here is only the backend-wrapper ruleguard rule.
+      - path: _test\.go$
+        linters: [gocritic]

core/application/startup.go

@@ -283,9 +283,15 @@ func New(opts ...config.AppOption) (*Application, error) {
 		// Wire ModelRouter so grpcModel() delegates to SmartRouter in distributed mode
 		application.modelLoader.SetModelRouter(distSvc.ModelAdapter.AsModelRouter())
 		// Wire DistributedModelStore so shutdown/list/watchdog can find remote models
+		// The client factory is not optional here. Without it the store builds
+		// remote models with no client, and pkg/model.Model.GRPC then dials the
+		// worker's raw address with gRPC's own dialer, which is the direct dial
+		// the tunnel replaces; ShutdownModel's Free and the backend monitor's
+		// Status both reach it.
 		distStore := nodes.NewDistributedModelStore(
 			model.NewInMemoryModelStore(),
 			distSvc.Registry,
+			distSvc.BackendClients,
 		)
 		application.modelLoader.SetModelStore(distStore)
 		// Drop the local stub when a model's last replica leaves the registry.

core/application/unadvertised_replica_test.go

@@ -0,0 +1,73 @@
+// SPDX-License-Identifier: MIT
+
+package application
+
+import (
+	"context"
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+// The runtime symptom for a deferral whose cost changed between phases.
+//
+// Not refusing to start without an advertised address stays deferred on
+// purpose: refusing would take out every single-host deployment. What is not
+// deferred is telling the operator, repeatedly, that this replica is invisible
+// and which workers that is costing - because the symptom it produces (a worker
+// that 5xxs from most of the fleet) reads as a worker problem, and a single
+// startup line has scrolled away long before anyone goes looking.
+var _ = Describe("the alarm for a replica with no advertised address", func() {
+	It("keeps firing for as long as the state lasts, and names the workers it costs", func() {
+		// Repetition is the property. A one-shot alarm is the startup line
+		// again, which is what was already there and was not enough.
+		ctx, cancel := context.WithCancel(context.Background())
+		DeferCleanup(cancel)
+
+		alarms := make(chan []string, 8)
+		go nagUnadvertisedReplica(ctx, func() []string { return []string{"w1", "w2"} },
+			time.Millisecond, func(held []string) { alarms <- held })
+
+		// Two, not one: the second is what a one-shot implementation fails.
+		var first, second []string
+		Eventually(alarms, "10s").Should(Receive(&first))
+		Eventually(alarms, "10s").Should(Receive(&second))
+		Expect(first).To(ConsistOf("w1", "w2"),
+			"the workers this is costing are the answer to the question the symptom provokes")
+		Expect(second).To(ConsistOf("w1", "w2"))
+	})
+
+	It("reads the held set on every tick rather than the one it started with", func() {
+		// A replica accumulates tunnels while it runs, so an alarm bound to the
+		// set at startup would name an empty list forever on exactly the
+		// deployment where the cost is real.
+		ctx, cancel := context.WithCancel(context.Background())
+		DeferCleanup(cancel)
+
+		workers := make(chan []string, 32)
+		for range 32 {
+			workers <- []string{"w-late"}
+		}
+		alarms := make(chan []string, 8)
+		go nagUnadvertisedReplica(ctx, func() []string { return <-workers },
+			time.Millisecond, func(held []string) { alarms <- held })
+
+		var got []string
+		Eventually(alarms, "10s").Should(Receive(&got))
+		Expect(got).To(ConsistOf("w-late"))
+	})
+
+	It("stops when the process context ends", func() {
+		ctx, cancel := context.WithCancel(context.Background())
+		stopped := make(chan struct{})
+		go func() {
+			defer GinkgoRecover()
+			nagUnadvertisedReplica(ctx, func() []string { return nil }, time.Hour, func([]string) {})
+			close(stopped)
+		}()
+
+		cancel()
+		Eventually(stopped, "10s").Should(BeClosed())
+	})
+})

core/cli/run.go

@@ -165,6 +165,7 @@ type RunCMD struct {
 	Distributed                  bool   `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" group:"distributed"`
 	InstanceID                   string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"`
 	NatsURL                      string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"`
+	DistributedAdvertiseAddr     string `env:"LOCALAI_DISTRIBUTED_ADVERTISE_ADDR" help:"host:port other frontend replicas dial to reach this one (peer link). Empty = derived from the local address that routes to PostgreSQL, which only works when the database is on another host." group:"distributed"`
 	StorageURL                   string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"`
 	StorageBucket                string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"`
 	StorageRegion                string `env:"LOCALAI_STORAGE_REGION" default:"us-east-1" help:"S3 region" group:"distributed"`
@@ -351,6 +352,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
 	if r.InstanceID != "" {
 		opts = append(opts, config.WithDistributedInstanceID(r.InstanceID))
 	}
+	if r.DistributedAdvertiseAddr != "" {
+		opts = append(opts, config.WithDistributedAdvertiseAddr(r.DistributedAdvertiseAddr))
+	}
 	if r.NatsURL != "" {
 		opts = append(opts, config.WithNatsURL(r.NatsURL))
 	}

core/cli/workerregistry/client.go

@@ -8,7 +8,9 @@ import (
 	"bytes"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"strings"
 	"sync"
@@ -58,9 +60,15 @@ func (c *RegistrationClient) setAuth(req *http.Request) {
 
 // RegisterResponse is the JSON body returned by /api/node/register.
 type RegisterResponse struct {
-	ID           string `json:"id"`
-	Status       string `json:"status,omitempty"` // "pending" until an admin approves the node
-	APIToken     string `json:"api_token,omitempty"`
+	ID       string `json:"id"`
+	Status   string `json:"status,omitempty"` // "pending" until an admin approves the node
+	APIToken string `json:"api_token,omitempty"`
+	// TunnelToken is this node's own credential for GET /api/cluster/connect.
+	// The frontend mints a fresh one on every registration and keeps only its
+	// hash, so this is the ONLY time the plaintext exists anywhere but in this
+	// worker's memory: a worker that discards it cannot get it back without
+	// registering again.
+	TunnelToken  string `json:"tunnel_token,omitempty"`
 	NatsJWT      string `json:"nats_jwt,omitempty"`
 	NatsUserSeed string `json:"nats_user_seed,omitempty"`
 }
@@ -87,7 +95,7 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a
 	defer resp.Body.Close()
 
 	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		return nil, fmt.Errorf("registration failed with status %d", resp.StatusCode)
+		return nil, registrationStatusError(resp)
 	}
 
 	var result RegisterResponse
@@ -97,6 +105,58 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a
 	return &result, nil
 }
 
+// ErrRegistrationRejected marks a registration the frontend REFUSED, as opposed
+// to one it could not answer.
+//
+// Retrying a refusal cannot change it: the request is wrong, or this worker is
+// not allowed to make it. The one that matters in practice is a worker of this
+// release registering against a frontend that predates it, which answers
+// "address is required for backend workers" with 400, because a worker no
+// longer has an address to send. Without this the retry ladder spends four
+// minutes on a verdict the frontend reached instantly, and the operator watches
+// it before being told anything.
+//
+// 408 and 429 are deliberately NOT rejections. Both are the frontend asking for
+// the same request again later, which is exactly what a retry does.
+var ErrRegistrationRejected = errors.New("the frontend refused this registration")
+
+// maxRegistrationErrorBody bounds how much of a refusal's body is quoted back.
+// Enough for a message, not enough for an HTML error page to bury the log line
+// it is meant to explain.
+const maxRegistrationErrorBody = 512
+
+// registrationStatusError turns a non-2xx response into an error that says WHY.
+//
+// The body is the point. The frontend explains its refusals there
+// ("address is required for backend workers", "invalid registration token"),
+// and discarding it left an operator with a bare status code: the one line that
+// would tell them which of several possible mistakes they made was read off the
+// socket and thrown away.
+func registrationStatusError(resp *http.Response) error {
+	detail, err := io.ReadAll(io.LimitReader(resp.Body, maxRegistrationErrorBody))
+	if err != nil {
+		xlog.Debug("Could not read the frontend's registration error body", "status", resp.StatusCode, "error", err)
+	}
+	msg := strings.Join(strings.Fields(string(detail)), " ")
+	base := fmt.Sprintf("registration failed with status %d", resp.StatusCode)
+	if msg != "" {
+		base = fmt.Sprintf("%s: %s", base, msg)
+	}
+	if isRegistrationRejection(resp.StatusCode) {
+		return fmt.Errorf("%s: %w", base, ErrRegistrationRejected)
+	}
+	return errors.New(base)
+}
+
+// isRegistrationRejection reports whether a status is a verdict rather than a
+// condition that may pass.
+func isRegistrationRejection(status int) bool {
+	if status == http.StatusRequestTimeout || status == http.StatusTooManyRequests {
+		return false
+	}
+	return status >= 400 && status < 500
+}
+
 // Register sends a single registration request and returns the node ID and
 // optional credentials (API token for agent workers, NATS JWT when configured).
 func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) (nodeID, apiToken, natsJWT, natsSeed string, err error) {
@@ -108,27 +168,48 @@ func (c *RegistrationClient) Register(ctx context.Context, body map[string]any)
 }
 
 // RegisterWithRetry retries registration with exponential backoff.
+//
+// It drops every field of the response it does not name, the tunnel credential
+// among them. Callers that need one use RegisterFullWithRetry.
 func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken, natsJWT, natsSeed string, err error) {
+	res, err := c.RegisterFullWithRetry(ctx, body, maxRetries)
+	if err != nil {
+		return "", "", "", "", err
+	}
+	return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil
+}
+
+// RegisterFullWithRetry retries registration with exponential backoff and
+// returns the whole response.
+func (c *RegistrationClient) RegisterFullWithRetry(ctx context.Context, body map[string]any, maxRetries int) (*RegisterResponse, error) {
 	backoff := 2 * time.Second
 	maxBackoff := 30 * time.Second
 
+	var err error
 	for attempt := 1; attempt <= maxRetries; attempt++ {
-		nodeID, apiToken, natsJWT, natsSeed, err = c.Register(ctx, body)
+		var res *RegisterResponse
+		res, err = c.RegisterFull(ctx, body)
 		if err == nil {
-			return nodeID, apiToken, natsJWT, natsSeed, nil
+			return res, nil
+		}
+		if errors.Is(err, ErrRegistrationRejected) {
+			// A verdict, not an outage. Reported on the first attempt so the
+			// reason the frontend gave is the first thing in the log rather
+			// than the last, after the ladder.
+			return nil, err
 		}
 		if attempt == maxRetries {
-			return "", "", "", "", fmt.Errorf("failed after %d attempts: %w", maxRetries, err)
+			return nil, fmt.Errorf("failed after %d attempts: %w", maxRetries, err)
 		}
 		xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err)
 		select {
 		case <-ctx.Done():
-			return "", "", "", "", ctx.Err()
+			return nil, ctx.Err()
 		case <-time.After(backoff):
 		}
 		backoff = min(backoff*2, maxBackoff)
 	}
-	return nodeID, apiToken, natsJWT, natsSeed, err
+	return nil, err
 }
 
 // Heartbeat sends a single heartbeat POST with the given body.