diff --git a/core/application/distributed.go b/core/application/distributed.go index 32269a44a..e22ac5534 100644 --- a/core/application/distributed.go +++ b/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 diff --git a/core/cli/run.go b/core/cli/run.go index 7d35fb693..53c8cbf47 100644 --- a/core/cli/run.go +++ b/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)) } diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 3cbbcf77b..5a48a84e9 100644 --- a/core/config/distributed_config.go +++ b/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 diff --git a/core/http/app.go b/core/http/app.go index 8b08af4ad..65e8ee917 100644 --- a/core/http/app.go +++ b/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) } } diff --git a/core/http/endpoints/localai/model_load_status.go b/core/http/endpoints/localai/model_load_status.go new file mode 100644 index 000000000..70cebdcaa --- /dev/null +++ b/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)) + } +} diff --git a/core/http/endpoints/localai/model_load_status_test.go b/core/http/endpoints/localai/model_load_status_test.go new file mode 100644 index 000000000..9f989a548 --- /dev/null +++ b/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))) + }) + }) +}) diff --git a/core/http/react-ui/e2e/chat-model-loading.spec.js b/core/http/react-ui/e2e/chat-model-loading.spec.js new file mode 100644 index 000000000..e8e564d21 --- /dev/null +++ b/core/http/react-ui/e2e/chat-model-loading.spec.js @@ -0,0 +1,87 @@ +import { test, expect } from './coverage-fixtures.js' + +// A model that is cold-loading onto a worker answers 503 with a `loading` +// object rather than an error. The chat must show that as progress and retry +// itself once the load finishes — the whole point of the change: an operator +// watching a 35 GB model stage normally used to see only repeated failures. +test.describe('Chat - model loading', () => { + test('renders staging progress on 503 and retries when the model is ready', async ({ page }) => { + await page.route('**/api/models/capabilities', (route) => { + route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ data: [{ id: 'test-model', capabilities: ['FLAG_CHAT'] }] }), + }) + }) + + let completions = 0 + await page.route('**/v1/chat/completions', (route) => { + completions++ + if (completions === 1) { + route.fulfill({ + status: 503, + contentType: 'application/json', + headers: { 'Retry-After': '5' }, + body: JSON.stringify({ + error: { + message: 'model test-model is staging on node nvidia-thor (41%)', + type: 'model_loading', + code: 'model_loading', + }, + loading: { + model: 'test-model', + state: 'staging', + node: 'nvidia-thor', + progress: 41.2, + bytes_sent: 14730000000, + total_bytes: 35776484480, + file_index: 1, + total_files: 2, + eta_seconds: 660, + }, + }), + }) + return + } + route.fulfill({ + status: 200, + contentType: 'text/event-stream', + body: + 'data: {"choices":[{"delta":{"content":"loaded at last"}}]}\n\n' + + 'data: [DONE]\n\n', + }) + }) + + // First poll still staging, then the job is gone: the model is ready. + let polls = 0 + await page.route('**/api/models/*/load-status', (route) => { + polls++ + if (polls === 1) { + route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ + model: 'test-model', state: 'staging', node: 'nvidia-thor', + progress: 62.5, bytes_sent: 22000000000, total_bytes: 35776484480, + file_index: 2, total_files: 2, eta_seconds: 300, + }), + }) + return + } + route.fulfill({ status: 404, contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/chat') + await expect(page.getByRole('button', { name: 'test-model' })).toBeVisible({ timeout: 10_000 }) + + await page.locator('.chat-input').fill('Hello') + await page.locator('.chat-send-btn').click() + + // The 503 is shown as progress, not as an error. + await expect(page.locator('.chat-staging-progress')).toBeVisible({ timeout: 10_000 }) + await expect(page.locator('.chat-staging-label')).toContainText('nvidia-thor') + await expect(page.locator('.chat-staging-pct')).toContainText('41%') + + // ...and the request retries itself once the load reports ready. + await expect(page.getByText('loaded at last')).toBeVisible({ timeout: 25_000 }) + expect(completions).toBeGreaterThan(1) + }) +}) diff --git a/core/http/react-ui/public/locales/en/chat.json b/core/http/react-ui/public/locales/en/chat.json index 0e7af136b..f004cee56 100644 --- a/core/http/react-ui/public/locales/en/chat.json +++ b/core/http/react-ui/public/locales/en/chat.json @@ -81,7 +81,15 @@ }, "streaming": { "transferring": "Transferring model...", - "transferringTo": "Transferring model to {{node}}..." + "transferringTo": "Transferring model to {{node}}...", + "onNode": "on {{node}}", + "eta": "~{{value}} left", + "modelState": { + "pending": "Preparing to load the model...", + "installing": "Installing the backend...", + "staging": "Staging model files...", + "loading": "Loading the model..." + } }, "tokens": { "perSec": "{{count}} tok/s", diff --git a/core/http/react-ui/src/hooks/useChat.js b/core/http/react-ui/src/hooks/useChat.js index 2bcf955ce..567cf7f55 100644 --- a/core/http/react-ui/src/hooks/useChat.js +++ b/core/http/react-ui/src/hooks/useChat.js @@ -16,6 +16,68 @@ async function extractHttpError(response) { return errorMsg } +// How long the UI keeps waiting for a model that is staging onto a worker. +// Staging a multi-GB checkpoint runs for tens of minutes — far longer than the +// server's own per-request wait budget — so the poll has to outlive several +// 503s. Bounded, so a load that never finishes eventually surfaces as an error +// rather than as a spinner nobody questions. +const MODEL_LOAD_POLL_INTERVAL = 3000 +const MODEL_LOAD_MAX_ATTEMPTS = 3 +const MODEL_LOAD_MAX_POLLS = 600 // ~30 min per attempt + +// readModelLoading returns the `loading` object from the 503 the server sends +// while a model is still cold-loading, or null for any other failure. It reads +// a clone so the caller can still parse the body for its error message. +async function readModelLoading(response) { + if (response.status !== 503) return null + try { + const data = await response.clone().json() + if (data?.error?.type !== 'model_loading') return null + return { ...(data.loading || {}), message: data.error?.message } + } catch { + return null + } +} + +// waitForModelReady polls load-status until the model finishes loading (404 — +// no job left), the load fails, or the caller aborts. Returns true when it is +// worth re-sending the request. +async function waitForModelReady(modelID, onProgress, signal) { + for (let i = 0; i < MODEL_LOAD_MAX_POLLS; i++) { + await new Promise(resolve => setTimeout(resolve, MODEL_LOAD_POLL_INTERVAL)) + if (signal?.aborted) return false + try { + const res = await fetch(apiUrl(API_CONFIG.endpoints.modelLoadStatus(modelID)), { signal }) + if (res.status === 404) return true // job gone: loaded, or worth one retry + if (!res.ok) return false + const status = await res.json() + if (status?.state === 'failed') return false + onProgress(status) + } catch { + return false + } + } + return false +} + +// fetchWithModelLoadWait issues the request and, when the model is still +// staging onto a worker, waits for it rather than surfacing an error. The +// server answers 503 within its own wait budget so no connection is held for +// the whole load; the UI picks the wait back up here and retries once ready. +async function fetchWithModelLoadWait(url, init, modelID, onLoading, signal) { + let response = await fetch(url, init) + for (let attempt = 0; attempt < MODEL_LOAD_MAX_ATTEMPTS; attempt++) { + const loading = await readModelLoading(response) + if (!loading) return response + onLoading(loading) + const ready = await waitForModelReady(loading.model || modelID, onLoading, signal) + onLoading(null) + if (!ready) return response + response = await fetch(url, init) + } + return response +} + function extractThinking(text) { let regularContent = '' let thinkingContent = '' @@ -125,6 +187,10 @@ export function useChat(initialModel = '') { const [streamingToolCalls, setStreamingToolCalls] = useState([]) const [tokensPerSecond, setTokensPerSecond] = useState(null) const [maxTokensPerSecond, setMaxTokensPerSecond] = useState(null) + // Live progress of a cold load the request is waiting on, so the composer can + // say "staging to nvidia-thor, 41%" instead of showing an error for a model + // that is loading exactly as it should. + const [modelLoading, setModelLoading] = useState(null) const abortControllerRef = useRef(null) const startTimeRef = useRef(null) const tokenCountRef = useRef(0) @@ -360,12 +426,12 @@ export function useChat(initialModel = '') { // Legacy MCP SSE streaming (custom event types from /v1/mcp/chat/completions) try { const timeoutId = setTimeout(() => controller.abort(), 300000) // 5 min timeout - const response = await fetch(apiUrl(endpoint), { + const response = await fetchWithModelLoadWait(apiUrl(endpoint), { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(requestBody), signal: controller.signal, - }) + }, requestBody.model, setModelLoading, controller.signal) clearTimeout(timeoutId) if (!response.ok) { @@ -506,12 +572,12 @@ export function useChat(initialModel = '') { let fullToolCalls = [] // Tool calls with id for agentic loop try { - const response = await fetch(apiUrl(endpoint), { + const response = await fetchWithModelLoadWait(apiUrl(endpoint), { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(loopBody), signal: controller.signal, - }) + }, loopBody.model, setModelLoading, controller.signal) if (!response.ok) { throw new Error(await extractHttpError(response)) @@ -746,6 +812,7 @@ export function useChat(initialModel = '') { // Finalize setIsStreaming(false) setStreamingChatId(null) + setModelLoading(null) abortControllerRef.current = null setStreamingContent('') setStreamingReasoning('') @@ -821,6 +888,7 @@ export function useChat(initialModel = '') { streamingToolCalls: isActiveStreaming ? streamingToolCalls : [], tokensPerSecond, maxTokensPerSecond, + modelLoading: isActiveStreaming ? modelLoading : null, addChat, forkChat, switchChat, diff --git a/core/http/react-ui/src/pages/Chat.jsx b/core/http/react-ui/src/pages/Chat.jsx index 55c6b672e..5866a2514 100644 --- a/core/http/react-ui/src/pages/Chat.jsx +++ b/core/http/react-ui/src/pages/Chat.jsx @@ -294,6 +294,17 @@ function editableMessageText(message) { return typeof textBlock?.text === 'string' ? textBlock.text : null } +// formatLoadEta renders the server's remaining-seconds estimate. The server +// omits it entirely until its observed transfer rate is meaningful, so anything +// arriving here is worth showing. +function formatLoadEta(seconds) { + if (!Number.isFinite(seconds) || seconds <= 0) return '' + if (seconds < 60) return `${Math.round(seconds)}s` + const minutes = Math.round(seconds / 60) + if (minutes < 60) return `${minutes} min` + return `${Math.floor(minutes / 60)}h ${minutes % 60}m` +} + function withEditedMessageText(message, text) { if (typeof message.content === 'string') return { ...message, content: text } const textIndex = message.content.findIndex(block => block?.type === 'text') @@ -315,7 +326,7 @@ export default function Chat() { const { operations } = useOperations() const { chats, activeChat, activeChatId, isStreaming, streamingChatId, streamingContent, - streamingReasoning, streamingToolCalls, tokensPerSecond, maxTokensPerSecond, + streamingReasoning, streamingToolCalls, tokensPerSecond, maxTokensPerSecond, modelLoading, addChat, forkChat, switchChat, deleteChat, deleteAllChats, renameChat, updateChatSettings, sendMessage, stopGeneration, clearHistory, getContextUsagePercent, addMessage, } = useChat(urlModel || '') @@ -326,6 +337,34 @@ export default function Chat() { return operations.find(op => op.taskType === 'staging' && op.name === activeChat.model) || null }, [operations, isStreaming, activeChat?.model]) + // What to show instead of the thinking dots while the model is not up yet. + // The load job wins over the staging operation: it is authoritative across + // frontend replicas and names the phase (installing / staging / loading), + // where the operation only knows about a byte transfer this replica is + // performing. The operation stays as the fallback for a transfer with no job + // attached to this request (a reconciler scale-up, for instance). + const loadProgress = useMemo(() => { + if (modelLoading) { + const eta = formatLoadEta(modelLoading.eta_seconds) + return { + label: t(`streaming.modelState.${modelLoading.state}`, t('streaming.transferring')) + + (modelLoading.node ? ` ${t('streaming.onNode', { node: modelLoading.node })}` : ''), + progress: modelLoading.progress || 0, + detail: eta ? t('streaming.eta', { value: eta }) : '', + } + } + if (stagingOp) { + return { + label: stagingOp.nodeName + ? t('streaming.transferringTo', { node: stagingOp.nodeName }) + : t('streaming.transferring'), + progress: stagingOp.progress || 0, + detail: stagingOp.message || '', + } + } + return null + }, [modelLoading, stagingOp, t]) + const [input, setInput] = useState('') const [files, setFiles] = useState([]) const [showSettings, setShowSettings] = useState(false) @@ -1347,21 +1386,21 @@ export default function Chat() {
- {stagingOp ? ( + {loadProgress ? (
- {stagingOp.nodeName ? t('streaming.transferringTo', { node: stagingOp.nodeName }) : t('streaming.transferring')} + {loadProgress.label}
- {stagingOp.progress > 0 && ( + {loadProgress.progress > 0 && (
-
+
- {Math.round(stagingOp.progress)}% + {Math.round(loadProgress.progress)}%
)} - {stagingOp.message && ( -
{stagingOp.message}
+ {loadProgress.detail && ( +
{loadProgress.detail}
)}
) : ( diff --git a/core/http/react-ui/src/utils/config.js b/core/http/react-ui/src/utils/config.js index 80d437c91..c72e7e91b 100644 --- a/core/http/react-ui/src/utils/config.js +++ b/core/http/react-ui/src/utils/config.js @@ -66,6 +66,9 @@ export const API_CONFIG = { cancelAgentJob: (id) => `/api/agent/jobs/${id}/cancel`, executeAgentJob: '/api/agent/jobs/execute', + // Progress of a cold load still staging a model onto a worker + modelLoadStatus: (id) => `/api/models/${encodeURIComponent(id)}/load-status`, + // OpenAI-compatible endpoints chatCompletions: '/v1/chat/completions', mcpChatCompletions: '/v1/mcp/chat/completions', diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index bfed2d9ed..30a2803f7 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -11,6 +11,7 @@ import ( "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/monitoring" + "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/templates" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/model" @@ -154,6 +155,18 @@ func RegisterLocalAIRoutes(router *echo.Echo, // Forget does not load a voice model — it only needs the registry. router.POST("/v1/voice/forget", localai.VoiceForgetEndpoint(app.VoiceRegistry())) + // Progress of an in-flight cold load. Standard auth only: it explains a 503 + // the caller just received, so gating it behind admin (or a per-modality + // feature) would hide the explanation from exactly the client that needs it. + // Resolved per request, not at registration: distributed services are wired + // during startup and a snapshot taken here could be nil forever. + router.GET("/api/models/:id/load-status", localai.ModelLoadStatusEndpoint(func() nodes.LoadJobStore { + if d := app.Distributed(); d != nil && d.Registry != nil { + return d.Registry + } + return nil + })) + voiceProfiles := app.VoiceProfileStore() router.GET("/api/voice-profiles", localai.ListVoiceProfilesEndpoint(voiceProfiles)) router.GET("/api/voice-profiles/:id/audio", localai.ServeVoiceProfileAudioEndpoint(voiceProfiles)) @@ -309,6 +322,7 @@ func RegisterLocalAIRoutes(router *echo.Echo, "config_patch": "/api/models/config-json/:name", "autocomplete": "/api/models/config-metadata/autocomplete/:provider", "vram_estimate": "/api/models/vram-estimate", + "model_load_status": "/api/models/:id/load-status", "tts": "/tts", "voice_profiles": "/api/voice-profiles", "transcription": "/v1/audio/transcriptions", @@ -344,6 +358,7 @@ func RegisterLocalAIRoutes(router *echo.Echo, "import": "/models/import", "reload": "/models/reload", "list_aliases": "/api/aliases", + "load_status": "/api/models/:id/load-status", }, "ai_functions": map[string]string{ "tts": "/tts", diff --git a/core/schema/model_loading.go b/core/schema/model_loading.go new file mode 100644 index 000000000..3e2927508 --- /dev/null +++ b/core/schema/model_loading.go @@ -0,0 +1,28 @@ +package schema + +// ModelLoadingStatus describes a cold load that is still in progress. In +// distributed mode a model can take tens of minutes to stage onto a worker, +// which is far longer than a request may be held; a caller that runs out of +// wait budget gets this instead of an anonymous hang or a misleading error. +type ModelLoadingStatus struct { + Model string `json:"model"` + State string `json:"state"` + Node string `json:"node,omitempty"` + Progress float64 `json:"progress"` + BytesSent int64 `json:"bytes_sent"` + TotalBytes int64 `json:"total_bytes"` + FileIndex int `json:"file_index"` + TotalFiles int `json:"total_files"` + // ETASeconds is omitted rather than guessed until enough bytes have moved + // for the observed rate to mean anything. A confidently wrong ETA on a + // twenty-minute wait is worse than none. + ETASeconds int `json:"eta_seconds,omitempty"` +} + +// ModelLoadingResponse is the 503 body served while a model is still loading. +// The `error` envelope keeps OpenAI-client compatibility; `loading` is additive, +// so existing clients ignore it and load-aware ones can render real progress. +type ModelLoadingResponse struct { + Error *APIError `json:"error,omitempty"` + Loading *ModelLoadingStatus `json:"loading,omitempty"` +} diff --git a/core/services/advisorylock/advisorylock.go b/core/services/advisorylock/advisorylock.go index a37403c31..ccaea2d96 100644 --- a/core/services/advisorylock/advisorylock.go +++ b/core/services/advisorylock/advisorylock.go @@ -180,6 +180,16 @@ func WithLockCtx(ctx context.Context, db *gorm.DB, key int64, fn func() error) e // Restore the session default before this pooled connection is reused. defer func() { _, _ = conn.ExecContext(context.Background(), "RESET lock_timeout") }() + // statement_timeout aborts the same blocking pg_advisory_lock() call + // independently of lock_timeout, with SQLSTATE 57014. Deployments that set a + // short statement_timeout on the role (60s is a common default) would + // otherwise kill every waiter regardless of the lock_timeout override above. + if _, err := conn.ExecContext(ctx, + fmt.Sprintf("SET statement_timeout = %d", waitBudget.Milliseconds())); err != nil { + return fmt.Errorf("advisorylock: setting statement_timeout: %w", err) + } + defer func() { _, _ = conn.ExecContext(context.Background(), "RESET statement_timeout") }() + if _, err := conn.ExecContext(ctx, "SELECT pg_advisory_lock($1)", key); err != nil { return fmt.Errorf("advisorylock: acquiring lock %d: %w", key, err) } diff --git a/core/services/advisorylock/advisorylock_test.go b/core/services/advisorylock/advisorylock_test.go index a77df56c4..f1bd3e75e 100644 --- a/core/services/advisorylock/advisorylock_test.go +++ b/core/services/advisorylock/advisorylock_test.go @@ -205,6 +205,54 @@ var _ = Describe("AdvisoryLock", func() { <-released }) + It("waits out a short server-side statement_timeout instead of failing with 57014", func() { + const lockKey int64 = 705 + + // Same shape as the lock_timeout case above, but for the *other* + // server-side bound that aborts the very same blocking + // pg_advisory_lock() statement. Production roles routinely carry + // statement_timeout=60s; a cold model load holds the lock far longer, + // so every concurrent caller died with SQLSTATE 57014 ("canceling + // statement due to statement timeout") rather than waiting its turn. + Expect(db.Exec("ALTER DATABASE testdb SET statement_timeout = '300ms'").Error).ToNot(HaveOccurred()) + sqlDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + // Drop pooled connections so subsequent ones reconnect and inherit + // the new database-level statement_timeout default. + sqlDB.SetMaxIdleConns(0) + + holding := make(chan struct{}) + released := make(chan struct{}) + go func() { + defer GinkgoRecover() + herr := WithLockCtx(context.Background(), db, lockKey, func() error { + close(holding) + // Hold well past the 300ms server statement_timeout. + time.Sleep(1 * time.Second) + return nil + }) + Expect(herr).ToNot(HaveOccurred()) + close(released) + }() + + <-holding // ensure the holder owns the lock before we contend + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + executed := false + start := time.Now() + werr := WithLockCtx(ctx, db, lockKey, func() error { + executed = true + return nil + }) + Expect(werr).ToNot(HaveOccurred(), + "waiter should wait out the in-progress hold, not fail with statement_timeout (57014)") + Expect(executed).To(BeTrue()) + Expect(time.Since(start)).To(BeNumerically(">=", 400*time.Millisecond), + "waiter should have actually waited for the holder to release") + <-released + }) + It("bounds a deadline-less waiter with the backstop instead of waiting forever", func() { const lockKey int64 = 704 diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 399ad74f8..1be5f20e8 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -39,6 +39,18 @@ type ModelRouter interface { GetNodeLabels(ctx context.Context, nodeID string) ([]NodeLabel, error) FindNodesWithModel(ctx context.Context, modelName string) ([]BackendNode, error) LoadedReplicaStats(ctx context.Context, modelName string, candidateNodeIDs []string) ([]ReplicaCandidate, error) + LoadJobStore +} + +// LoadJobStore is the durable cold-load job record SmartRouter uses to +// de-duplicate concurrent loaders across replicas without holding the per-model +// advisory lock for the whole load. See ModelLoadJob. +type LoadJobStore interface { + ClaimLoadJob(ctx context.Context, trackingKey, owner string) (*ModelLoadJob, bool, error) + GetLoadJob(ctx context.Context, trackingKey string) (*ModelLoadJob, error) + UpdateLoadJob(ctx context.Context, trackingKey string, u LoadJobUpdate) error + FailLoadJob(ctx context.Context, trackingKey, msg string) error + DeleteLoadJob(ctx context.Context, trackingKey string) error } // ConcurrencyConflictResolver returns the names of configured models that diff --git a/core/services/nodes/load_job_phase.go b/core/services/nodes/load_job_phase.go new file mode 100644 index 000000000..ece0d099d --- /dev/null +++ b/core/services/nodes/load_job_phase.go @@ -0,0 +1,64 @@ +package nodes + +import ( + "context" + "sync" +) + +// loadPhaseReporter carries the current phase of a cold load from the code that +// performs it back to the job runner's heartbeat, without threading a job +// handle through every scheduling function. +// +// It rides on the context the same way the cold-load deadline does (see +// load_deadline.go), so the single-host paths and every test that constructs a +// router directly stay untouched: with no reporter on the context, the report +// calls are no-ops. +type loadPhaseReporter struct { + mu sync.Mutex + state string + nodeID string + nodeName string + replicaIndex int +} + +type loadPhaseKey struct{} + +func newLoadPhaseReporter() *loadPhaseReporter { + return &loadPhaseReporter{state: LoadJobStatePending} +} + +func withLoadPhaseReporter(ctx context.Context, p *loadPhaseReporter) context.Context { + return context.WithValue(ctx, loadPhaseKey{}, p) +} + +// snapshot returns the phase as a job update. Byte counts are filled in by the +// caller from the staging tracker. +func (p *loadPhaseReporter) snapshot() LoadJobUpdate { + p.mu.Lock() + defer p.mu.Unlock() + return LoadJobUpdate{ + State: p.state, + NodeID: p.nodeID, + NodeName: p.nodeName, + ReplicaIndex: p.replicaIndex, + } +} + +func (p *loadPhaseReporter) set(state string, node *BackendNode, replicaIndex int) { + p.mu.Lock() + defer p.mu.Unlock() + p.state = state + if node != nil { + p.nodeID, p.nodeName, p.replicaIndex = node.ID, node.Name, replicaIndex + } +} + +// reportLoadPhase records which phase of a cold load is running, so a waiting +// request can be told "staging to nvidia-thor" rather than nothing at all. A +// context without a reporter (non-distributed loads, reconciler scale-ups, +// tests) is a no-op. +func reportLoadPhase(ctx context.Context, state string, node *BackendNode, replicaIndex int) { + if p, ok := ctx.Value(loadPhaseKey{}).(*loadPhaseReporter); ok { + p.set(state, node, replicaIndex) + } +} diff --git a/core/services/nodes/load_job_runner.go b/core/services/nodes/load_job_runner.go new file mode 100644 index 000000000..bfb18b1e5 --- /dev/null +++ b/core/services/nodes/load_job_runner.go @@ -0,0 +1,299 @@ +package nodes + +import ( + "context" + "fmt" + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/xlog" +) + +// maxColdLoadRounds bounds how many times a request may claim-or-wait before +// giving up. A round ends when the job reaches a terminal state; a second round +// only happens when the model was evicted between the job finishing and the +// waiter re-checking, which is rare and must not become a spin. +const maxColdLoadRounds = 3 + +// routeViaLoadJob serves a request whose model is not loaded, in distributed +// mode. The cold load itself becomes a durable job owned by whichever replica +// claims it; every other request for the same model — on this replica or any +// other — attaches as a waiter and is served the moment the model is ready. +// +// The per-model advisory lock still de-duplicates loaders, but it is held only +// for the claim. Before this split it wrapped the whole load, so a 35.7 GB +// staging run pinned it for ~20 minutes and every concurrent request died at +// the role's 60s statement_timeout with SQLSTATE 57014. +func (r *SmartRouter) routeViaLoadJob(ctx context.Context, att *routeAttempt) (*RouteResult, error) { + // A held HTTP request cannot survive real infrastructure: an ingress or LB + // idle timeout kills a twenty-minute request regardless of what LocalAI + // does. So the wait is bounded, and expiry produces a structured answer + // carrying live progress rather than letting the connection die anonymously. + budget := r.loadWaitBudget() + waitCtx := ctx + if budget > 0 { + var cancelWait context.CancelFunc + waitCtx, cancelWait = context.WithTimeout(ctx, budget) + defer cancelWait() + } + + for range maxColdLoadRounds { + // Register interest BEFORE claiming, so a job that finishes immediately + // cannot close the channel before this waiter exists. + waiter := r.loadWaiterChan(att.trackingKey) + + job, claimed, err := r.registry.ClaimLoadJob(ctx, att.trackingKey, ReplicaID()) + if err != nil { + // A broken job table must not make the model unroutable: fall back + // to loading inline, which is what every release before this did. + xlog.Warn("Claiming the model load job failed; loading inline instead", + "model", att.trackingKey, "error", err) + loadCtx, cancelLoad := r.newColdLoadContext(context.WithoutCancel(ctx)) + defer cancelLoad() + return r.coldLoad(loadCtx, att, 1) + } + + switch { + case claimed: + // The model may have been loaded between this request's warm-path + // check and the claim — the check the old code did after acquiring + // the lock. Without it the claim would schedule a second copy of a + // model that is already up. + if result := r.tryWarmPath(ctx, att); result != nil { + r.finishLoadJob(ctx, att.trackingKey) + return result, nil + } + r.startLoadJob(ctx, att) + case job != nil && job.State == LoadJobStateFailed: + // Inside the failure grace window: report the real cause rather + // than silently starting a fresh load of a model that just failed. + return nil, fmt.Errorf("loading model %s: %s", att.trackingKey, job.LastError) + default: + xlog.Info("Model is already loading on another replica; waiting for it", + "model", att.trackingKey, "state", job.State, "node", job.NodeName, "owner", job.OwnerReplica) + } + + if err := r.waitForLoadJob(waitCtx, att.trackingKey, waiter); err != nil { + // The caller's own context is still live, so it was the wait budget + // that ran out, not the client giving up: answer with progress. + if ctx.Err() == nil && waitCtx.Err() != nil { + return nil, r.loadingAnswer(ctx, att.trackingKey, budget) + } + return nil, err + } + + // The signal is not the authority — the model may have been evicted + // between ready and wake, so re-run the warm path. + if result := r.tryWarmPath(ctx, att); result != nil { + return result, nil + } + } + return nil, fmt.Errorf("loading model %s: the load finished but the model is not available", att.trackingKey) +} + +// loadWaitBudget resolves the configured wait into a duration, where 0 means +// "no timer — wait as long as the load takes". +func (r *SmartRouter) loadWaitBudget() time.Duration { + switch { + case r.modelLoadWait < 0: // LOCALAI_MODEL_LOAD_WAIT=0 + return 0 + case r.modelLoadWait == 0: // unset + return config.DefaultModelLoadWait + default: + return r.modelLoadWait + } +} + +// loadingAnswer builds the 503 payload for a caller whose wait budget expired, +// reading the job row for live progress. A job that finished in the meantime +// leaves nothing to report, so the caller is told to retry against a plain +// deadline instead. +func (r *SmartRouter) loadingAnswer(ctx context.Context, trackingKey string, budget time.Duration) error { + // The wait context is spent; read on a fresh, short-lived one. + readCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + + job, err := r.registry.GetLoadJob(readCtx, trackingKey) + if err != nil || job == nil { + return fmt.Errorf("timed out waiting for model %s to load", trackingKey) + } + if job.State == LoadJobStateFailed { + return fmt.Errorf("loading model %s: %s", trackingKey, job.LastError) + } + return newModelLoadingError(job, budget) +} + +// startLoadJob runs the claimed cold load in the background, detached from the +// request that triggered it. The job is owned by its record, not by that +// request: the client may disconnect, be retried onto another replica, or time +// out, and the transfer keeps going. +func (r *SmartRouter) startLoadJob(ctx context.Context, att *routeAttempt) { + trackingKey := att.trackingKey + // Keep the request's context VALUES (prefix chain and friends) but none of + // its cancellation — see newColdLoadContext. + parent := context.WithoutCancel(ctx) + + go func() { + loadCtx, cancelLoad := r.newColdLoadContext(parent) + defer cancelLoad() + + phase := newLoadPhaseReporter() + loadCtx = withLoadPhaseReporter(loadCtx, phase) + + stopHeartbeat := r.startLoadJobHeartbeat(parent, trackingKey, phase) + + _, err := r.coldLoad(loadCtx, att, 0) + + stopHeartbeat() + + // Bookkeeping must survive the load context, which may be exactly what + // just expired. + bookCtx, cancelBook := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second) + defer cancelBook() + + if err != nil { + xlog.Error("Cold load job failed", "model", trackingKey, "error", err) + if ferr := r.registry.FailLoadJob(bookCtx, trackingKey, err.Error()); ferr != nil { + xlog.Warn("Failed to record cold load failure", "model", trackingKey, "error", ferr) + } + r.closeLoadWaiters(trackingKey) + // Keep the row briefly so a request arriving right now reports this + // failure instead of starting a duplicate load. Deleting it + // immediately turns a failure into a retry storm. + time.AfterFunc(loadJobFailureGrace, func() { + delCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if derr := r.registry.DeleteLoadJob(delCtx, trackingKey); derr != nil { + xlog.Warn("Failed to clear failed cold load job", "model", trackingKey, "error", derr) + } + }) + return + } + + r.finishLoadJob(bookCtx, trackingKey) + }() +} + +// finishLoadJob ends a job that succeeded. The NodeModel row (state `loaded`) +// is the record from here, so the job row is dropped BEFORE waiters are woken: +// they re-run the warm path and must not find a job that is really done. +func (r *SmartRouter) finishLoadJob(ctx context.Context, trackingKey string) { + if err := r.registry.DeleteLoadJob(ctx, trackingKey); err != nil { + xlog.Warn("Failed to clear completed cold load job", "model", trackingKey, "error", err) + } + r.closeLoadWaiters(trackingKey) +} + +// startLoadJobHeartbeat keeps the job row's liveness and progress fresh while +// the load runs, and returns a function that stops it. +// +// The heartbeat is deliberately time-driven rather than byte-driven: a +// checkpoint load moves no bytes for many minutes, and a job that only wrote a +// row when bytes moved would look orphaned and be reclaimed mid-load. Byte +// progress is copied in from the staging tracker, which already debounces the +// per-chunk callbacks, so the row is written at most once per interval. +func (r *SmartRouter) startLoadJobHeartbeat(parent context.Context, trackingKey string, phase *loadPhaseReporter) func() { + done := make(chan struct{}) + stopped := make(chan struct{}) + + go func() { + defer close(stopped) + ticker := time.NewTicker(loadJobHeartbeatInterval) + defer ticker.Stop() + var startedAt time.Time + for { + select { + case <-done: + return + case <-ticker.C: + u := phase.snapshot() + if st := r.stagingTracker.Get(trackingKey); st != nil { + u.BytesSent, u.TotalBytes = st.BytesSent, st.TotalBytes + u.FileIndex, u.TotalFiles = st.FileIndex, st.TotalFiles + if u.BytesSent > 0 && startedAt.IsZero() { + startedAt = time.Now() + } + u.StartedAt = startedAt + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(parent), loadJobHeartbeatInterval*5) + if err := r.registry.UpdateLoadJob(ctx, trackingKey, u); err != nil { + xlog.Debug("Failed to heartbeat cold load job", "model", trackingKey, "error", err) + } + cancel() + } + } + }() + + return func() { + close(done) + <-stopped + } +} + +// waitForLoadJob blocks until the cold load of trackingKey reaches a terminal +// state, the job's failure is known, or the caller gives up. +// +// Waiters share one broadcast rather than an ordered queue: they all want the +// identical outcome — the model loaded — so ordering them would add fairness +// machinery that changes no result. The local channel wakes same-replica +// waiters instantly; the DB poll is the authority, because a waiter on another +// replica has no channel to close and NATS broadcasts are fire-and-forget, so a +// missed terminal event must not strand it. +func (r *SmartRouter) waitForLoadJob(ctx context.Context, trackingKey string, waiter <-chan struct{}) error { + ticker := time.NewTicker(loadJobPollInterval) + defer ticker.Stop() + + for { + select { + case <-waiter: + return nil + case <-ctx.Done(): + // The client gave up. The job is unaffected: it is owned by the job + // record, not by this request. + return ctx.Err() + case <-ticker.C: + job, err := r.registry.GetLoadJob(ctx, trackingKey) + if err != nil { + xlog.Debug("Polling the model load job failed", "model", trackingKey, "error", err) + continue + } + if job == nil { + // Terminal: either it succeeded, or it was reaped. Either way + // the caller re-checks the warm path. + return nil + } + if job.State == LoadJobStateFailed { + return fmt.Errorf("loading model %s: %s", trackingKey, job.LastError) + } + } + } +} + +// loadWaiterChan returns the broadcast channel for trackingKey, creating it on +// first use. Same shape as advisorylock.localLocks: N local requests share one +// wait and wake together. +func (r *SmartRouter) loadWaiterChan(trackingKey string) <-chan struct{} { + r.loadWaitersMu.Lock() + defer r.loadWaitersMu.Unlock() + if r.loadWaiters == nil { + r.loadWaiters = map[string]chan struct{}{} + } + ch, ok := r.loadWaiters[trackingKey] + if !ok { + ch = make(chan struct{}) + r.loadWaiters[trackingKey] = ch + } + return ch +} + +// closeLoadWaiters wakes every local waiter on trackingKey. A waiter that +// registers after this sees a fresh channel and falls back to the DB poll. +func (r *SmartRouter) closeLoadWaiters(trackingKey string) { + r.loadWaitersMu.Lock() + ch, ok := r.loadWaiters[trackingKey] + delete(r.loadWaiters, trackingKey) + r.loadWaitersMu.Unlock() + if ok { + close(ch) + } +} diff --git a/core/services/nodes/model_load_job.go b/core/services/nodes/model_load_job.go new file mode 100644 index 000000000..eaf50c789 --- /dev/null +++ b/core/services/nodes/model_load_job.go @@ -0,0 +1,262 @@ +package nodes + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + "github.com/google/uuid" + "github.com/mudler/LocalAI/core/services/advisorylock" + "gorm.io/gorm" +) + +// Cold-load job states. `pending` covers node selection and replica +// allocation, which report nothing a waiter could act on; the rest name the +// phase the load is actually in. There is no terminal `ready` state — a +// successful job deletes its row and leaves the NodeModel row as the record. +const ( + LoadJobStatePending = "pending" + LoadJobStateInstalling = "installing" + LoadJobStateStaging = "staging" + LoadJobStateLoading = "loading" + LoadJobStateFailed = "failed" +) + +const ( + // loadJobHeartbeatInterval is how often a running job touches LastProgress. + // It matches the staging broadcast debounce so a job writes at most one row + // per second regardless of how many 32 KB chunks land in it. + loadJobHeartbeatInterval = stagingBroadcastInterval + + // loadJobOrphanWindow is how long a job may go without a heartbeat before + // another replica may reclaim it. Generous relative to the 1s heartbeat: a + // frontend under GC pressure or a stalled DB write must not have its + // perfectly healthy multi-GB transfer stolen and restarted from zero. + loadJobOrphanWindow = 60 * time.Second + + // loadJobFailureGrace is how long a failed job row is kept before deletion. + // Without it a waiter polling just after the failure finds no row, concludes + // "not loading", and starts a duplicate load of a model that just failed — + // a retry storm dressed as recovery. + loadJobFailureGrace = 15 * time.Second + + // loadJobPollInterval is how often a waiter on a non-owning replica polls + // the job row. The DB is the authority: NATS staging broadcasts are + // fire-and-forget, so a missed terminal event must not strand a waiter. + loadJobPollInterval = 2 * time.Second +) + +// loadJobLockPrefix namespaces the per-model advisory lock key. It is the same +// key the whole cold load used to hold; only the guarded section changed. +const loadJobLockPrefix = "model-load:" + +var ( + replicaIDOnce sync.Once + replicaIDValue string +) + +// ReplicaID returns this process's identity, generated once at startup and held +// for the process lifetime. It is recorded on jobs for diagnostics only, never +// for correctness decisions: a replica cannot be assumed alive just because its +// ID is on a row, which is what the LastProgress heartbeat is for. +func ReplicaID() string { + replicaIDOnce.Do(func() { replicaIDValue = uuid.New().String() }) + return replicaIDValue +} + +// IsOrphaned reports whether the job's owner has stopped heartbeating and the +// job may be reclaimed by another replica. +func (j *ModelLoadJob) IsOrphaned(now time.Time) bool { + return now.Sub(j.LastProgress) > loadJobOrphanWindow +} + +// Progress returns overall completion as a percentage, or 0 when the job has +// not reported enough to compute one. +func (j *ModelLoadJob) Progress() float64 { + if j.TotalBytes <= 0 { + return 0 + } + filePct := float64(j.BytesSent) / float64(j.TotalBytes) * 100 + if j.TotalFiles <= 1 || j.FileIndex <= 0 { + return filePct + } + return (float64(j.FileIndex-1)*100 + filePct) / float64(j.TotalFiles) +} + +// ETA returns the estimated time remaining for the transfer, and false when the +// job has not moved enough bytes for the observed rate to mean anything. A +// confidently wrong ETA on a twenty-minute wait is worse than none, so this +// omits rather than guesses. +func (j *ModelLoadJob) ETA(now time.Time) (time.Duration, bool) { + if j.State != LoadJobStateStaging || j.BytesSent <= 0 || j.TotalBytes <= j.BytesSent { + return 0, false + } + if j.StartedAt.IsZero() { + return 0, false + } + elapsed := now.Sub(j.StartedAt) + if elapsed < loadJobHeartbeatInterval { + return 0, false + } + rate := float64(j.BytesSent) / elapsed.Seconds() + if rate <= 0 { + return 0, false + } + return time.Duration(float64(j.TotalBytes-j.BytesSent)/rate) * time.Second, true +} + +// LoadJobUpdate is a partial update to a running job. Empty node fields are +// left untouched so a heartbeat does not erase the placement the runner +// reported earlier. +type LoadJobUpdate struct { + State string + NodeID string + NodeName string + ReplicaIndex int + BytesSent int64 + TotalBytes int64 + FileIndex int + TotalFiles int + // StartedAt anchors the rate the ETA is derived from. Set by the runner the + // first time the transfer reports bytes; zero leaves the stored value alone. + StartedAt time.Time +} + +// ClaimLoadJob decides, under the per-model advisory lock, whether this replica +// owns the cold load of trackingKey. It returns the live job and claimed=false +// when another replica is already loading it (or it just failed and is inside +// its grace window), or a fresh `pending` job with claimed=true when this +// replica took the work. +// +// The lock is held only across these statements — no network, file, or gRPC I/O +// happens inside it, which is the entire point of the job row. The primary key +// on TrackingKey is the real guard: if the lock were somehow bypassed the +// INSERT fails rather than producing two loaders. +func (r *NodeRegistry) ClaimLoadJob(ctx context.Context, trackingKey, owner string) (*ModelLoadJob, bool, error) { + var ( + job *ModelLoadJob + claimed bool + ) + lockKey := advisorylock.KeyFromString(loadJobLockPrefix + trackingKey) + err := advisorylock.WithLockCtx(ctx, r.db, lockKey, func() error { + var existing ModelLoadJob + err := r.db.WithContext(ctx).First(&existing, "tracking_key = ?", trackingKey).Error + switch { + case err == nil: + if !existing.IsOrphaned(time.Now()) { + job, claimed = &existing, false + return nil + } + // The owning replica died mid-load. Without this a crashed frontend + // would wedge the model permanently: every later request would find + // a job row that nobody is running and wait for a load that will + // never progress. + if err := r.db.WithContext(ctx).Delete(&ModelLoadJob{}, "tracking_key = ?", trackingKey).Error; err != nil { + return fmt.Errorf("deleting orphaned model load job: %w", err) + } + case errors.Is(err, gorm.ErrRecordNotFound): + default: + return fmt.Errorf("reading model load job: %w", err) + } + + now := time.Now() + fresh := &ModelLoadJob{ + TrackingKey: trackingKey, + State: LoadJobStatePending, + OwnerReplica: owner, + CreatedAt: now, + UpdatedAt: now, + LastProgress: now, + } + if err := r.db.WithContext(ctx).Create(fresh).Error; err != nil { + return fmt.Errorf("creating model load job: %w", err) + } + job, claimed = fresh, true + return nil + }) + if err != nil { + return nil, false, err + } + return job, claimed, nil +} + +// GetLoadJob returns the active job for trackingKey, or (nil, nil) when none is +// active. Callers on a non-owning replica poll this; it is the authority for +// both readiness and failure. +func (r *NodeRegistry) GetLoadJob(ctx context.Context, trackingKey string) (*ModelLoadJob, error) { + var job ModelLoadJob + err := r.db.WithContext(ctx).First(&job, "tracking_key = ?", trackingKey).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("reading model load job: %w", err) + } + return &job, nil +} + +// UpdateLoadJob applies a phase transition or heartbeat. LastProgress is always +// touched: it is the liveness signal the orphan check reads, and it must tick +// even during phases that move no bytes at all. +func (r *NodeRegistry) UpdateLoadJob(ctx context.Context, trackingKey string, u LoadJobUpdate) error { + now := time.Now() + fields := map[string]any{ + "last_progress": now, + "updated_at": now, + "bytes_sent": u.BytesSent, + "total_bytes": u.TotalBytes, + "file_index": u.FileIndex, + "total_files": u.TotalFiles, + } + if u.State != "" { + fields["state"] = u.State + } + if u.NodeID != "" { + fields["node_id"] = u.NodeID + fields["replica_index"] = u.ReplicaIndex + } + if u.NodeName != "" { + fields["node_name"] = u.NodeName + } + if !u.StartedAt.IsZero() { + fields["started_at"] = u.StartedAt + } + res := r.db.WithContext(ctx).Model(&ModelLoadJob{}). + Where("tracking_key = ?", trackingKey).Updates(fields) + if res.Error != nil { + return fmt.Errorf("updating model load job: %w", res.Error) + } + return nil +} + +// FailLoadJob records the real failure on the job row so every waiter — local +// or on another replica — reports the same cause instead of an anonymous +// timeout. The row is deleted after loadJobFailureGrace by the runner. +func (r *NodeRegistry) FailLoadJob(ctx context.Context, trackingKey, msg string) error { + now := time.Now() + res := r.db.WithContext(ctx).Model(&ModelLoadJob{}). + Where("tracking_key = ?", trackingKey). + Updates(map[string]any{ + "state": LoadJobStateFailed, + "last_error": msg, + "last_progress": now, + "updated_at": now, + }) + if res.Error != nil { + return fmt.Errorf("failing model load job: %w", res.Error) + } + return nil +} + +// DeleteLoadJob removes a terminal job row. Success deletes immediately (the +// NodeModel row is the record of a loaded model); failures delete after their +// grace window. +func (r *NodeRegistry) DeleteLoadJob(ctx context.Context, trackingKey string) error { + if err := r.db.WithContext(ctx). + Delete(&ModelLoadJob{}, "tracking_key = ?", trackingKey).Error; err != nil { + return fmt.Errorf("deleting model load job: %w", err) + } + return nil +} diff --git a/core/services/nodes/model_load_job_test.go b/core/services/nodes/model_load_job_test.go new file mode 100644 index 000000000..a304a6b3d --- /dev/null +++ b/core/services/nodes/model_load_job_test.go @@ -0,0 +1,214 @@ +package nodes + +import ( + "context" + "runtime" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/testutil" + "gorm.io/gorm" +) + +var _ = Describe("ModelLoadJob", func() { + var ( + db *gorm.DB + registry *NodeRegistry + ctx context.Context + ) + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + ctx = context.Background() + }) + + Describe("ClaimLoadJob", func() { + It("claims a model that has no job yet", func() { + job, claimed, err := registry.ClaimLoadJob(ctx, "qwen3", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + Expect(job.State).To(Equal(LoadJobStatePending)) + Expect(job.OwnerReplica).To(Equal("replica-a")) + }) + + It("hands a second caller the live job instead of a second claim", func() { + _, claimed, err := registry.ClaimLoadJob(ctx, "qwen3", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + + job, claimed, err := registry.ClaimLoadJob(ctx, "qwen3", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeFalse(), "the second caller must attach as a waiter, not start a duplicate load") + Expect(job.OwnerReplica).To(Equal("replica-a")) + }) + + It("gives exactly one of many concurrent claimers the job", func() { + const claimers = 8 + var claims int32 + var wg sync.WaitGroup + for range claimers { + wg.Go(func() { + defer GinkgoRecover() + _, claimed, err := registry.ClaimLoadJob(ctx, "contended", "replica") + Expect(err).ToNot(HaveOccurred()) + if claimed { + atomic.AddInt32(&claims, 1) + } + }) + } + wg.Wait() + Expect(claims).To(Equal(int32(1))) + }) + + It("returns promptly while another replica's job is running", func() { + // The whole point of the split: a claim is a decision that takes + // milliseconds, so a concurrent request never waits behind the + // minutes-long load the owner is running. + _, claimed, err := registry.ClaimLoadJob(ctx, "slow-model", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + + running := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer GinkgoRecover() + close(running) + // Stand in for a multi-GB staging run owned by replica-a. + time.Sleep(2 * time.Second) + Expect(registry.DeleteLoadJob(ctx, "slow-model")).To(Succeed()) + close(done) + }() + <-running + + start := time.Now() + _, claimed, err = registry.ClaimLoadJob(ctx, "slow-model", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeFalse()) + Expect(time.Since(start)).To(BeNumerically("<", 100*time.Millisecond), + "claiming must not block behind the running job") + <-done + }) + + It("reclaims a job whose owner stopped heartbeating", func() { + _, claimed, err := registry.ClaimLoadJob(ctx, "orphan", "dead-replica") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + + // Backdate the heartbeat past the orphan window, as a replica killed + // mid-load would leave it. + stale := time.Now().Add(-2 * loadJobOrphanWindow) + Expect(db.Model(&ModelLoadJob{}).Where("tracking_key = ?", "orphan"). + Update("last_progress", stale).Error).ToNot(HaveOccurred()) + + job, claimed, err := registry.ClaimLoadJob(ctx, "orphan", "live-replica") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue(), "a dead replica must not wedge the model permanently") + Expect(job.OwnerReplica).To(Equal("live-replica")) + }) + }) + + Describe("lifecycle", func() { + It("records progress and clears the row on completion", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m1", "replica-a") + Expect(err).ToNot(HaveOccurred()) + + started := time.Now() + Expect(registry.UpdateLoadJob(ctx, "m1", LoadJobUpdate{ + State: LoadJobStateStaging, NodeID: "node-1", NodeName: "nvidia-thor", + ReplicaIndex: 2, BytesSent: 500, TotalBytes: 1000, FileIndex: 1, TotalFiles: 1, + StartedAt: started, + })).To(Succeed()) + + job, err := registry.GetLoadJob(ctx, "m1") + Expect(err).ToNot(HaveOccurred()) + Expect(job.State).To(Equal(LoadJobStateStaging)) + Expect(job.NodeName).To(Equal("nvidia-thor")) + Expect(job.ReplicaIndex).To(Equal(2)) + Expect(job.Progress()).To(BeNumerically("~", 50, 0.01)) + + Expect(registry.DeleteLoadJob(ctx, "m1")).To(Succeed()) + job, err = registry.GetLoadJob(ctx, "m1") + Expect(err).ToNot(HaveOccurred()) + Expect(job).To(BeNil()) + }) + + It("keeps the placement across a byte-less heartbeat", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m2", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(registry.UpdateLoadJob(ctx, "m2", LoadJobUpdate{ + State: LoadJobStateStaging, NodeID: "node-1", NodeName: "nvidia-thor", + })).To(Succeed()) + + before, err := registry.GetLoadJob(ctx, "m2") + Expect(err).ToNot(HaveOccurred()) + + // A checkpoint load moves no bytes for minutes; the heartbeat must + // still tick, and must not erase where the model is loading. + time.Sleep(10 * time.Millisecond) + Expect(registry.UpdateLoadJob(ctx, "m2", LoadJobUpdate{State: LoadJobStateLoading})).To(Succeed()) + + after, err := registry.GetLoadJob(ctx, "m2") + Expect(err).ToNot(HaveOccurred()) + Expect(after.NodeName).To(Equal("nvidia-thor")) + Expect(after.State).To(Equal(LoadJobStateLoading)) + Expect(after.LastProgress.After(before.LastProgress)).To(BeTrue()) + }) + + It("records the failure cause for waiters to read", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m3", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(registry.FailLoadJob(ctx, "m3", "no available nodes")).To(Succeed()) + + job, err := registry.GetLoadJob(ctx, "m3") + Expect(err).ToNot(HaveOccurred()) + Expect(job.State).To(Equal(LoadJobStateFailed)) + Expect(job.LastError).To(Equal("no available nodes")) + }) + + It("hands a request arriving inside the failure grace the real error", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m4", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(registry.FailLoadJob(ctx, "m4", "worker out of VRAM")).To(Succeed()) + + job, claimed, err := registry.ClaimLoadJob(ctx, "m4", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeFalse(), "a fresh load must not silently start on top of a just-failed one") + Expect(job.LastError).To(Equal("worker out of VRAM")) + }) + }) + + Describe("ETA", func() { + It("is omitted until the observed rate means something", func() { + job := &ModelLoadJob{State: LoadJobStateStaging, BytesSent: 0, TotalBytes: 1000} + _, ok := job.ETA(time.Now()) + Expect(ok).To(BeFalse()) + + job = &ModelLoadJob{State: LoadJobStateStaging, BytesSent: 10, TotalBytes: 1000, StartedAt: time.Now()} + _, ok = job.ETA(time.Now()) + Expect(ok).To(BeFalse(), "less than one broadcast interval of data is not a rate") + }) + + It("derives the remaining time from the job's own rate", func() { + now := time.Now() + job := &ModelLoadJob{ + State: LoadJobStateStaging, + BytesSent: 1000, TotalBytes: 3000, + StartedAt: now.Add(-10 * time.Second), + } + eta, ok := job.ETA(now) + Expect(ok).To(BeTrue()) + // 100 B/s observed, 2000 B left. + Expect(eta).To(BeNumerically("~", 20*time.Second, time.Second)) + }) + }) +}) diff --git a/core/services/nodes/model_loading_error.go b/core/services/nodes/model_loading_error.go new file mode 100644 index 000000000..a7c8d7fb4 --- /dev/null +++ b/core/services/nodes/model_loading_error.go @@ -0,0 +1,72 @@ +package nodes + +import ( + "fmt" + "time" + + "github.com/mudler/LocalAI/core/schema" +) + +const ( + // retryAfterFloor and retryAfterCeiling clamp the Retry-After we hand a + // client. Below the floor a client hammers a load that cannot possibly be + // done yet; above the ceiling it stops polling long enough that a model + // which became ready in the meantime sits idle. + retryAfterFloor = 5 * time.Second + retryAfterCeiling = 300 * time.Second +) + +// ModelLoadingError reports that the request's model is still cold-loading and +// the caller's wait budget ran out. It carries live progress so the answer is +// actionable — "staging to nvidia-thor, 41%, ETA ~11m" — rather than an +// anonymous timeout, which is what every UI retry produced before. +type ModelLoadingError struct { + Status schema.ModelLoadingStatus + RetryAfter time.Duration +} + +func (e *ModelLoadingError) Error() string { + msg := fmt.Sprintf("model %s is %s", e.Status.Model, e.Status.State) + if e.Status.Node != "" { + msg += " on node " + e.Status.Node + } + if e.Status.Progress > 0 { + msg += fmt.Sprintf(" (%.0f%%", e.Status.Progress) + if e.Status.ETASeconds > 0 { + msg += fmt.Sprintf(", ETA ~%s", (time.Duration(e.Status.ETASeconds) * time.Second).Round(time.Minute)) + } + msg += ")" + } + return msg +} + +// LoadingStatus renders a job row as the API's `loading` object. +func LoadingStatus(job *ModelLoadJob) schema.ModelLoadingStatus { + status := schema.ModelLoadingStatus{ + Model: job.TrackingKey, + State: job.State, + Node: job.NodeName, + Progress: job.Progress(), + BytesSent: job.BytesSent, + TotalBytes: job.TotalBytes, + FileIndex: job.FileIndex, + TotalFiles: job.TotalFiles, + } + if eta, ok := job.ETA(time.Now()); ok { + status.ETASeconds = int(eta.Seconds()) + } + return status +} + +// newModelLoadingError builds the 503 answer for a caller whose wait budget +// expired. Retry-After is the ETA when the job has one, clamped so it stays a +// useful poll interval, and the caller's own budget otherwise. +func newModelLoadingError(job *ModelLoadJob, budget time.Duration) *ModelLoadingError { + status := LoadingStatus(job) + retryAfter := budget + if status.ETASeconds > 0 { + retryAfter = time.Duration(status.ETASeconds) * time.Second + } + retryAfter = min(max(retryAfter, retryAfterFloor), retryAfterCeiling) + return &ModelLoadingError{Status: status, RetryAfter: retryAfter} +} diff --git a/core/services/nodes/model_router_test.go b/core/services/nodes/model_router_test.go index f68caa143..b31e9dc06 100644 --- a/core/services/nodes/model_router_test.go +++ b/core/services/nodes/model_router_test.go @@ -14,6 +14,7 @@ import ( // --- fakeModelRouterForSmartRouter implements ModelRouter --- type fakeModelRouterForSmartRouter struct { + fakeLoadJobStore mu sync.Mutex node *BackendNode nodeModel *NodeModel diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 0b85a1f44..1a59e361f 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -241,6 +241,46 @@ type PendingBackendOp struct { NextRetryAt time.Time `gorm:"index" json:"next_retry_at"` } +// ModelLoadJob is one in-flight cold load of a model. Exactly one row per +// trackingKey may be active at a time; that uniqueness — not the lifetime of an +// advisory lock — is what de-duplicates concurrent loaders across replicas. +// +// Before this table the whole cold load (backend install, multi-GB staging, +// checkpoint load) ran inside the per-model advisory lock, so every other +// replica's request for the model blocked on pg_advisory_lock for tens of +// minutes and was killed by the role's statement_timeout. The job row lets the +// lock shrink to the claim while the work itself runs unlocked and observable. +// +// Terminal rows are deleted rather than retained: NodeModel is already the +// record of what is loaded, and keeping finished jobs would create a second +// source of truth about it. +type ModelLoadJob struct { + TrackingKey string `gorm:"primaryKey;size:255" json:"tracking_key"` + State string `gorm:"size:16;not null;index" json:"state"` + OwnerReplica string `gorm:"size:64" json:"owner_replica"` + NodeID string `gorm:"size:36" json:"node_id"` + NodeName string `gorm:"size:255" json:"node_name"` + ReplicaIndex int `json:"replica_index"` + BytesSent int64 `json:"bytes_sent"` + TotalBytes int64 `json:"total_bytes"` + FileIndex int `json:"file_index"` + TotalFiles int `json:"total_files"` + LastError string `gorm:"type:text" json:"last_error,omitempty"` + // StartedAt is when the job first reported bytes, and is what the ETA rate + // is measured from. Distinct from CreatedAt, which also covers node + // selection and backend install — phases that move no bytes and would skew + // the derived rate low. + StartedAt time.Time `json:"started_at"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + // LastProgress is a heartbeat, not a byte counter: the runner touches it on + // a fixed interval for as long as it is alive, whether or not bytes are + // moving. A checkpoint load legitimately transfers zero bytes for many + // minutes, so a reaper keyed on byte movement would reclaim a healthy job + // mid-load. Byte progress is measured separately, by load_deadline.go. + LastProgress time.Time `gorm:"index" json:"last_progress_at"` +} + // Op constants mirror the operation names used by DistributedBackendManager // so callers don't repeat stringly-typed values. const ( @@ -343,7 +383,7 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}) + return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index fd051e954..dbcd23bb8 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -14,7 +14,6 @@ import ( "time" "github.com/mudler/LocalAI/core/config" - "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/pkg/distributedhdr" grpc "github.com/mudler/LocalAI/pkg/grpc" @@ -106,6 +105,11 @@ type SmartRouterOptions struct { // arriving, so a peer trickling bytes forever cannot pin the advisory lock // indefinitely. Zero selects modelLoadAbsoluteMax (24h). ModelLoadAbsoluteMax time.Duration + // ModelLoadWait bounds how long a REQUEST waits for a cold load that is + // already running before it is answered with live progress. It bounds the + // caller, never the load: the job keeps running either way. Zero selects + // config.DefaultModelLoadWait; config.ModelLoadWaitUnbounded waits forever. + ModelLoadWait time.Duration } // modelLoadStagingMargin is the slack ModelLoadCeilingFor adds on top of the @@ -188,6 +192,15 @@ type SmartRouter struct { // hard countdown into a progress-extended hold (see load_deadline.go). stagingStallWindow time.Duration modelLoadAbsoluteMax time.Duration + // modelLoadWait bounds the REQUEST's wait for a running cold load, not the + // load itself (see SmartRouterOptions.ModelLoadWait). + modelLoadWait time.Duration + // loadWaiters is one broadcast channel per model being cold-loaded, closed + // when the job reaches a terminal state. Same-model waiters all want the + // identical outcome, so they share one wait instead of queueing. See + // load_job_runner.go. + loadWaitersMu sync.Mutex + loadWaiters map[string]chan struct{} } // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is @@ -239,6 +252,8 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter // ceiling, so nothing to normalize here. stagingStallWindow: opts.StagingStallWindow, modelLoadAbsoluteMax: opts.ModelLoadAbsoluteMax, + modelLoadWait: opts.ModelLoadWait, + loadWaiters: map[string]chan struct{}{}, } } @@ -329,6 +344,7 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking 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) } + reportLoadPhase(ctx, LoadJobStateStaging, node, replicaIndex) lifecycleSettled := false defer func() { if lifecycleSettled { @@ -367,6 +383,7 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking 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) } + reportLoadPhase(ctx, LoadJobStateLoading, node, replicaIndex) // The cold-load hold above this call extends on STAGING progress, and // the remote LoadModel reports none — so once the last byte lands the @@ -555,138 +572,149 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType // below. Both are nil (no-op) when prefix-cache routing is disabled. pref, observeChain := r.buildPreference(ctx, trackingKey, candidateNodeIDs, sched) + att := &routeAttempt{ + trackingKey: trackingKey, + modelName: modelName, + backendType: backendType, + modelOpts: modelOpts, + parallel: parallel, + sched: sched, + candidateNodeIDs: candidateNodeIDs, + pref: pref, + observeChain: observeChain, + } + // Step 1: Find and atomically lock a node with this model loaded - node, nm, err := r.registry.FindAndLockNodeWithModel(ctx, trackingKey, candidateNodeIDs, pref) - if err == nil && node != nil { - modelAddr := node.Address - if nm.Address != "" { - modelAddr = nm.Address - } - replicaIdx := nm.ReplicaIndex - - // Verify the backend process is still alive via gRPC health check - if !r.probeHealth(ctx, node, modelAddr) { - // Stale — roll back the increment, remove the specific replica row, fall through - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - r.registry.RemoveNodeModel(ctx, node.ID, trackingKey, replicaIdx) - xlog.Warn("Backend not reachable for cached model, falling through to reload", - "node", node.Name, "model", modelName, "replica", replicaIdx) - } else { - // Verify node still matches scheduling constraints - if !r.nodeMatchesScheduling(ctx, node, sched) { - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - xlog.Info("Cached model on node that no longer matches selector, falling through", - "node", node.Name, "model", trackingKey, "replica", replicaIdx) - // Fall through to step 2 (scheduleNewModel) - } else { - // Node is alive — FindAndLockNodeWithModel already incremented in-flight as a - // reservation. InFlightTrackingClient handles per-inference tracking, and its - // onFirstComplete callback releases the reservation after the first inference - // call finishes, so in-flight returns to 0 when idle. - r.registry.TouchNodeModel(ctx, node.ID, trackingKey, replicaIdx) - r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx}) - grpcClient := r.buildClientForAddr(node, modelAddr, parallel) - tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, trackingKey, replicaIdx) - return r.newRouteResult(node, trackingKey, replicaIdx, grpcClient, tracked), nil - } - } - } - - // Step 2: Model not loaded — schedule loading with distributed lock to prevent duplicates. - // - // Detach the cold-load from the caller's context. Staging a model can - // transfer multiple GB to a worker, which takes far longer than any client - // keeps its HTTP request open — a browser refresh, an ingress/LB idle - // timeout, or a round-robined retry landing on another replica all cancel - // the request context. If staging were bound to it, the multi-GB upload - // aborts with "context canceled" mid-transfer and large models can never - // finish staging (the model-load outage). WithoutCancel keeps the request's - // values (prefix chain, etc.) but drops its cancellation/deadline. - // - // Detaching from the caller is necessary, but it must not be unbounded: the - // load runs while holding the per-model advisory lock, and a worker that - // dies mid-install (its backend.install never replies) would otherwise pin - // that lock (and every other replica's request for the same model) until - // the NATS install deadline alone expires. Re-impose a single hard ceiling - // over the whole sequence so the lock is always released in bounded time, - // even if a sub-step wedges. Each long step still has its own (tighter) - // bound; this only backstops them. The per-model advisory lock below - // de-dupes concurrent loaders across replicas. - // The backstop is progress-based, not wall-clock: staging time is bytes over - // bandwidth, so a fixed ceiling is a model-size cliff (a 70 GB checkpoint - // transferring healthily at 26 MB/s needs ~45m and was killed at exactly - // 25m00s). The hold instead extends while the transfer reports bytes and - // expires a stall window after they stop. See load_deadline.go. - loadCtx, cancelLoad := newLoadDeadlineContext(context.WithoutCancel(ctx), - r.modelLoadCeiling, r.stagingStallWindow, r.modelLoadAbsoluteMax) - defer cancelLoad() - loadModel := func(ctx context.Context) (*RouteResult, error) { - // Re-check after acquiring lock — another request may have loaded it - node, nm, err := r.registry.FindAndLockNodeWithModel(ctx, trackingKey, candidateNodeIDs, pref) - if err == nil && node != nil { - modelAddr := node.Address - if nm.Address != "" { - modelAddr = nm.Address - } - replicaIdx := nm.ReplicaIndex - - // Verify the backend process is still alive via gRPC health check - if !r.probeHealth(ctx, node, modelAddr) { - // Stale — roll back the increment, remove the specific replica row, continue loading - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - r.registry.RemoveNodeModel(ctx, node.ID, trackingKey, replicaIdx) - xlog.Warn("Backend not reachable for cached model inside lock, proceeding to load", - "node", node.Name, "model", modelName, "replica", replicaIdx) - } else { - // Verify node still matches scheduling constraints - if !r.nodeMatchesScheduling(ctx, node, sched) { - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - xlog.Info("Cached model on node that no longer matches selector, falling through", - "node", node.Name, "model", trackingKey, "replica", replicaIdx) - // Fall through to scheduling below - } else { - // Model loaded while we waited — FindAndLockNodeWithModel already incremented - // in-flight as a reservation. Release it after the first inference completes. - r.registry.TouchNodeModel(ctx, node.ID, trackingKey, replicaIdx) - r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx}) - grpcClient := r.buildClientForAddr(node, modelAddr, parallel) - tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, trackingKey, replicaIdx) - return r.newRouteResult(node, trackingKey, replicaIdx, grpcClient, tracked), nil - } - } - } - - // Still not loaded — use shared schedule-and-load logic, which picks - // both the node and the replica slot. - result, err := r.scheduleAndLoad(ctx, backendType, trackingKey, modelName, modelOpts, parallel, 1) - if err != nil { - return nil, err - } - - // Cold load landed on result.Node replica result.ReplicaIndex: record the - // assignment so subsequent requests with the same prefix prefer it. - r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: result.Node.ID, Replica: result.ReplicaIndex}) - - replicaIdx := result.ReplicaIndex - tracked := NewInFlightTrackingClient(result.Client, r.registry, result.Node.ID, trackingKey, replicaIdx) - return r.newRouteResult(result.Node, trackingKey, replicaIdx, result.Client, tracked), nil - } - - if r.db != nil { - lockKey := advisorylock.KeyFromString("model-load:" + trackingKey) - var result *RouteResult - lockErr := advisorylock.WithLockCtx(loadCtx, r.db, lockKey, func() error { - var err error - result, err = loadModel(loadCtx) - return err - }) - if lockErr != nil { - return nil, fmt.Errorf("loading model %s: %w", trackingKey, lockErr) - } + if result := r.tryWarmPath(ctx, att); result != nil { return result, nil } - // No DB (non-distributed) — proceed without lock - return loadModel(loadCtx) + + // Step 2: model not loaded — it has to be cold-loaded. + // + // In distributed mode that runs as a durable job (see load_job_runner.go): + // the per-model advisory lock now guards only the claim, and the transfer + // itself runs unlocked so a concurrent request for the same model never + // blocks on pg_advisory_lock for the tens of minutes a multi-GB stage takes. + if r.db != nil { + return r.routeViaLoadJob(ctx, att) + } + + // No DB (non-distributed): there is no other replica to coordinate with, so + // the load stays inline on the request exactly as before. + loadCtx, cancelLoad := r.newColdLoadContext(context.WithoutCancel(ctx)) + defer cancelLoad() + if result := r.tryWarmPath(loadCtx, att); result != nil { + return result, nil + } + return r.coldLoad(loadCtx, att, 1) +} + +// routeAttempt is the per-request routing state shared by the warm path, the +// cold-load job runner, and the waiter loop. Bundling it keeps those from +// drifting apart on which candidate set or preference they used. +type routeAttempt struct { + trackingKey string + modelName string + backendType string + modelOpts *pb.ModelOptions + parallel bool + sched *ModelSchedulingConfig + candidateNodeIDs []string + pref *RoutePreference + observeChain []uint64 +} + +// tryWarmPath returns a route to an already-loaded, reachable replica, or nil +// when the model has to be cold-loaded. It is the authority on readiness: a +// waiter woken by a finished job re-runs it rather than trusting the signal, +// because the model may have been evicted between ready and wake. +func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *RouteResult { + node, nm, err := r.registry.FindAndLockNodeWithModel(ctx, att.trackingKey, att.candidateNodeIDs, att.pref) + if err != nil || node == nil { + return nil + } + modelAddr := node.Address + if nm.Address != "" { + modelAddr = nm.Address + } + replicaIdx := nm.ReplicaIndex + + // Verify the backend process is still alive via gRPC health check + if !r.probeHealth(ctx, node, modelAddr) { + // Stale — roll back the increment, remove the specific replica row, fall through + if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to release stale routing reservation", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + if err := r.registry.RemoveNodeModel(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to remove stale model from registry", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + xlog.Warn("Backend not reachable for cached model, falling through to reload", + "node", node.Name, "model", att.modelName, "replica", replicaIdx) + return nil + } + + // Verify node still matches scheduling constraints + if !r.nodeMatchesScheduling(ctx, node, att.sched) { + if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to release unmatched routing reservation", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + xlog.Info("Cached model on node that no longer matches selector, falling through", + "node", node.Name, "model", att.trackingKey, "replica", replicaIdx) + return nil + } + + // Node is alive — FindAndLockNodeWithModel already incremented in-flight as a + // reservation. InFlightTrackingClient handles per-inference tracking, and its + // onFirstComplete callback releases the reservation after the first inference + // call finishes, so in-flight returns to 0 when idle. + r.registry.TouchNodeModel(ctx, node.ID, att.trackingKey, replicaIdx) + r.observePrefix(att.trackingKey, att.observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx}) + grpcClient := r.buildClientForAddr(node, modelAddr, att.parallel) + tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, att.trackingKey, replicaIdx) + return r.newRouteResult(node, att.trackingKey, replicaIdx, grpcClient, tracked) +} + +// coldLoad schedules the model onto a node and loads it, returning a route to +// the replica it landed on. initialInFlight reserves the slot for the calling +// request; the job runner passes 0 because it is loading on nobody's behalf. +func (r *SmartRouter) coldLoad(ctx context.Context, att *routeAttempt, initialInFlight int) (*RouteResult, error) { + result, err := r.scheduleAndLoad(ctx, att.backendType, att.trackingKey, att.modelName, att.modelOpts, att.parallel, initialInFlight) + if err != nil { + return nil, err + } + + // Cold load landed on result.Node replica result.ReplicaIndex: record the + // assignment so subsequent requests with the same prefix prefer it. + r.observePrefix(att.trackingKey, att.observeChain, prefixcache.ReplicaKey{NodeID: result.Node.ID, Replica: result.ReplicaIndex}) + + tracked := NewInFlightTrackingClient(result.Client, r.registry, result.Node.ID, att.trackingKey, result.ReplicaIndex) + return r.newRouteResult(result.Node, att.trackingKey, result.ReplicaIndex, result.Client, tracked), nil +} + +// newColdLoadContext builds the detached, progress-extended context a cold load +// runs under. +// +// Detach the cold load from the caller's context. Staging a model can transfer +// multiple GB to a worker, which takes far longer than any client keeps its +// HTTP request open — a browser refresh, an ingress/LB idle timeout, or a +// round-robined retry landing on another replica all cancel the request +// context. If staging were bound to it, the multi-GB upload aborts with +// "context canceled" mid-transfer and large models can never finish staging +// (the model-load outage). The caller passes context.WithoutCancel, which keeps +// the request's values (prefix chain, etc.) but drops its cancellation. +// +// Detaching must not be unbounded either: a worker that dies mid-install (its +// backend.install never replies) would otherwise leave the job wedged until the +// NATS install deadline alone expires. The backstop is progress-based, not +// wall-clock: staging time is bytes over bandwidth, so a fixed ceiling is a +// model-size cliff (a 70 GB checkpoint transferring healthily at 26 MB/s needs +// ~45m and was killed at exactly 25m00s). The hold extends while the transfer +// reports bytes and expires a stall window after they stop. See load_deadline.go. +func (r *SmartRouter) newColdLoadContext(parent context.Context) (context.Context, context.CancelFunc) { + return newLoadDeadlineContext(parent, r.modelLoadCeiling, r.stagingStallWindow, r.modelLoadAbsoluteMax) } // parseSelectorJSON decodes a JSON node selector string into a map. @@ -1188,6 +1216,7 @@ func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNod if r.unloader == nil { return "", fmt.Errorf("no NATS connection for backend installation") } + reportLoadPhase(ctx, LoadJobStateInstalling, node, replicaIndex) key := fmt.Sprintf("%s|%s|%s|%d", node.ID, backendType, modelID, replicaIndex) // DoChan rather than Do so this wait honors ctx cancellation. InstallBackend diff --git a/core/services/nodes/router_load_job_test.go b/core/services/nodes/router_load_job_test.go new file mode 100644 index 000000000..f67e1399b --- /dev/null +++ b/core/services/nodes/router_load_job_test.go @@ -0,0 +1,278 @@ +package nodes + +import ( + "context" + "errors" + "runtime" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "gorm.io/gorm" +) + +// These specs cover the claim/run split: a cold load runs as a durable job +// outside the per-model advisory lock, and concurrent requests attach to it as +// waiters instead of blocking on pg_advisory_lock (which the production role's +// statement_timeout killed at 60s). +var _ = Describe("Route cold-load jobs", func() { + var ( + db *gorm.DB + registry *NodeRegistry + backend *stubBackend + factory *stubClientFactory + unloader *fakeUnloader + node *BackendNode + ) + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + + node = &BackendNode{ + Name: "worker-1", + NodeType: NodeTypeBackend, + Address: "10.0.0.1:50051", + TotalVRAM: 64_000_000_000, + AvailableVRAM: 64_000_000_000, + } + Expect(registry.Register(context.Background(), node, true)).To(Succeed()) + + backend = &stubBackend{healthResult: true, loadResult: &pb.Result{Success: true}} + factory = &stubClientFactory{client: backend} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}, + } + }) + + newRouter := func() *SmartRouter { + return NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + DB: db, + }) + } + + It("serves a concurrent request for a loading model without a duplicate load", func() { + // The load takes far longer than a request would tolerate holding a + // lock. Both callers must still be served, from ONE load. + release := make(chan struct{}) + unloader.installHook = func() { <-release } + + router := newRouter() + + first := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + first <- err + }() + + // Give the first request time to claim and start its job. + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "big-model") + return job + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + second := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + second <- err + }() + + // Neither caller may be stuck behind a lock: the second request must + // still be waiting (not failed) while the load runs. + Consistently(second, 300*time.Millisecond).ShouldNot(Receive()) + + close(release) + + var firstErr, secondErr error + Eventually(first, 15*time.Second).Should(Receive(&firstErr)) + Eventually(second, 15*time.Second).Should(Receive(&secondErr)) + Expect(firstErr).ToNot(HaveOccurred()) + Expect(secondErr).ToNot(HaveOccurred()) + + unloader.mu.Lock() + installs := len(unloader.installCalls) + unloader.mu.Unlock() + Expect(installs).To(Equal(1), "the waiter must attach to the running job, not start a second load") + + // The job row is the record of an IN-FLIGHT load only; NodeModel is the + // record of a loaded model. + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "big-model") + return job + }, 5*time.Second, 50*time.Millisecond).Should(BeNil()) + }) + + It("reports the load's real failure to every waiter", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + unloader.installReply = &messaging.BackendInstallReply{Success: false, Error: "worker out of disk"} + + router := newRouter() + + first := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/doomed.gguf"}, false) + first <- err + }() + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "doomed") + return job + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + second := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/doomed.gguf"}, false) + second <- err + }() + + close(release) + + var firstErr, secondErr error + Eventually(first, 15*time.Second).Should(Receive(&firstErr)) + Eventually(second, 15*time.Second).Should(Receive(&secondErr)) + Expect(firstErr).To(HaveOccurred()) + Expect(secondErr).To(HaveOccurred()) + Expect(secondErr.Error()).To(ContainSubstring("worker out of disk"), + "a waiter must learn the real cause, not an anonymous timeout") + }) + + It("returns immediately when the client disconnects, leaving the job running", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + defer close(release) + + router := newRouter() + + go func() { + defer GinkgoRecover() + _, _ = router.Route(context.Background(), "detached", "models/detached.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/detached.gguf"}, false) + }() + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "detached") + return job + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + ctx, cancel := context.WithCancel(context.Background()) + waiter := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(ctx, "detached", "models/detached.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/detached.gguf"}, false) + waiter <- err + }() + time.Sleep(100 * time.Millisecond) + cancel() + + var waitErr error + Eventually(waiter, 3*time.Second).Should(Receive(&waitErr)) + Expect(waitErr).To(HaveOccurred()) + + // The job is owned by its record, not by the request that triggered it. + job, err := registry.GetLoadJob(context.Background(), "detached") + Expect(err).ToNot(HaveOccurred()) + Expect(job).ToNot(BeNil(), "cancelling a waiter must not abort the load") + }) + + It("answers with live progress once the wait budget is spent", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + defer close(release) + + router := NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + DB: db, + ModelLoadWait: 500 * time.Millisecond, + }) + + start := time.Now() + _, err := router.Route(context.Background(), "slow-model", "models/slow.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/slow.gguf"}, false) + Expect(err).To(HaveOccurred()) + Expect(time.Since(start)).To(BeNumerically("<", 10*time.Second)) + + var loadingErr *ModelLoadingError + Expect(errors.As(err, &loadingErr)).To(BeTrue(), + "a caller out of budget must get a structured answer, not an anonymous timeout") + Expect(loadingErr.Status.Model).To(Equal("slow-model")) + Expect(loadingErr.Status.State).ToNot(BeEmpty()) + Expect(loadingErr.RetryAfter).To(BeNumerically(">", 0)) + + // The job is untouched: the caller gave up, the load did not. + job, err := registry.GetLoadJob(context.Background(), "slow-model") + Expect(err).ToNot(HaveOccurred()) + Expect(job).ToNot(BeNil()) + }) + + It("waits unbounded when the budget is explicitly disabled", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + + router := NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + DB: db, + ModelLoadWait: config.ModelLoadWaitUnbounded, + }) + + done := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "patient", "models/patient.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/patient.gguf"}, false) + done <- err + }() + + // Well past any default budget shrunk for tests; the caller must still + // be waiting rather than 503-ing. + Consistently(done, 2*time.Second).ShouldNot(Receive()) + close(release) + Eventually(done, 15*time.Second).Should(Receive(BeNil())) + }) + + It("heartbeats the job row so a live load is never mistaken for an orphan", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + defer close(release) + + router := newRouter() + go func() { + defer GinkgoRecover() + _, _ = router.Route(context.Background(), "beating", "models/beating.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/beating.gguf"}, false) + }() + + var first *ModelLoadJob + Eventually(func() *ModelLoadJob { + first, _ = registry.GetLoadJob(context.Background(), "beating") + return first + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + Eventually(func() bool { + job, _ := registry.GetLoadJob(context.Background(), "beating") + return job != nil && job.LastProgress.After(first.LastProgress) + }, 5*time.Second, 200*time.Millisecond).Should(BeTrue(), + "the runner must heartbeat even while no bytes move") + }) +}) diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index a3f63cc4b..6e1d86d47 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -60,6 +60,8 @@ func (f *fakeFileStager) ListRemoteDir(_ context.Context, _, _ string) ([]string // fakeModelRouter implements ModelRouter with configurable return values. type fakeModelRouter struct { + fakeLoadJobStore + // FindAndLockNodeWithModel returns findAndLockNode *BackendNode findAndLockNM *NodeModel @@ -148,6 +150,86 @@ func (f *fakeModelRouter) LoadedReplicaStats(_ context.Context, modelName string return f.loadedReplicaStatsByName[modelName], nil } +// fakeLoadJobStore is an in-memory LoadJobStore so tests that build a +// SmartRouter over a fake registry get the real claim/wait semantics without a +// database. +type fakeLoadJobStore struct { + mu sync.Mutex + jobs map[string]*ModelLoadJob +} + +func (s *fakeLoadJobStore) ClaimLoadJob(_ context.Context, trackingKey, owner string) (*ModelLoadJob, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.jobs == nil { + s.jobs = map[string]*ModelLoadJob{} + } + if existing, ok := s.jobs[trackingKey]; ok && !existing.IsOrphaned(time.Now()) { + cp := *existing + return &cp, false, nil + } + now := time.Now() + job := &ModelLoadJob{TrackingKey: trackingKey, State: LoadJobStatePending, OwnerReplica: owner, CreatedAt: now, UpdatedAt: now, LastProgress: now} + s.jobs[trackingKey] = job + cp := *job + return &cp, true, nil +} + +func (s *fakeLoadJobStore) GetLoadJob(_ context.Context, trackingKey string) (*ModelLoadJob, error) { + s.mu.Lock() + defer s.mu.Unlock() + job, ok := s.jobs[trackingKey] + if !ok { + return nil, nil + } + cp := *job + return &cp, nil +} + +func (s *fakeLoadJobStore) UpdateLoadJob(_ context.Context, trackingKey string, u LoadJobUpdate) error { + s.mu.Lock() + defer s.mu.Unlock() + job, ok := s.jobs[trackingKey] + if !ok { + return nil + } + if u.State != "" { + job.State = u.State + } + if u.NodeID != "" { + job.NodeID = u.NodeID + job.ReplicaIndex = u.ReplicaIndex + } + if u.NodeName != "" { + job.NodeName = u.NodeName + } + if !u.StartedAt.IsZero() { + job.StartedAt = u.StartedAt + } + job.BytesSent, job.TotalBytes = u.BytesSent, u.TotalBytes + job.FileIndex, job.TotalFiles = u.FileIndex, u.TotalFiles + job.LastProgress = time.Now() + return nil +} + +func (s *fakeLoadJobStore) FailLoadJob(_ context.Context, trackingKey, msg string) error { + s.mu.Lock() + defer s.mu.Unlock() + if job, ok := s.jobs[trackingKey]; ok { + job.State = LoadJobStateFailed + job.LastError = msg + job.LastProgress = time.Now() + } + return nil +} + +func (s *fakeLoadJobStore) DeleteLoadJob(_ context.Context, trackingKey string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.jobs, trackingKey) + return nil +} + func (f *fakeModelRouter) DecrementInFlight(_ context.Context, nodeID, modelName string, _ int) error { f.decrementCalls = append(f.decrementCalls, nodeID+":"+modelName) return nil diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index f3a0f6534..bf5a87261 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -75,6 +75,7 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--backend-install-timeout` | `LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT` | `15m` | How long the frontend waits for a worker to acknowledge a backend install before considering the request stalled. Raise it when workers pull large backend images over slow links. If a worker takes longer than this, the operation shows as "still installing in background" in the admin UI and clears once the worker finishes. | | `--backend-upgrade-timeout` | `LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT` | `15m` | Same as the install timeout, applied to backend upgrades (force-reinstall). | | `--model-load-timeout` | `LOCALAI_NATS_MODEL_LOAD_TIMEOUT` | *(derived from checkpoint size)* | Pins the deadline for the `LoadModel` gRPC call the frontend issues to a worker. Leave it unset: by default the deadline is **derived from the checkpoint's on-disk size** (see below), which is what the worker actually spends its load time reading. Set it only to pin a specific budget — the value is then used verbatim, including when it is *shorter* than the derived one, so an operator who wants fast failure gets it. | +| *(env only)* | `LOCALAI_MODEL_LOAD_WAIT` | `60s` | 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. 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. See [Requests for a model that is still loading](#requests-for-a-model-that-is-still-loading). | | `--expose-node-header` | `LOCALAI_EXPOSE_NODE_HEADER` | `false` | When enabled, inference responses carry an `X-LocalAI-Node` header with the ID of the worker node that served the request. Coverage spans the OpenAI-compatible endpoints (chat completions, completions, embeddings, audio transcriptions, audio speech / TTS, image generations, image inpainting), the Jina rerank endpoint (`/v1/rerank`), the VAD endpoints (`/v1/vad`, `/vad`), and the Anthropic Messages (`/v1/messages`) and Ollama (`/api/chat`, `/api/generate`, `/api/embed`) shims. Useful for debugging, observability and load-balancer attribution. Off by default: the node ID reveals internal cluster topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency for the same model across multiple replicas, the header may reflect a recent routing decision rather than this exact request's. Acceptable for observability and debugging. | ### The model load deadline scales with the checkpoint @@ -111,6 +112,48 @@ While **model files are staging**, however, the deadline extends every time stag An absolute cap of 24h ends the hold even if progress keeps arriving, so a degenerate peer trickling a few bytes at a time cannot pin the lock forever. No configuration is needed for either value; both are sized well above any legitimate transfer. +### Requests for a model that is still loading + +A cold load in distributed mode is a long-running background job: install the backend, stage multi-GB model files to the worker, then load the checkpoint. Staging a 35.7 GB GGUF onto a fresh worker takes roughly twenty minutes on a fast LAN — far longer than any HTTP request can be held open. + +So the load does **not** run on the request. The first request for an unloaded model claims a durable **model load job** — that claim takes milliseconds and is the only part that holds the per-model advisory lock — and the job then runs in the background on the frontend replica that claimed it. Every other request for the same model, on any replica, attaches to that job as a waiter: + +- It is **served the moment the model is ready**, with no client-side retry. A model already 90% staged usually needs no second request. +- It never starts a duplicate load and never blocks on the database lock. (Before this split, concurrent requests blocked on `pg_advisory_lock` for the whole load and were killed by the PostgreSQL role's `statement_timeout` — `SQLSTATE 57014` — so from the operator's seat the model simply never loaded.) +- If the load fails, the waiter gets the *real* cause (`worker out of disk`), not an anonymous timeout. +- If the client disconnects, the load keeps going. It belongs to the job record, not to the request. + +When the wait budget (`LOCALAI_MODEL_LOAD_WAIT`, default `60s`) runs out, the request is answered with `503`, a `Retry-After` header, and a body that says exactly where the load is: + +```json +{ + "error": { + "message": "model Qwen3.6-27B-MTP-GGUF is staging on node nvidia-thor (41%, ETA ~11m)", + "type": "model_loading", + "code": "model_loading" + }, + "loading": { + "model": "Qwen3.6-27B-MTP-GGUF", + "state": "staging", + "node": "nvidia-thor", + "progress": 41.2, + "bytes_sent": 14730000000, + "total_bytes": 35776484480, + "file_index": 1, + "total_files": 2, + "eta_seconds": 660 + } +} +``` + +The `error` envelope keeps OpenAI clients working unchanged; `loading` is additive, so a client that understands it renders progress instead of an error. `eta_seconds` is derived from the job's own observed transfer rate and is **omitted rather than guessed** until enough bytes have moved for that rate to mean anything — a confidently wrong ETA on a twenty-minute wait is worse than none. `state` is one of `pending` (choosing a node), `installing`, `staging` (transferring files) or `loading` (the worker is reading the checkpoint). + +The chat UI renders this state inline and retries automatically once the model reports ready. Poll `GET /api/models/{id}/load-status` for the same `loading` object at any time. + +{{% notice note %}} +A frontend replica that dies mid-load does not wedge the model: the job row carries a heartbeat and another replica reclaims a job whose heartbeat has stopped. The heartbeat is time-based, not byte-based, because a checkpoint load legitimately transfers zero bytes for many minutes. +{{% /notice %}} + ### NATS JWT authentication (recommended for production) By default, NATS connections are anonymous: any client that can reach port `4222` may publish control-plane subjects such as `nodes..backend.install`. Enable JWT auth to scope workers to their own node subjects and give the frontend a dedicated service credential. diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 6e0abee69..fd54d358f 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -469,7 +469,10 @@ func (ml *ModelLoader) loadModel(modelID, modelName, modelFileName string, loade modelFile := filepath.Join(ml.ModelPath, modelFileName) model, err := loader(modelID, modelName, modelFile) if err != nil { - return nil, fmt.Errorf("failed to route model with internal loader: %s", err) + // %w, not %s: the router reports a still-loading model as a typed + // error that the HTTP layer turns into 503 plus live progress, and + // that only survives an unbroken chain. + return nil, fmt.Errorf("failed to route model with internal loader: %w", err) } if model == nil { return nil, fmt.Errorf("loader didn't return a model") diff --git a/swagger/docs.go b/swagger/docs.go index 14b0f3352..d785faa68 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -1097,6 +1097,41 @@ const docTemplate = `{ } } }, + "/api/models/{id}/load-status": { + "get": { + "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.", + "produces": [ + "application/json" + ], + "tags": [ + "models" + ], + "summary": "Report the progress of an in-flight model load.", + "parameters": [ + { + "type": "string", + "description": "Model ID", + "name": "id", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "Live load progress", + "schema": { + "$ref": "#/definitions/schema.ModelLoadingStatus" + } + }, + "404": { + "description": "No load is running for this model", + "schema": { + "$ref": "#/definitions/schema.ErrorResponse" + } + } + } + } + }, "/api/models/{name}/{action}": { "put": { "description": "Enable or disable a model from being loaded on demand. Disabled models remain installed but cannot be loaded.", @@ -6064,6 +6099,39 @@ const docTemplate = `{ } } }, + "schema.ModelLoadingStatus": { + "type": "object", + "properties": { + "bytes_sent": { + "type": "integer" + }, + "eta_seconds": { + "description": "ETASeconds is omitted rather than guessed until enough bytes have moved\nfor the observed rate to mean anything. A confidently wrong ETA on a\ntwenty-minute wait is worse than none.", + "type": "integer" + }, + "file_index": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "node": { + "type": "string" + }, + "progress": { + "type": "number" + }, + "state": { + "type": "string" + }, + "total_bytes": { + "type": "integer" + }, + "total_files": { + "type": "integer" + } + } + }, "schema.ModelsDataResponse": { "type": "object", "properties": { diff --git a/swagger/swagger.json b/swagger/swagger.json index 4181a0e82..8b7f4803a 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -1094,6 +1094,41 @@ } } }, + "/api/models/{id}/load-status": { + "get": { + "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.", + "produces": [ + "application/json" + ], + "tags": [ + "models" + ], + "summary": "Report the progress of an in-flight model load.", + "parameters": [ + { + "type": "string", + "description": "Model ID", + "name": "id", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "Live load progress", + "schema": { + "$ref": "#/definitions/schema.ModelLoadingStatus" + } + }, + "404": { + "description": "No load is running for this model", + "schema": { + "$ref": "#/definitions/schema.ErrorResponse" + } + } + } + } + }, "/api/models/{name}/{action}": { "put": { "description": "Enable or disable a model from being loaded on demand. Disabled models remain installed but cannot be loaded.", @@ -6061,6 +6096,39 @@ } } }, + "schema.ModelLoadingStatus": { + "type": "object", + "properties": { + "bytes_sent": { + "type": "integer" + }, + "eta_seconds": { + "description": "ETASeconds is omitted rather than guessed until enough bytes have moved\nfor the observed rate to mean anything. A confidently wrong ETA on a\ntwenty-minute wait is worse than none.", + "type": "integer" + }, + "file_index": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "node": { + "type": "string" + }, + "progress": { + "type": "number" + }, + "state": { + "type": "string" + }, + "total_bytes": { + "type": "integer" + }, + "total_files": { + "type": "integer" + } + } + }, "schema.ModelsDataResponse": { "type": "object", "properties": { diff --git a/swagger/swagger.yaml b/swagger/swagger.yaml index 8b66f6892..f944d88f8 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -1583,6 +1583,31 @@ definitions: an error). type: string type: object + schema.ModelLoadingStatus: + properties: + bytes_sent: + type: integer + eta_seconds: + description: |- + ETASeconds is omitted rather than guessed until enough bytes have moved + for the observed rate to mean anything. A confidently wrong ETA on a + twenty-minute wait is worse than none. + type: integer + file_index: + type: integer + model: + type: string + node: + type: string + progress: + type: number + state: + type: string + total_bytes: + type: integer + total_files: + type: integer + type: object schema.ModelsDataResponse: properties: data: @@ -3454,6 +3479,32 @@ paths: summary: Get an instruction's API guide or OpenAPI fragment tags: - instructions + /api/models/{id}/load-status: + get: + 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. + parameters: + - description: Model ID + in: path + name: id + required: true + type: string + produces: + - application/json + responses: + "200": + description: Live load progress + schema: + $ref: '#/definitions/schema.ModelLoadingStatus' + "404": + description: No load is running for this model + schema: + $ref: '#/definitions/schema.ErrorResponse' + summary: Report the progress of an in-flight model load. + tags: + - models /api/models/{name}/{action}: put: description: Enable or disable a model from being loaded on demand. Disabled