Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion backend/cpp/bonsai/Makefile
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@

# Pinned to the HEAD of the `prism` branch on https://github.com/PrismML-Eng/llama.cpp.
# Auto-bumped nightly by .github/workflows/bump_deps.yaml.
BONSAI_VERSION?=79697f23a2c8f3aa2ccb2fd7406095a8dbfbb454
BONSAI_VERSION?=9fcaed763ccda38ea81068ad9d7f991aaddca451
LLAMA_REPO?=https://github.com/PrismML-Eng/llama.cpp

CMAKE_ARGS?=
Expand Down
2 changes: 1 addition & 1 deletion backend/cpp/llama-cpp/Makefile
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@

LLAMA_VERSION?=6bdd77f13cf11b264b4231d320afc404f48d576e
LLAMA_VERSION?=571d0d540df04f25298d0e159e520d9fc62ed121
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp

CMAKE_ARGS?=
Expand Down
2 changes: 1 addition & 1 deletion backend/go/crispasr/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)

# CrispASR version (release tag)
CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR
CRISPASR_VERSION?=c96281d6d409a7f97edbce62c12a6dd2f4da6a92
CRISPASR_VERSION?=1e6f3ad962dc46d86422c3baa4f3c1110d037e4d
SO_TARGET?=libgocrispasr.so

CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
Expand Down
2 changes: 1 addition & 1 deletion backend/go/stablediffusion-ggml/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)

# stablediffusion.cpp (ggml)
STABLEDIFFUSION_GGML_REPO?=https://github.com/leejet/stable-diffusion.cpp
STABLEDIFFUSION_GGML_VERSION?=b2906939774dc73453467215c80390404d0a2701
STABLEDIFFUSION_GGML_VERSION?=ea4e566ccffa10f853ecc3f29e74b1820bc91beb

CMAKE_ARGS+=-DGGML_MAX_NAME=128

Expand Down
10 changes: 9 additions & 1 deletion core/services/messaging/subjects.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
58 changes: 38 additions & 20 deletions core/services/nodes/unloader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
//
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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) {
Expand All @@ -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).
Expand Down
16 changes: 12 additions & 4 deletions core/services/nodes/unloader_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand All @@ -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())
})
})

Expand Down
2 changes: 2 additions & 0 deletions core/services/worker/addr_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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),
)
})

Expand Down
44 changes: 30 additions & 14 deletions core/services/worker/file_staging.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
}
Expand All @@ -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"`
Expand All @@ -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)})
Expand All @@ -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"`
}
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading