fix(distributed): run cold model loads as durable jobs instead of holding the advisory lock

core/application/distributed.go

@@ -378,6 +378,9 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
 			cfg.Distributed.BackendInstallTimeoutOrDefault(),
 			cfg.Distributed.ModelLoadTimeoutOrDefault(),
 		),
+		// Bounds the REQUEST, not the load: a caller out of budget gets 503 with
+		// live staging progress while the job keeps running underneath.
+		ModelLoadWait: cfg.Distributed.ModelLoadWait,
 	})
 
 	// Wire staging-progress broadcasting so file-staging shows up on every

core/cli/run.go

@@ -179,6 +179,7 @@ type RunCMD struct {
 	BackendInstallTimeout        string `env:"LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT" help:"NATS round-trip timeout for backend.install requests sent to worker nodes (default 15m). Increase for slow links pulling multi-GB images." group:"distributed"`
 	BackendUpgradeTimeout        string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"NATS round-trip timeout for backend.upgrade requests (default 15m)." group:"distributed"`
 	ModelLoadTimeout             string `env:"LOCALAI_NATS_MODEL_LOAD_TIMEOUT" help:"Fixed gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged. Unset (the default), the deadline is derived from the checkpoint size instead: 5m plus 20s per GiB, capped at 6h, so multi-tens-of-GB diffusion/video checkpoints get the minutes they need without a fixed cliff. Set this only to pin a specific budget; the value is used verbatim, including when it is shorter than the derived one." group:"distributed"`
+	ModelLoadWait                string `env:"LOCALAI_MODEL_LOAD_WAIT" help:"How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with 503, a Retry-After header and live staging progress (default 60s). The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to 0 to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front." group:"distributed"`
 	NatsAccountSeed              string `env:"LOCALAI_NATS_ACCOUNT_SEED" help:"NATS account signing seed (SU...) used to mint per-node worker JWTs at registration" group:"distributed"`
 	NatsServiceJWT               string `env:"LOCALAI_NATS_SERVICE_JWT" help:"NATS user JWT for the frontend (and agent workers) to publish control-plane messages" group:"distributed"`
 	NatsServiceSeed              string `env:"LOCALAI_NATS_SERVICE_SEED" help:"NATS user signing seed (SU...) paired with LOCALAI_NATS_SERVICE_JWT" group:"distributed"`
@@ -386,6 +387,13 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
 		}
 		opts = append(opts, config.WithModelLoadTimeout(d))
 	}
+	if r.ModelLoadWait != "" {
+		d, err := parseDistributedDuration("LOCALAI_MODEL_LOAD_WAIT", r.ModelLoadWait)
+		if err != nil {
+			return err
+		}
+		opts = append(opts, config.WithModelLoadWait(d))
+	}
 	if r.RegistrationToken != "" {
 		opts = append(opts, config.WithRegistrationToken(r.RegistrationToken))
 	}

core/config/distributed_config.go

@@ -82,6 +82,16 @@ type DistributedConfig struct {
 	// pipeline init, which for a multi-tens-of-GB diffusion/video checkpoint on
 	// unified memory can far exceed the 5m default.
 	ModelLoadTimeout time.Duration // gRPC deadline for remote LoadModel (default 5m)
+	// ModelLoadWait bounds how long an inference request waits for a model that
+	// is being cold-loaded before it is answered with 503 plus live progress. A
+	// held HTTP request cannot survive real infrastructure — an ingress or LB
+	// idle timeout kills a 20-minute request regardless of what LocalAI does —
+	// so the wait is bounded by default.
+	//
+	// Zero means unset (DefaultModelLoadWait applies); ModelLoadWaitUnbounded
+	// records the operator asking for unbounded waiting with
+	// LOCALAI_MODEL_LOAD_WAIT=0.
+	ModelLoadWait time.Duration
 
 	MaxUploadSize int64 // Maximum upload body size in bytes (default 50 GB)
 
@@ -315,6 +325,18 @@ func WithModelLoadTimeout(d time.Duration) AppOption {
 	}
 }
 
+// WithModelLoadWait sets how long a request waits for a cold-loading model. A
+// zero d records the operator asking for unbounded waiting: "set the knob to
+// zero" cannot sensibly mean "use the default".
+func WithModelLoadWait(d time.Duration) AppOption {
+	return func(o *ApplicationConfig) {
+		if d == 0 {
+			d = ModelLoadWaitUnbounded
+		}
+		o.Distributed.ModelLoadWait = d
+	}
+}
+
 var EnableAutoApproveNodes = func(o *ApplicationConfig) {
 	o.Distributed.AutoApproveNodes = true
 }
@@ -379,6 +401,7 @@ const (
 	FlagBackendInstallTimeout = "backend-install-timeout"
 	FlagBackendUpgradeTimeout = "backend-upgrade-timeout"
 	FlagModelLoadTimeout      = "model-load-timeout"
+	FlagModelLoadWait         = "model-load-wait"
 	// FlagDiskHeadroomCheck names the disk-headroom toggle. It is quoted in
 	// the warning the check emits while disabled, so the operator reading a
 	// log line knows exactly which knob produced it.
@@ -397,8 +420,19 @@ const (
 	DefaultBackendInstallTimeout = 15 * time.Minute
 	DefaultBackendUpgradeTimeout = 15 * time.Minute
 	DefaultModelLoadTimeout      = 5 * time.Minute
+	// DefaultModelLoadWait is how long a request waits for a cold-loading model
+	// before it is answered with 503 and live progress. Chosen to sit under the
+	// idle timeout of typical ingress/LB defaults, so the answer comes from
+	// LocalAI (with progress the client can act on) rather than from a proxy
+	// dropping the connection.
+	DefaultModelLoadWait = 60 * time.Second
 )
 
+// ModelLoadWaitUnbounded records LOCALAI_MODEL_LOAD_WAIT=0 — "wait as long as
+// it takes" — which a plain zero cannot express, since zero also means "unset,
+// use the default". Only deployments with no proxy in front should use it.
+const ModelLoadWaitUnbounded = -1 * time.Second
+
 // DefaultMaxUploadSize is the default maximum upload body size (50 GB).
 const DefaultMaxUploadSize int64 = 50 << 30
 

core/http/app.go

@@ -78,6 +78,53 @@ func applyModelLoadCooldown(err error, code int, c echo.Context) int {
 	return http.StatusServiceUnavailable
 }
 
+// respondModelLoading answers a request whose model is still cold-loading with
+// 503, a Retry-After header and the live `loading` object, reporting true when
+// it handled the error.
+//
+// The distinction from applyModelLoadCooldown matters: a cooldown means "the
+// last load FAILED, back off", this means "the load is progressing, here is how
+// far it got". Both used to look like the same anonymous error, so an operator
+// watching a 35 GB model stage normally onto a new worker saw only failures.
+func respondModelLoading(err error, c echo.Context) bool {
+	var loadErr *nodes.ModelLoadingError
+	if !errors.As(err, &loadErr) {
+		return false
+	}
+	setModelLoadingRetryAfter(loadErr, c)
+	status := loadErr.Status
+	if jerr := c.JSON(http.StatusServiceUnavailable, schema.ModelLoadingResponse{
+		Error: &schema.APIError{
+			Message: loadErr.Error(),
+			Code:    "model_loading",
+			Type:    "model_loading",
+		},
+		Loading: &status,
+	}); jerr != nil {
+		xlog.Debug("Failed to write model-loading response", "error", jerr)
+	}
+	return true
+}
+
+// applyModelLoading is the body-less half of respondModelLoading, for the
+// opaque-errors handler: status and Retry-After only.
+func applyModelLoading(err error, code int, c echo.Context) int {
+	var loadErr *nodes.ModelLoadingError
+	if !errors.As(err, &loadErr) {
+		return code
+	}
+	setModelLoadingRetryAfter(loadErr, c)
+	return http.StatusServiceUnavailable
+}
+
+func setModelLoadingRetryAfter(loadErr *nodes.ModelLoadingError, c echo.Context) {
+	secs := int(math.Ceil(loadErr.RetryAfter.Seconds()))
+	if secs < 1 {
+		secs = 1
+	}
+	c.Response().Header().Set("Retry-After", strconv.Itoa(secs))
+}
+
 // @title LocalAI API
 // @version 2.0.0
 // @description The LocalAI Rest API.
@@ -141,6 +188,9 @@ func API(application *application.Application) (*echo.Echo, error) {
 	// Set error handler
 	if !application.ApplicationConfig().OpaqueErrors {
 		e.HTTPErrorHandler = func(err error, c echo.Context) {
+			if respondModelLoading(err, c) {
+				return
+			}
 			code := http.StatusInternalServerError
 			var he *echo.HTTPError
 			if errors.As(err, &he) {
@@ -175,6 +225,9 @@ func API(application *application.Application) (*echo.Echo, error) {
 				code = he.Code
 			}
 			code = applyModelLoadCooldown(err, code, c)
+			// Opaque errors deliberately withhold the body, so a still-loading
+			// model gets the status and Retry-After but no progress detail.
+			code = applyModelLoading(err, code, c)
 			c.NoContent(code)
 		}
 	}

core/http/endpoints/localai/model_load_status.go

@@ -0,0 +1,66 @@
+package localai
+
+import (
+	"net/http"
+
+	"github.com/labstack/echo/v4"
+	"github.com/mudler/LocalAI/core/schema"
+	"github.com/mudler/LocalAI/core/services/nodes"
+)
+
+// ModelLoadStatusEndpoint reports the progress of a cold load that is
+// currently running for a model, so a client that got a 503 while the model
+// stages onto a worker can poll rather than blind-retry.
+//
+// Read-only and observability-shaped: it is not admin-gated (a caller allowed
+// to ask for inference on a model may see why it is not answering yet) and it
+// is not feature-gated, since a per-capability gate would make the explanation
+// for a 503 depend on which modality the model happens to be.
+//
+// @Summary Report the progress of an in-flight model load.
+// @Description Returns the live state of a distributed cold load — phase, node, byte progress and ETA — or 404 when no load is running for the model. This is the same `loading` object the 503 response carries while a model is still staging.
+// @Tags models
+// @Produce json
+// @Param id path string true "Model ID"
+// @Success 200 {object} schema.ModelLoadingStatus "Live load progress"
+// @Failure 404 {object} schema.ErrorResponse "No load is running for this model"
+// @Router /api/models/{id}/load-status [get]
+func ModelLoadStatusEndpoint(loadJobs func() nodes.LoadJobStore) echo.HandlerFunc {
+	return func(c echo.Context) error {
+		modelID := c.Param("id")
+		if modelID == "" {
+			return c.JSON(http.StatusBadRequest, schema.ErrorResponse{
+				Error: &schema.APIError{Message: "model id is required", Code: http.StatusBadRequest, Type: "invalid_request_error"},
+			})
+		}
+
+		notLoading := schema.ErrorResponse{
+			Error: &schema.APIError{
+				Message: "no load is running for model " + modelID,
+				Code:    http.StatusNotFound,
+				Type:    "not_found_error",
+			},
+		}
+
+		// Cold-load jobs are a distributed-mode concept: a single-host load is
+		// synchronous and has no job to report on.
+		var store nodes.LoadJobStore
+		if loadJobs != nil {
+			store = loadJobs()
+		}
+		if store == nil {
+			return c.JSON(http.StatusNotFound, notLoading)
+		}
+
+		job, err := store.GetLoadJob(c.Request().Context(), modelID)
+		if err != nil {
+			return c.JSON(http.StatusInternalServerError, schema.ErrorResponse{
+				Error: &schema.APIError{Message: err.Error(), Code: http.StatusInternalServerError, Type: "server_error"},
+			})
+		}
+		if job == nil {
+			return c.JSON(http.StatusNotFound, notLoading)
+		}
+		return c.JSON(http.StatusOK, nodes.LoadingStatus(job))
+	}
+}

core/http/endpoints/localai/model_load_status_test.go

@@ -0,0 +1,75 @@
+package localai_test
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+
+	"github.com/labstack/echo/v4"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/http/endpoints/localai"
+	"github.com/mudler/LocalAI/core/schema"
+	"github.com/mudler/LocalAI/core/services/nodes"
+	"github.com/mudler/LocalAI/core/services/testutil"
+)
+
+var _ = Describe("ModelLoadStatusEndpoint", func() {
+	get := func(store func() nodes.LoadJobStore, modelID string) *httptest.ResponseRecorder {
+		e := echo.New()
+		e.GET("/api/models/:id/load-status", localai.ModelLoadStatusEndpoint(store))
+		req := httptest.NewRequest(http.MethodGet, "/api/models/"+modelID+"/load-status", nil)
+		rec := httptest.NewRecorder()
+		e.ServeHTTP(rec, req)
+		return rec
+	}
+
+	It("404s when the server is not running distributed", func() {
+		rec := get(nil, "some-model")
+		Expect(rec.Code).To(Equal(http.StatusNotFound))
+	})
+
+	Context("with a registry", func() {
+		var registry *nodes.NodeRegistry
+
+		BeforeEach(func() {
+			db := testutil.SetupTestDB()
+			var err error
+			registry, err = nodes.NewNodeRegistry(db)
+			Expect(err).ToNot(HaveOccurred())
+		})
+
+		store := func(r *nodes.NodeRegistry) func() nodes.LoadJobStore {
+			return func() nodes.LoadJobStore { return r }
+		}
+
+		It("404s when no load is running for the model", func() {
+			rec := get(store(registry), "idle-model")
+			Expect(rec.Code).To(Equal(http.StatusNotFound))
+		})
+
+		It("reports the live progress of a running load", func() {
+			ctx := context.Background()
+			_, claimed, err := registry.ClaimLoadJob(ctx, "big-model", "replica-a")
+			Expect(err).ToNot(HaveOccurred())
+			Expect(claimed).To(BeTrue())
+			Expect(registry.UpdateLoadJob(ctx, "big-model", nodes.LoadJobUpdate{
+				State: nodes.LoadJobStateStaging, NodeID: "n1", NodeName: "nvidia-thor",
+				BytesSent: 1000, TotalBytes: 4000, FileIndex: 1, TotalFiles: 1,
+			})).To(Succeed())
+
+			rec := get(store(registry), "big-model")
+			Expect(rec.Code).To(Equal(http.StatusOK))
+
+			var status schema.ModelLoadingStatus
+			Expect(json.Unmarshal(rec.Body.Bytes(), &status)).To(Succeed())
+			Expect(status.Model).To(Equal("big-model"))
+			Expect(status.State).To(Equal(nodes.LoadJobStateStaging))
+			Expect(status.Node).To(Equal("nvidia-thor"))
+			Expect(status.Progress).To(BeNumerically("~", 25, 0.01))
+			Expect(status.TotalBytes).To(Equal(int64(4000)))
+		})
+	})
+})