diff --git a/.gitignore b/.gitignore index 666b81df934e..689c043d1b61 100644 --- a/.gitignore +++ b/.gitignore @@ -41,7 +41,12 @@ models/* test-models/ test-dir/ tests/e2e-aio/backends -mock-backend +# The mock backend binary built by `make build-mock-backend`. Anchored to its +# full path: a bare `mock-backend` also matched the *directory* holding the +# source, so git would not descend into it and adding a file there needed -f. +# tests/e2e/mock-backend/.gitignore covers the same binary; kept here too so +# the artifact stays ignored if that scoped file is ever removed. +/tests/e2e/mock-backend/mock-backend release/ diff --git a/core/application/application.go b/core/application/application.go index ab7f34a41f54..df904c02d524 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -94,9 +94,30 @@ type Application struct { // is set; otherwise initialised in start() after galleryService. localAIAssistant *mcpTools.LocalAIAssistantHolder + // startupComplete flips to true once New() has finished its whole startup + // sequence. It backs the /readyz probe. + // + // The expensive step is the model preload: since #10949 it materializes + // HuggingFace artifacts for managed backends, which is tens of GB for a + // large model (31 GB observed on a live cluster). Tracking it explicitly + // means readiness reports lifecycle state instead of "the handler was + // reachable", so a replica that is still starting can be kept out of a + // load balancer's rotation. + startupComplete atomic.Bool + shutdownOnce sync.Once } +// Ready reports whether the application has finished starting up and can serve +// traffic. It backs the /readyz probe and is safe to call from any goroutine, +// including while startup is still running. +func (a *Application) Ready() bool { return a.startupComplete.Load() } + +// markStartupComplete flips the application to ready. Called once, at the very +// end of New(), so every startup step — model preload included — has finished +// before the process advertises itself as able to serve. +func (a *Application) markStartupComplete() { a.startupComplete.Store(true) } + func newApplication(appConfig *config.ApplicationConfig) *Application { ml := model.NewModelLoader(appConfig.SystemState) diff --git a/core/application/readiness_test.go b/core/application/readiness_test.go new file mode 100644 index 000000000000..328d16341760 --- /dev/null +++ b/core/application/readiness_test.go @@ -0,0 +1,122 @@ +package application + +import ( + "context" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/system" +) + +// blockingMaterializer parks inside Ensure until release is closed, so a spec +// can hold the startup preload open and observe readiness while it is still +// running. This mirrors how a real managed-artifact preload behaves: since +// #10949 it downloads a HuggingFace snapshot, which for a large model is tens +// of GB and can take a long time. +type blockingMaterializer struct { + seen chan struct{} + release chan struct{} +} + +func (b *blockingMaterializer) Ensure(ctx context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { + select { + case b.seen <- struct{}{}: + default: + } + select { + case <-b.release: + case <-ctx.Done(): + return modelartifacts.Result{}, ctx.Err() + } + spec.Resolved = &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + } + return modelartifacts.Result{Spec: spec}, nil +} + +var _ = Describe("application readiness", func() { + var modelsPath string + var state *system.SystemState + + BeforeEach(func() { + modelsPath = GinkgoT().TempDir() + var err error + state, err = system.GetSystemState( + system.WithModelPath(modelsPath), + system.WithBackendPath(GinkgoT().TempDir()), + ) + Expect(err).NotTo(HaveOccurred()) + }) + + It("is not ready before the startup sequence has run", func() { + app := newApplication(&config.ApplicationConfig{ + Context: context.Background(), + SystemState: state, + }) + + // A constructed-but-unstarted application must never advertise itself + // as able to serve traffic: /readyz reads this, and a load balancer + // that believes it will route requests to a process that cannot answer. + Expect(app.Ready()).To(BeFalse()) + }) + + It("stays not-ready while the startup model preload is still running", func() { + blocker := &blockingMaterializer{ + seen: make(chan struct{}, 1), + release: make(chan struct{}), + } + app := newApplication(&config.ApplicationConfig{ + Context: context.Background(), + SystemState: state, + ModelArtifactMaterializer: blocker, + }) + + Expect(os.WriteFile(filepath.Join(modelsPath, "managed.yaml"), []byte(` +name: managed +backend: transformers +artifacts: + - source: {type: huggingface, repo: owner/repo} +parameters: {model: owner/repo} +`), 0644)).To(Succeed()) + Expect(app.ModelConfigLoader().LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + + preloadDone := make(chan error, 1) + go func() { + preloadDone <- app.ModelConfigLoader().PreloadWithContext(context.Background(), modelsPath) + }() + + // Preload is now parked inside artifact materialization — exactly the + // window in which a Kubernetes Service would otherwise send this + // replica half the cluster's traffic. + Eventually(blocker.seen).Should(Receive()) + Consistently(app.Ready).Should(BeFalse()) + + close(blocker.release) + Expect(<-preloadDone).NotTo(HaveOccurred()) + + // Preload finished, but the rest of the startup sequence has not, so + // the application is still not ready. + Expect(app.Ready()).To(BeFalse()) + }) + + It("becomes ready once the startup sequence completes", func() { + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + + app, err := New( + config.WithContext(ctx), + config.WithSystemState(state), + ) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { _ = app.Shutdown() }) + + Expect(app.Ready()).To(BeTrue()) + }) +}) diff --git a/core/application/startup.go b/core/application/startup.go index be0bc2d18081..455d78c13744 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -496,6 +496,12 @@ func New(opts ...config.AppOption) (*Application, error) { // Watch the configuration directory startWatcher(options) + // Everything that must happen before this process can serve a request has + // happened. Flip readiness last, and only on the success path — the early + // `return nil, err` exits above abort startup, and an application that + // never finished starting must never report itself ready. + application.markStartupComplete() + xlog.Info("core/startup process completed!") return application, nil } diff --git a/core/http/app.go b/core/http/app.go index a5b13ed59a01..4d969b8fdb65 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -263,7 +263,7 @@ func API(application *application.Application) (*echo.Echo, error) { } // Health Checks should always be exempt from auth, so register these first - routes.HealthRoutes(e) + routes.HealthRoutes(e, application.Ready) // Build auth middleware: use the new auth.Middleware when auth is enabled or // as a unified replacement for the legacy key-auth middleware. diff --git a/core/http/routes/health.go b/core/http/routes/health.go index 5b03953733d8..343584155be1 100644 --- a/core/http/routes/health.go +++ b/core/http/routes/health.go @@ -1,15 +1,40 @@ package routes import ( + "net/http" + "github.com/labstack/echo/v4" ) -func HealthRoutes(app *echo.Echo) { - // Service health checks - ok := func(c echo.Context) error { - return c.NoContent(200) - } +// HealthRoutes registers the liveness (/healthz) and readiness (/readyz) +// probes. +// +// ready reports whether the process has finished starting up. It is consulted +// per request rather than sampled once at registration time, because readiness +// is a lifecycle property that changes after the router has been built. A nil +// ready fails open: an embedder that does not wire readiness keeps the +// historical always-200 behaviour instead of being stuck out of rotation +// forever. +func HealthRoutes(app *echo.Echo, ready func() bool) { + // Liveness: "is this process alive?". Deliberately independent of + // readiness — a long startup preload must not read as a hung process, or + // an orchestrator restarts the pod mid-download and the preload can never + // finish. + app.GET("/healthz", func(c echo.Context) error { + return c.NoContent(http.StatusOK) + }) - app.GET("/healthz", ok) - app.GET("/readyz", ok) + // Readiness: "should this replica receive traffic?". Answering 200 while + // the process is still starting up makes a Kubernetes Service add the pod + // to its endpoints early, and traffic round-robins onto a replica that + // cannot serve it. + app.GET("/readyz", func(c echo.Context) error { + if ready != nil && !ready() { + return c.JSON(http.StatusServiceUnavailable, map[string]string{ + "status": "starting", + "reason": "startup preload in progress", + }) + } + return c.NoContent(http.StatusOK) + }) } diff --git a/core/http/routes/health_test.go b/core/http/routes/health_test.go new file mode 100644 index 000000000000..e3a4fa57b17d --- /dev/null +++ b/core/http/routes/health_test.go @@ -0,0 +1,61 @@ +package routes_test + +import ( + "net/http" + "net/http/httptest" + "sync/atomic" + + "github.com/labstack/echo/v4" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/http/routes" +) + +var _ = Describe("Health and readiness probes", func() { + var e *echo.Echo + var ready atomic.Bool + + get := func(path string) *httptest.ResponseRecorder { + rec := httptest.NewRecorder() + e.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, path, nil)) + return rec + } + + BeforeEach(func() { + e = echo.New() + ready.Store(false) + routes.HealthRoutes(e, ready.Load) + }) + + It("reports /readyz as unavailable while startup is in progress", func() { + Expect(get("/readyz").Code).To(Equal(http.StatusServiceUnavailable)) + }) + + It("reports /readyz as available once startup completes", func() { + ready.Store(true) + Expect(get("/readyz").Code).To(Equal(http.StatusOK)) + }) + + It("keeps /healthz green during startup", func() { + // Liveness and readiness answer different questions. Failing liveness + // during a long preload would make Kubernetes restart the pod and the + // preload would never finish. + Expect(get("/healthz").Code).To(Equal(http.StatusOK)) + + ready.Store(true) + Expect(get("/healthz").Code).To(Equal(http.StatusOK)) + }) + + It("treats a nil readiness source as ready", func() { + // Fail open: an embedder that does not wire readiness keeps the + // historical always-200 behaviour rather than being permanently + // out of rotation. + plain := echo.New() + routes.HealthRoutes(plain, nil) + + rec := httptest.NewRecorder() + plain.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/readyz", nil)) + Expect(rec.Code).To(Equal(http.StatusOK)) + }) +}) diff --git a/coverage-baseline.txt b/coverage-baseline.txt index 0b320ba1d7a6..ac1025b8090e 100644 --- a/coverage-baseline.txt +++ b/coverage-baseline.txt @@ -1 +1 @@ -48.5 +54.2