From 5d57c08c6eaedb8ac4c53897ec9508b3b15f61c6 Mon Sep 17 00:00:00 2001 From: "mudler's LocalAI [bot]" <139863280+localai-bot@users.noreply.github.com> Date: Sat, 25 Jul 2026 16:05:44 +0200 Subject: [PATCH] feat(distributed): cache staged-artifact hashes and publish the model load lifecycle (#11121) feat(distributed): cache staged-artifact hashes and publish load lifecycle Every load request re-hashed every staged artifact on the controller (probeExisting and the upload path both re-read the full file), which for a large multi-file model on NAS-backed storage is minutes of pure re-reading per request even when nothing changed - observed as ~9 minutes of "Upload skipped (file already exists with matching hash)" before every avatar generation. Cache the local hash in the same .sha256 sidecar the worker-side transfer server already maintains, invalidated whenever the sidecar is older than the file. The whole staging+loading phase was also invisible: the NodeModel row was only written after LoadModel succeeded, so /api/nodes and the UI showed nothing while a cold load spent 10+ minutes staging - indistinguishable from nothing happening. Publish the lifecycle instead: "staging" as soon as the node is chosen, "loading" when the checkpoint load starts, and the existing "loaded" on success, with the row removed on any failure so a dead load does not leave a phantom replica. The early row also reserves the replica slot against concurrent schedulers. The nodes view already renders non-loaded states on model chips; style "staging" like "loading". Audited every state-filtered registry/router query: eviction, routing, reconciler and idle-model queries all filter state='loaded' explicitly, so the new transitional rows are visible to observability surfaces but inert to scheduling decisions (except slot occupancy, intentionally). Co-authored-by: Ettore Di Giacinto Co-authored-by: Claude Fable 5 --- .../src/components/nodes/nodeStatus.js | 1 + core/services/nodes/file_stager_http.go | 37 ++++++++++++++++++- core/services/nodes/registry.go | 2 +- core/services/nodes/router.go | 29 ++++++++++++++- core/services/nodes/router_test.go | 11 ++++-- 5 files changed, 73 insertions(+), 7 deletions(-) diff --git a/core/http/react-ui/src/components/nodes/nodeStatus.js b/core/http/react-ui/src/components/nodes/nodeStatus.js index 799a7c13f..3fd61046e 100644 --- a/core/http/react-ui/src/components/nodes/nodeStatus.js +++ b/core/http/react-ui/src/components/nodes/nodeStatus.js @@ -9,6 +9,7 @@ export const statusConfig = { export const modelStateConfig = { loaded: { bg: 'var(--color-success-light)', color: 'var(--color-success)', border: 'var(--color-success-border)' }, + staging: { bg: 'var(--color-primary-light)', color: 'var(--color-primary)', border: 'var(--color-primary-border)' }, loading: { bg: 'var(--color-primary-light)', color: 'var(--color-primary)', border: 'var(--color-primary-border)' }, unloading: { bg: 'var(--color-warning-light)', color: 'var(--color-warning)', border: 'var(--color-warning-border)' }, idle: { bg: 'var(--color-bg-tertiary)', color: 'var(--color-text-muted)', border: 'var(--color-border-subtle)' }, diff --git a/core/services/nodes/file_stager_http.go b/core/services/nodes/file_stager_http.go index 9982a6a13..79047aad6 100644 --- a/core/services/nodes/file_stager_http.go +++ b/core/services/nodes/file_stager_http.go @@ -107,7 +107,7 @@ func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, ke // attempt — the server uses it to detect mid-flight content drift and // reject (409) if a partial upload claims a new identity, forcing a clean // restart. - localHash, err := hashFileWithActivity(ctx, localPath) + localHash, err := hashLocalCached(ctx, localPath) if err != nil { if ctx.Err() != nil { // The cold load was cancelled or expired while hashing. Uploading @@ -473,7 +473,7 @@ func (h *HTTPFileStager) probeExisting(ctx context.Context, addr, localPath, key // legitimately skipped shards accrues no activity at all. observeLoadProgress(ctx) - localHash, err := hashFileWithActivity(ctx, localPath) + localHash, err := hashLocalCached(ctx, localPath) if err != nil { return "", false } @@ -538,6 +538,39 @@ func hashFileWithActivity(ctx context.Context, path string) (string, error) { return hex.EncodeToString(h.Sum(nil)), nil } +// hashLocalCached returns the SHA-256 of path, consulting the same .sha256 +// sidecar the worker-side transfer server maintains (computeAndCacheHash). +// The stager re-hashes every staged artifact on every load request; for a +// large multi-file model on slow storage (a NAS-mounted models dir) that is +// minutes of pure re-reading per request even when nothing changed. The +// sidecar is trusted only when it is not older than the file itself, so a +// rewritten file invalidates it. +func hashLocalCached(ctx context.Context, path string) (string, error) { + sidecar := path + hashSidecarSuffix + + fileStat, err := os.Stat(path) + if err != nil { + return "", err + } + if sidecarStat, err := os.Stat(sidecar); err == nil && !sidecarStat.ModTime().Before(fileStat.ModTime()) { + if data, err := os.ReadFile(sidecar); err == nil { + cached := strings.TrimSpace(string(data)) + if len(cached) == 64 { + return cached, nil + } + } + } + + hashHex, err := hashFileWithActivity(ctx, path) + if err != nil { + return "", err + } + if err := os.WriteFile(sidecar, []byte(hashHex), 0640); err != nil { + xlog.Warn("Failed to write hash sidecar", "path", sidecar, "error", err) + } + return hashHex, nil +} + // progressReader wraps an io.Reader and logs upload progress periodically. // If a StagingProgressCallback is present in the context, it also calls it // for UI-visible progress updates. diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 1f5c42ef7..b9c78ee03 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -128,7 +128,7 @@ type NodeModel struct { ModelName string `gorm:"index;size:255" json:"model_name"` ReplicaIndex int `gorm:"column:replica_index;default:0;index" json:"replica_index"` Address string `gorm:"size:255" json:"address"` // gRPC address for this replica's backend process - State string `gorm:"size:32;default:idle" json:"state"` // loading, loaded, unloading, idle + State string `gorm:"size:32;default:idle" json:"state"` // staging, loading, loaded, unloading, idle InFlight int `json:"in_flight"` // number of active requests on this replica LastUsed time.Time `json:"last_used"` LoadingBy string `gorm:"size:36" json:"loading_by,omitempty"` // frontend ID that triggered loading diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 5b6aed953..62bf302ac 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -319,6 +319,26 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking // that built modelOpts may have no GPU at all in distributed mode. applyNodeHardwareDefaults(modelOpts, node, backendType) + // Publish the load lifecycle from the moment the node is chosen: staging a + // large model takes minutes, and without a registry row the whole phase is + // invisible to /api/nodes and the UI — a cold load looks exactly like + // nothing happening. The row also reserves the replica slot against + // concurrent schedulers. Removed on any failure below so a dead load does + // not leave a phantom replica. + if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "staging", backendAddr, 0); err != nil { + xlog.Warn("Failed to record staging state", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) + } + lifecycleSettled := false + defer func() { + if lifecycleSettled { + return + } + cleanupCtx := context.WithoutCancel(ctx) + if err := r.registry.RemoveNodeModel(cleanupCtx, node.ID, trackingKey, replicaIndex); err != nil { + xlog.Warn("Failed to clear lifecycle row after failed load", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) + } + }() + // Size the remote load budget BEFORE staging: stageModelFiles rewrites the // path fields to their remote equivalents on a clone, and only the local // paths can be stat'ed here. @@ -342,6 +362,11 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr, "payloadBytes", payloadBytes, "loadBudget", loadTimeout) + // Staging is done; the checkpoint load on the worker begins. + if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "loading", backendAddr, 0); err != nil { + xlog.Warn("Failed to record loading state", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) + } + // The cold-load hold above this call extends on STAGING progress, and // the remote LoadModel reports none — so once the last byte lands the // hold expires a stall window later and would cancel a load that is @@ -380,7 +405,9 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking } } - // Record the model as loaded on this node (specific replica slot). + // Record the model as loaded on this node (specific replica slot). From + // here the row is authoritative; the failure-cleanup defer must not touch it. + lifecycleSettled = true if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "loaded", backendAddr, initialInFlight); err != nil { xlog.Warn("Failed to record model on node", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) } diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index 8f0b2f0fd..422822857 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -487,9 +487,14 @@ var _ = Describe("SmartRouter", func() { Expect(result).ToNot(BeNil()) Expect(result.Node.ID).To(Equal("n2")) - // SetNodeModel should record the model as loaded on the node - Expect(reg2.setCalls).To(HaveLen(1)) - Expect(reg2.setCalls[0]).To(ContainSubstring("n2:some-model:loaded")) + // The load lifecycle is published: a staging row appears as soon + // as the node is chosen (what makes a multi-minute cold load + // visible in /api/nodes and the UI), then the final loaded row. + // This path passes nil model options, so the checkpoint-load + // phase (and its "loading" state) is skipped. + Expect(reg2.setCalls).To(HaveLen(2)) + Expect(reg2.setCalls[0]).To(ContainSubstring("n2:some-model:staging")) + Expect(reg2.setCalls[1]).To(ContainSubstring("n2:some-model:loaded")) }) })