fix(model): make backend shutdown model-scoped (#10865)

Avoid holding the global loader lock across backend lifecycle waits and propagate forced shutdown through distributed workers. Track parallel requests with in-flight counters and reserve worker ports until process termination.

Add focused race tests and an authoritative FizzBee lifecycle model with a fail-closed conformance target.

Assisted-by: Codex:GPT-5 [FizzBee] [Ginkgo]

Signed-off-by: Richard Palethorpe <io@richiejp.com>

Author: richiejp

Makefile

@@ -412,8 +412,13 @@ test-realtime: build-mock-backend
 test-realtime-conformance:
 	GOCMD=$(GOCMD) ./scripts/realtime-conformance.sh
 
+# Verify the shared model-loader shutdown behavior independently of any API
+# modality (focused loader/gRPC/distributed/worker tests under -race + FizzBee).
+test-model-lifecycle-conformance:
+	GOCMD=$(GOCMD) ./scripts/model-lifecycle-conformance.sh
+
 # Install the pinned, checksum-verified FizzBee model checker (into .tools/,
-# gitignored) used by test-realtime-conformance. Idempotent; no-op if present.
+# gitignored) used by the conformance targets. Idempotent; no-op if present.
 install-fizzbee:
 	./scripts/install-fizzbee.sh
 

core/services/messaging/subjects.go

@@ -258,9 +258,17 @@ type NodeBackendInfo struct {
 	Digest  string `json:"digest,omitempty"`
 }
 
+// BackendStopRequest controls worker-side process shutdown. Force skips the
+// best-effort Free RPC so a backend stuck serving a request can still be
+// terminated by the watchdog.
+type BackendStopRequest struct {
+	Backend string `json:"backend"`
+	Force   bool   `json:"force,omitempty"`
+}
+
 // SubjectNodeBackendStop tells a worker node to stop its gRPC backend process.
 // Equivalent to the local deleteProcess(). The node will:
-// 1. Best-effort Free() via gRPC
+// 1. Best-effort bounded Free() via gRPC (unless Force is true)
 // 2. Kill the backend process
 // 3. Can be restarted via another backend.start event.
 func SubjectNodeBackendStop(nodeID string) string {

core/services/nodes/unloader.go

@@ -15,11 +15,6 @@ import (
 	"github.com/mudler/xlog"
 )
 
-// backendStopRequest is the request payload for backend.stop (fire-and-forget).
-type backendStopRequest struct {
-	Backend string `json:"backend"`
-}
-
 // NodeCommandSender abstracts NATS-based commands to worker nodes.
 // Used by HTTP endpoint handlers to avoid coupling to the concrete RemoteUnloaderAdapter.
 //
@@ -78,28 +73,44 @@ func (a *RemoteUnloaderAdapter) InstallTimeout() time.Duration {
 
 // UnloadRemoteModel finds the node(s) hosting the given model and tells them
 // to stop their backend process via NATS backend.stop event.
-// The worker process handles: Free() → kill process.
+// The worker process handles a bounded Free() followed by process termination;
+// forced shutdown skips Free().
 // This is called by ModelLoader.deleteProcess() when process == nil (remote model).
 func (a *RemoteUnloaderAdapter) UnloadRemoteModel(modelName string) error {
-	ctx := context.Background()
+	return a.UnloadRemoteModelContext(context.Background(), modelName, false)
+}
+
+// UnloadRemoteModelContext is the cancellation-aware extension used by the
+// model loader to preserve forced shutdown across the distributed boundary.
+func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error {
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	nodes, err := a.registry.FindNodesWithModel(ctx, modelName)
-	if err != nil || len(nodes) == 0 {
+	if err != nil {
+		return fmt.Errorf("finding nodes with model %q: %w", modelName, err)
+	}
+	if len(nodes) == 0 {
 		xlog.Debug("No remote nodes found with model", "model", modelName)
 		return nil
 	}
 
+	var unloadErr error
 	for _, node := range nodes {
-		xlog.Info("Sending NATS backend.stop to node", "model", modelName, "node", node.Name, "nodeID", node.ID)
-		if err := a.StopBackend(node.ID, modelName); err != nil {
+		xlog.Info("Sending NATS backend.stop to node", "model", modelName, "node", node.Name, "nodeID", node.ID, "force", force)
+		if err := a.stopBackend(node.ID, modelName, force); err != nil {
 			xlog.Warn("Failed to send backend.stop", "node", node.Name, "error", err)
+			unloadErr = errors.Join(unloadErr, fmt.Errorf("stopping model on node %s: %w", node.ID, err))
 			continue
 		}
 		// Remove every replica of this model on the node — the worker will
 		// handle the actual process cleanup.
-		a.registry.RemoveAllNodeModelReplicas(ctx, node.ID, modelName)
+		if err := a.registry.RemoveAllNodeModelReplicas(ctx, node.ID, modelName); err != nil {
+			unloadErr = errors.Join(unloadErr, fmt.Errorf("removing model replicas from node %s: %w", node.ID, err))
+		}
 	}
 
-	return nil
+	return unloadErr
 }
 
 // InstallBackend sends a backend.install request-reply to a worker node.
@@ -142,7 +153,9 @@ func (a *RemoteUnloaderAdapter) InstallBackend(
 	}, a.installTimeout)
 
 	if sub != nil {
-		_ = sub.Unsubscribe()
+		if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil {
+			xlog.Warn("Failed to unsubscribe from backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr)
+		}
 	}
 
 	if err != nil && isNATSTimeout(err) {
@@ -216,7 +229,9 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO
 	}, a.upgradeTimeout)
 
 	if sub != nil {
-		_ = sub.Unsubscribe()
+		if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil {
+			xlog.Warn("Failed to unsubscribe from backend upgrade progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr)
+		}
 	}
 
 	if err != nil && isNATSTimeout(err) {
@@ -250,7 +265,9 @@ func (a *RemoteUnloaderAdapter) installWithForceFallback(nodeID, backendType, ga
 	}, a.upgradeTimeout)
 
 	if sub != nil {
-		_ = sub.Unsubscribe()
+		if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil {
+			xlog.Warn("Failed to unsubscribe from legacy backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr)
+		}
 	}
 
 	if err != nil && isNATSTimeout(err) {
@@ -272,14 +289,15 @@ func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*messaging.BackendL
 // If backend is empty, the worker stops ALL backends.
 // The node stays registered and can receive another InstallBackend later.
 func (a *RemoteUnloaderAdapter) StopBackend(nodeID, backend string) error {
+	return a.stopBackend(nodeID, backend, false)
+}
+
+func (a *RemoteUnloaderAdapter) stopBackend(nodeID, backend string, force bool) error {
 	subject := messaging.SubjectNodeBackendStop(nodeID)
-	if backend == "" {
+	if backend == "" && !force {
 		return a.nats.Publish(subject, nil)
 	}
-	req := struct {
-		Backend string `json:"backend"`
-	}{Backend: backend}
-	return a.nats.Publish(subject, req)
+	return a.nats.Publish(subject, messaging.BackendStopRequest{Backend: backend, Force: force})
 }
 
 // DeleteBackend tells a worker node to delete a backend (stop + remove files).

core/services/nodes/unloader_test.go

@@ -160,14 +160,23 @@ var _ = Describe("RemoteUnloaderAdapter", func() {
 			failOnce := &failOnceMessagingClient{inner: mc, failOn: 0}
 			adapter = NewRemoteUnloaderAdapter(locator, failOnce, 3*time.Minute, 15*time.Minute)
 
-			Expect(adapter.UnloadRemoteModel("llama")).To(Succeed())
+			Expect(adapter.UnloadRemoteModel("llama")).To(HaveOccurred())
 
 			// The second node should still have been processed.
 			// The first node's StopBackend errored, so RemoveNodeModel was NOT called for it.
 			// The second node's StopBackend succeeded, so RemoveNodeModel WAS called.
 			Expect(locator.removedPairs).To(HaveLen(1))
 			Expect(locator.removedPairs[0].nodeID).To(Equal("node-ok"))
 		})
+
+		It("propagates forced shutdown to every worker", func() {
+			locator.nodes = []BackendNode{{ID: "node-1", Name: "worker-1"}}
+			Expect(adapter.UnloadRemoteModelContext(context.Background(), "llama", true)).To(Succeed())
+
+			var payload messaging.BackendStopRequest
+			Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed())
+			Expect(payload).To(Equal(messaging.BackendStopRequest{Backend: "llama", Force: true}))
+		})
 	})
 
 	Describe("StopBackend", func() {
@@ -182,11 +191,10 @@ var _ = Describe("RemoteUnloaderAdapter", func() {
 			Expect(adapter.StopBackend("node-1", "llama-backend")).To(Succeed())
 			Expect(mc.published).To(HaveLen(1))
 
-			var payload struct {
-				Backend string `json:"backend"`
-			}
+			var payload messaging.BackendStopRequest
 			Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed())
 			Expect(payload.Backend).To(Equal("llama-backend"))
+			Expect(payload.Force).To(BeFalse())
 		})
 	})
 

core/services/worker/addr_test.go

@@ -19,7 +19,9 @@ var _ = Describe("Worker address resolution", func() {
 			Entry("falls back to ServeAddr", "", "0.0.0.0:50051", 50051),
 			Entry("returns 50051 when neither set", "", "", 50051),
 			Entry("Addr with custom port", "10.0.0.5:7000", "", 7000),
+			Entry("supports bracketed IPv6", "[2001:db8::1]:7001", "", 7001),
 			Entry("invalid port in Addr falls through to ServeAddr", "host:notanumber", "0.0.0.0:9999", 9999),
+			Entry("out-of-range port in Addr falls through to ServeAddr", "host:70000", "0.0.0.0:9998", 9998),
 		)
 	})
 

core/services/worker/file_staging.go

@@ -59,7 +59,7 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
 	}
 
 	// Subscribe: files.ensure — download S3 key to local, reply with local path
-	natsClient.SubscribeReply(messaging.SubjectNodeFilesEnsure(nodeID), func(data []byte, reply func([]byte)) {
+	if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesEnsure(nodeID), func(data []byte, reply func([]byte)) {
 		var req struct {
 			Key string `json:"key"`
 		}
@@ -77,10 +77,12 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
 
 		xlog.Debug("File ensured locally", "key", req.Key, "path", localPath)
 		replyJSON(reply, map[string]string{"local_path": localPath})
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to files.ensure events: %w", err)
+	}
 
 	// Subscribe: files.stage — upload local path to S3, reply with key
-	natsClient.SubscribeReply(messaging.SubjectNodeFilesStage(nodeID), func(data []byte, reply func([]byte)) {
+	if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesStage(nodeID), func(data []byte, reply func([]byte)) {
 		var req struct {
 			LocalPath string `json:"local_path"`
 			Key       string `json:"key"`
@@ -107,10 +109,12 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
 
 		xlog.Debug("File staged to S3", "path", req.LocalPath, "key", req.Key)
 		replyJSON(reply, map[string]string{"key": req.Key})
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to files.stage events: %w", err)
+	}
 
 	// Subscribe: files.temp — allocate temp file, reply with local path
-	natsClient.SubscribeReply(messaging.SubjectNodeFilesTemp(nodeID), func(data []byte, reply func([]byte)) {
+	if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesTemp(nodeID), func(data []byte, reply func([]byte)) {
 		tmpDir := filepath.Join(cacheDir, "staging-tmp")
 		if err := os.MkdirAll(tmpDir, 0750); err != nil {
 			replyJSON(reply, map[string]string{"error": fmt.Sprintf("creating temp dir: %v", err)})
@@ -123,14 +127,19 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
 			return
 		}
 		localPath := f.Name()
-		f.Close()
+		if err := f.Close(); err != nil {
+			replyJSON(reply, map[string]string{"error": fmt.Sprintf("closing temp file: %v", err)})
+			return
+		}
 
 		xlog.Debug("Allocated temp file", "path", localPath)
 		replyJSON(reply, map[string]string{"local_path": localPath})
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to files.temp events: %w", err)
+	}
 
 	// Subscribe: files.listdir — list files in a local directory, reply with relative paths
-	natsClient.SubscribeReply(messaging.SubjectNodeFilesListDir(nodeID), func(data []byte, reply func([]byte)) {
+	if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesListDir(nodeID), func(data []byte, reply func([]byte)) {
 		var req struct {
 			KeyPrefix string `json:"key_prefix"`
 		}
@@ -163,22 +172,29 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no
 		}
 
 		var files []string
-		filepath.WalkDir(dirPath, func(path string, d os.DirEntry, err error) error {
+		if err := filepath.WalkDir(dirPath, func(path string, d os.DirEntry, err error) error {
 			if err != nil {
-				return nil
+				return err
 			}
 			if !d.IsDir() {
 				rel, err := filepath.Rel(dirPath, path)
-				if err == nil {
-					files = append(files, rel)
+				if err != nil {
+					return err
 				}
+				files = append(files, rel)
 			}
 			return nil
-		})
+		}); err != nil {
+			xlog.Error("Failed to list staged files", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "error", err)
+			replyJSON(reply, map[string]any{"error": err.Error()})
+			return
+		}
 
 		xlog.Debug("Listed remote dir", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "fileCount", len(files))
 		replyJSON(reply, map[string]any{"files": files})
-	})
+	}); err != nil {
+		return fmt.Errorf("subscribing to files.listdir events: %w", err)
+	}
 
 	xlog.Info("Subscribed to file staging NATS subjects", "nodeID", nodeID)
 	return nil

core/services/worker/install.go

@@ -78,7 +78,7 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest,
 			}
 			xlog.Warn("Stale process entry for backend (dead process); cleaning up before reinstall",
 				"backend", req.Backend, "model", req.ModelID, "replica", req.ReplicaIndex, "addr", addr)
-			s.stopBackendExact(processKey)
+			s.stopBackendExact(processKey, false)
 		}
 	} else {
 		// Upgrade path: stop every live process that shares this backend so the
@@ -95,17 +95,18 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest,
 		for _, key := range toStop {
 			xlog.Info("Force install: stopping running backend before reinstall",
 				"backend", req.Backend, "processKey", key)
-			s.stopBackendExact(key)
+			s.stopBackendExact(key, true)
 		}
 	}
 
 	// Parse galleries from request (override local config if provided)
 	galleries := s.galleries
 	if req.BackendGalleries != "" {
 		var reqGalleries []config.Gallery
-		if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err == nil {
-			galleries = reqGalleries
+		if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err != nil {
+			return "", fmt.Errorf("decoding backend galleries: %w", err)
 		}
+		galleries = reqGalleries
 	}
 
 	// When the master tagged this install with an OpID, stream the
@@ -147,7 +148,9 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest,
 			}
 		}
 		// Re-register after install and retry
-		gallery.RegisterBackends(s.systemState, s.ml)
+		if err := gallery.RegisterBackends(s.systemState, s.ml); err != nil {
+			return "", fmt.Errorf("refreshing registered backends after install: %w", err)
+		}
 		backendPath = s.findBackend(req.Backend)
 	}
 
@@ -175,15 +178,16 @@ func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest)
 	for _, key := range toStop {
 		xlog.Info("Upgrade: stopping running backend before reinstall",
 			"backend", req.Backend, "processKey", key)
-		s.stopBackendExact(key)
+		s.stopBackendExact(key, true)
 	}
 
 	galleries := s.galleries
 	if req.BackendGalleries != "" {
 		var reqGalleries []config.Gallery
-		if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err == nil {
-			galleries = reqGalleries
+		if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err != nil {
+			return fmt.Errorf("decoding backend galleries: %w", err)
 		}
+		galleries = reqGalleries
 	}
 
 	// When the master tagged this upgrade with an OpID, stream gallery download
@@ -215,7 +219,9 @@ func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest)
 		}
 	}
 
-	gallery.RegisterBackends(s.systemState, s.ml)
+	if err := gallery.RegisterBackends(s.systemState, s.ml); err != nil {
+		return fmt.Errorf("refreshing registered backends after upgrade: %w", err)
+	}
 	return nil
 }